from django.core.management.base import BaseCommand, CommandError
from django.db import transaction
from django.utils import timezone

from apps.agents.models import AgentDefinition, AgentStepBinding, AgentVersion
from apps.modelhub.models import ModelDeployment
from apps.skills.models import SkillDefinition, SkillVersion

SKILLS = [
    {
        "code": "issue-structure",
        "name": "问题结构化",
        "description": "将技术支持问题整理为统一结构。",
        "prompt": (
            "请分析技术支持问题：{{ input.problem }}，环境：{{ input.environment }}。只输出 JSON。"
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "problem": {"type": "string"},
                "environment": {"type": "string"},
            },
            "required": ["problem", "environment"],
        },
        "output_schema": {
            "type": "object",
            "properties": {
                "summary": {"type": "string"},
                "symptoms": {"type": "array", "items": {"type": "string"}},
            },
            "required": ["summary", "symptoms"],
        },
    },
    {
        "code": "troubleshooting",
        "name": "排查建议",
        "description": "基于结构化问题生成排查步骤。",
        "prompt": (
            "根据问题摘要 {{ input.summary }} 和现象 {{ input.symptoms }} "
            "生成排查建议，只输出 JSON。"
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "summary": {"type": "string"},
                "symptoms": {"type": "array", "items": {"type": "string"}},
            },
            "required": ["summary", "symptoms"],
        },
        "output_schema": {
            "type": "object",
            "properties": {
                "likely_cause": {"type": "string"},
                "steps": {"type": "array", "items": {"type": "string"}},
            },
            "required": ["likely_cause", "steps"],
        },
    },
    {
        "code": "customer-reply",
        "name": "客户回复草稿",
        "description": "将排查建议整理为客户可读回复。",
        "prompt": (
            "原因：{{ input.likely_cause }}；步骤：{{ input.steps }}。生成客户回复，只输出 JSON。"
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "likely_cause": {"type": "string"},
                "steps": {"type": "array", "items": {"type": "string"}},
            },
            "required": ["likely_cause", "steps"],
        },
        "output_schema": {
            "type": "object",
            "properties": {"reply": {"type": "string"}},
            "required": ["reply"],
        },
    },
]


class Command(BaseCommand):
    help = "Seed the technical support analysis Agent and its published Skills."

    def add_arguments(self, parser):
        parser.add_argument("--deployment-id", type=int)

    @transaction.atomic
    def handle(self, *args, **options):
        deployment_id = options.get("deployment_id")
        deployments = ModelDeployment.objects.filter(enabled=True, provider__enabled=True)
        deployment = (
            deployments.filter(id=deployment_id).first() if deployment_id else deployments.first()
        )
        if deployment is None:
            raise CommandError(
                "Configure an enabled model deployment before seeding the demo Agent."
            )

        versions = {}
        for spec in SKILLS:
            definition, _ = SkillDefinition.objects.get_or_create(
                code=spec["code"],
                defaults={"name": spec["name"], "description": spec["description"]},
            )
            version = definition.versions.filter(version=1).first()
            if version is None:
                version = SkillVersion.objects.create(
                    definition=definition,
                    version=1,
                    status="published",
                    prompt_template=spec["prompt"],
                    input_schema=spec["input_schema"],
                    output_schema=spec["output_schema"],
                    example_input={},
                    published_at=timezone.now(),
                )
            definition.active_version = version
            definition.save(update_fields=["active_version", "updated_at"])
            versions[spec["code"]] = version

        definition, _ = AgentDefinition.objects.get_or_create(
            code="support-analysis",
            defaults={
                "name": "技术支持分析 Agent",
                "description": "结构化问题、生成排查建议和客户回复草稿。",
            },
        )
        draft = definition.versions.filter(status="draft").first()
        if draft is None:
            next_version = (
                definition.versions.order_by("-version").values_list("version", flat=True).first()
                or 0
            ) + 1
            draft = AgentVersion.objects.create(
                definition=definition,
                version=next_version,
                status="draft",
                model_deployment=deployment,
                system_prompt="你是企业技术支持助手。输出必须严格符合当前步骤要求的 JSON。",
                input_schema={
                    "type": "object",
                    "properties": {
                        "problem": {"type": "string"},
                        "environment": {"type": "string"},
                    },
                    "required": ["problem", "environment"],
                },
                example_input={
                    "problem": "采集任务长时间没有更新",
                    "environment": "生产环境",
                },
            )
            AgentStepBinding.objects.bulk_create(
                [
                    AgentStepBinding(
                        agent_version=draft,
                        code="structure",
                        order=1,
                        skill_version=versions["issue-structure"],
                        input_mapping={
                            "problem": "$input.problem",
                            "environment": "$input.environment",
                        },
                    ),
                    AgentStepBinding(
                        agent_version=draft,
                        code="troubleshooting",
                        order=2,
                        skill_version=versions["troubleshooting"],
                        input_mapping={
                            "summary": "$steps.structure.output.summary",
                            "symptoms": "$steps.structure.output.symptoms",
                        },
                    ),
                    AgentStepBinding(
                        agent_version=draft,
                        code="reply",
                        order=3,
                        skill_version=versions["customer-reply"],
                        input_mapping={
                            "likely_cause": "$steps.troubleshooting.output.likely_cause",
                            "steps": "$steps.troubleshooting.output.steps",
                        },
                    ),
                ]
            )
        self.stdout.write(
            self.style.SUCCESS(
                f"Seeded Agent {definition.code} draft v{draft.version} "
                f"with deployment {deployment.id}."
            )
        )