from django.core.management.base import BaseCommand
from django.contrib.auth import get_user_model
from user_management.models import ChatHistory
from django.utils import timezone
import uuid

User = get_user_model()

class Command(BaseCommand):
    help = '创建聊天测试数据'

    def handle(self, *args, **kwargs):
        try:
            # 获取测试用户
            user = User.objects.first()
            if not user:
                self.stdout.write(self.style.ERROR('没有找到用户，请先创建用户'))
                return

            # 创建3个不同的对话
            for i in range(3):
                conversation_id = f"conv_{uuid.uuid4().hex[:10]}"
                dataset_id = f"dataset_{i+1}"
                dataset_name = f"测试数据集_{i+1}"

                # 每个对话创建3轮问答
                for j in range(3):
                    # 用户问题
                    ChatHistory.objects.create(
                        user=user,
                        dataset_id=dataset_id,
                        dataset_name=dataset_name,
                        conversation_id=conversation_id,
                        role='user',
                        content=f"这是第{i+1}个数据集的第{j+1}个问题？",
                        question=f"这是第{i+1}个数据集的第{j+1}个问题？",
                        answer="",
                    )

                    # AI回答
                    ChatHistory.objects.create(
                        user=user,
                        dataset_id=dataset_id,
                        dataset_name=dataset_name,
                        conversation_id=conversation_id,
                        role='assistant',
                        content=f"这是第{i+1}个数据集的第{j+1}个回答。",
                        question="",
                        answer=f"这是第{i+1}个数据集的第{j+1}个回答。",
                    )

            self.stdout.write(self.style.SUCCESS('成功创建测试聊天数据'))

        except Exception as e:
            self.stdout.write(self.style.ERROR(f'创建测试数据失败: {str(e)}'))