import json
from datetime import timedelta

from celery import shared_task
from django.conf import settings
from django.db import transaction
from django.utils import timezone
from jsonschema import Draft202012Validator, ValidationError

from apps.modelhub.models import ModelInvocation
from apps.modelhub.services import execute_invocation
from apps.modelhub.structured_output import (
    STRUCTURED_OUTPUT_TEST_EXPECTED,
    STRUCTURED_OUTPUT_TEST_SCHEMA,
)


@shared_task(name="apps.modelhub.tasks.run_model_test")
def run_model_test(invocation_id: int) -> None:
    invocation = ModelInvocation.objects.select_related("deployment__provider").get(
        id=invocation_id
    )
    messages = invocation.input_snapshot.get("messages", [])
    output_schema = invocation.input_snapshot.get("output_schema")
    execute_invocation(
        invocation,
        messages,
        output_schema=output_schema,
        schema_name=invocation.input_snapshot.get("schema_name"),
    )
    if invocation.invocation_type != "structured_output_test":
        return

    deployment = invocation.deployment
    if deployment is None:
        return
    message = invocation.error_message
    passed = invocation.status == "succeeded"
    if passed:
        try:
            output = json.loads(invocation.output_text)
            Draft202012Validator(STRUCTURED_OUTPUT_TEST_SCHEMA).validate(output)
            if output != STRUCTURED_OUTPUT_TEST_EXPECTED:
                raise ValidationError("测试值与预期不一致")
        except (json.JSONDecodeError, ValidationError) as error:
            passed = False
            invocation.status = "failed"
            invocation.error_code = "structured_output_invalid"
            invocation.error_message = f"结构化输出测试未通过：{error}"
            invocation.save(update_fields=["status", "error_code", "error_message"])
            message = invocation.error_message

    deployment.structured_output_test_status = "passed" if passed else "failed"
    deployment.structured_output_verified_at = timezone.now() if passed else None
    deployment.structured_output_test_message = (
        "结构化输出测试已通过" if passed else (message or "结构化输出测试失败")
    )
    deployment.save(
        update_fields=[
            "structured_output_test_status",
            "structured_output_verified_at",
            "structured_output_test_message",
            "updated_at",
        ]
    )
    from apps.audit.models import OperationAuditLog

    OperationAuditLog.objects.create(
        user=invocation.created_by,
        module="modelhub",
        action="structured_output_test_result",
        target_type="model_deployment",
        target_id=str(deployment.id),
        request_snapshot={
            "invocation_id": str(invocation.public_id),
            "status": deployment.structured_output_test_status,
            "error_code": invocation.error_code,
        },
        result="success" if passed else "failed",
    )


def _retained_input_metadata(snapshot: dict) -> dict:
    structured = snapshot.get("structured_output")
    retained: dict = {}
    if isinstance(structured, dict):
        retained["structured_output"] = {
            key: structured.get(key)
            for key in (
                "mode",
                "schema_name",
                "schema_hash",
                "response_format_sent",
            )
            if key in structured
        }
    if snapshot.get("schema_name"):
        retained["schema_name"] = snapshot["schema_name"]
    retained["diagnostics_scrubbed"] = True
    return retained


@shared_task(name="apps.modelhub.tasks.scrub_expired_model_invocation_diagnostics")
def scrub_expired_model_invocation_diagnostics() -> int:
    """擦除过期诊断正文，保留调用元数据、Schema 证据、Token 与终态。"""

    cutoff = timezone.now() - timedelta(
        days=settings.MODEL_INVOCATION_DIAGNOSTIC_RETENTION_DAYS
    )
    from apps.agents.models import AgentRunAttempt

    scrubbed_count = 0
    while True:
        invocation_ids = list(
            ModelInvocation.objects.filter(
                created_at__lt=cutoff,
                diagnostics_scrubbed_at__isnull=True,
            ).values_list("id", flat=True)[:1000]
        )
        if not invocation_ids:
            break
        now = timezone.now()
        with transaction.atomic():
            invocations = list(
                ModelInvocation.objects.select_for_update().filter(id__in=invocation_ids)
            )
            for invocation in invocations:
                invocation.input_snapshot = _retained_input_metadata(
                    invocation.input_snapshot or {}
                )
                invocation.output_text = ""
                invocation.diagnostics_scrubbed_at = now
            ModelInvocation.objects.bulk_update(
                invocations,
                ["input_snapshot", "output_text", "diagnostics_scrubbed_at"],
            )
            AgentRunAttempt.objects.filter(
                model_invocation_id__in=invocation_ids
            ).update(raw_output="")
        scrubbed_count += len(invocations)
    return scrubbed_count