import json
import re
from copy import deepcopy
from time import monotonic
from typing import Any

from django.db import transaction
from django.utils import timezone

from apps.agents.models import (
    AgentDefinition,
    AgentRun,
    AgentRunAttempt,
    AgentRunStep,
    AgentStepBinding,
    AgentVersion,
)
from apps.audit.models import OperationAuditLog
from apps.modelhub.models import ModelInvocation
from apps.modelhub.services import execute_invocation
from apps.skills.services import (
    SkillValidationError,
    render_prompt,
    validate_json_schema,
    validate_payload,
)

SOURCE_PATH = re.compile(
    r"^\$(?:input(?:\.[A-Za-z_][A-Za-z0-9_]*)*|steps\.[A-Za-z_][A-Za-z0-9_-]*\.output(?:\.[A-Za-z_][A-Za-z0-9_]*)*)$"
)


class AgentValidationError(ValueError):
    pass


def validate_agent_version(version: AgentVersion, *, publishing: bool = False) -> None:
    validate_json_schema(version.input_schema, "Agent 输入 Schema")
    if version.example_input:
        validate_payload(version.input_schema, version.example_input, "Agent 示例输入")
    if version.temperature_override is not None and not 0 <= version.temperature_override <= 2:
        raise AgentValidationError("Temperature 必须在 0 到 2 之间")
    if (
        version.max_output_tokens_override is not None
        and not 1 <= version.max_output_tokens_override <= 32768
    ):
        raise AgentValidationError("最大输出 Token 超出允许范围")
    steps = list(version.steps.select_related("skill_version__definition"))
    if not steps:
        raise AgentValidationError("Agent 至少需要一个 Skill 步骤")
    if len(steps) > 10:
        raise AgentValidationError("Agent 最多允许 10 个步骤")
    known_steps: set[str] = set()
    for step in steps:
        if publishing and step.skill_version.status != "published":
            raise AgentValidationError(f"步骤 {step.code} 必须绑定已发布 Skill")
        for target, source in step.input_mapping.items():
            if not isinstance(target, str) or not target:
                raise AgentValidationError(f"步骤 {step.code} 的映射字段无效")
            if not isinstance(source, str) or not SOURCE_PATH.fullmatch(source):
                raise AgentValidationError(f"步骤 {step.code} 的映射来源无效: {source}")
            if source.startswith("$steps."):
                source_step = source.split(".")[1]
                if source_step not in known_steps:
                    raise AgentValidationError(
                        f"步骤 {step.code} 只能引用之前步骤的输出: {source_step}"
                    )
        known_steps.add(step.code)


def _resolve_source(path: str, run_input: dict, step_outputs: dict[str, dict]) -> Any:
    parts = path.removeprefix("$").split(".")
    if parts[0] == "input":
        value: Any = run_input
        tail = parts[1:]
    else:
        value = step_outputs.get(parts[1], {}).get("output")
        tail = parts[3:]
    for part in tail:
        if not isinstance(value, dict) or part not in value:
            raise AgentValidationError(f"映射来源不存在: {path}")
        value = value[part]
    return deepcopy(value)


def map_step_input(
    binding: AgentStepBinding,
    run_input: dict,
    step_outputs: dict[str, dict],
) -> dict:
    return {
        target: _resolve_source(source, run_input, step_outputs)
        for target, source in binding.input_mapping.items()
    }


def ensure_agent_draft(definition: AgentDefinition, user=None) -> AgentVersion:
    existing = definition.versions.filter(status="draft").first()
    if existing:
        return existing
    source = definition.active_version
    if source is None:
        raise AgentValidationError("Agent 没有可克隆的已发布版本")
    latest = definition.versions.order_by("-version").values_list("version", flat=True).first() or 0
    draft = AgentVersion.objects.create(
        definition=definition,
        version=latest + 1,
        status="draft",
        model_deployment=source.model_deployment,
        system_prompt=source.system_prompt,
        input_schema=deepcopy(source.input_schema),
        example_input=deepcopy(source.example_input),
        temperature_override=source.temperature_override,
        max_output_tokens_override=source.max_output_tokens_override,
        created_by=user,
    )
    AgentStepBinding.objects.bulk_create(
        [
            AgentStepBinding(
                agent_version=draft,
                code=step.code,
                order=step.order,
                skill_version=step.skill_version,
                input_mapping=deepcopy(step.input_mapping),
            )
            for step in source.steps.all()
        ]
    )
    return draft


def build_config_snapshot(version: AgentVersion) -> dict:
    deployment = version.model_deployment
    return {
        "agent": {
            "id": version.definition_id,
            "code": version.definition.code,
            "version": version.version,
            "status": version.status,
        },
        "model": {
            "deployment_id": deployment.id,
            "provider_name": deployment.provider.name,
            "model_name": deployment.model_name,
            "temperature": version.temperature_override
            if version.temperature_override is not None
            else deployment.default_temperature,
            "max_output_tokens": version.max_output_tokens_override or deployment.max_output_tokens,
        },
        "steps": [
            {
                "code": step.code,
                "order": step.order,
                "skill_id": step.skill_version.definition_id,
                "skill_code": step.skill_version.definition.code,
                "skill_version": step.skill_version.version,
                "input_mapping": step.input_mapping,
            }
            for step in version.steps.select_related("skill_version__definition")
        ],
    }


def _parse_json_output(raw_output: str) -> Any:
    value = raw_output.strip()
    if value.startswith("```"):
        value = re.sub(r"^```(?:json)?\s*", "", value)
        value = re.sub(r"\s*```$", "", value)
    try:
        return json.loads(value)
    except json.JSONDecodeError as error:
        raise SkillValidationError(f"模型输出不是合法 JSON: {error.msg}") from error


def _record_attempt(
    step: AgentRunStep,
    invocation: ModelInvocation,
    attempt_type: str,
    validation_error: str = "",
) -> AgentRunAttempt:
    return AgentRunAttempt.objects.create(
        step=step,
        model_invocation=invocation,
        attempt_type=attempt_type,
        status=invocation.status,
        raw_output=invocation.output_text,
        validation_error=validation_error,
        input_tokens=invocation.input_tokens or 0,
        output_tokens=invocation.output_tokens or 0,
        duration_ms=invocation.duration_ms,
    )


def _invoke_for_step(
    run: AgentRun,
    step: AgentRunStep,
    prompt: str,
    attempt_type: str,
) -> ModelInvocation:
    version = run.agent_version
    messages = []
    if version.system_prompt:
        messages.append({"role": "system", "content": version.system_prompt})
    messages.append({"role": "user", "content": prompt})
    invocation = ModelInvocation.objects.create(
        deployment=version.model_deployment,
        invocation_type="agent_step",
        input_snapshot={"messages": messages, "run_id": str(run.public_id), "step": step.code},
        created_by=run.created_by,
    )
    execute_invocation(
        invocation,
        messages,
        temperature=version.temperature_override,
        max_output_tokens=version.max_output_tokens_override,
    )
    _record_attempt(step, invocation, attempt_type)
    if invocation.retry_count:
        AgentRunAttempt.objects.create(
            step=step,
            model_invocation=invocation,
            attempt_type="transient_retry",
            status=invocation.status,
            raw_output=invocation.output_text,
            input_tokens=0,
            output_tokens=0,
            duration_ms=0,
        )
    return invocation


def _repair_output(
    run: AgentRun,
    step: AgentRunStep,
    raw_output: str,
    validation_error: str,
    output_schema: dict,
) -> ModelInvocation:
    prompt = (
        "请只输出修正后的 JSON，不要解释。\n"
        f"输出 Schema：{json.dumps(output_schema, ensure_ascii=False)}\n"
        f"校验错误：{validation_error}\n"
        f"原始输出：{raw_output}"
    )
    invocation = _invoke_for_step(run, step, prompt, "structure_repair")
    attempt = step.attempts.order_by("-created_at").first()
    if attempt:
        attempt.validation_error = validation_error
        attempt.save(update_fields=["validation_error"])
    return invocation


def execute_agent_run(run_id: int) -> AgentRun:
    run = AgentRun.objects.select_related(
        "definition",
        "agent_version__model_deployment__provider",
        "created_by",
    ).get(id=run_id)
    started = monotonic()
    run.status = "running"
    run.started_at = timezone.now()
    run.save(update_fields=["status", "started_at"])
    step_outputs: dict[str, dict] = {}
    try:
        bindings = list(
            run.agent_version.steps.select_related(
                "skill_version__definition",
            )
        )
        for binding in bindings:
            run.refresh_from_db(fields=["cancellation_requested_at"])
            if run.cancellation_requested_at:
                run.status = "cancelled"
                break
            run.current_step = binding.code
            run.save(update_fields=["current_step"])
            step = AgentRunStep.objects.create(
                run=run,
                code=binding.code,
                name=binding.skill_version.definition.name,
                order=binding.order,
                skill_version=binding.skill_version,
                status="running",
                started_at=timezone.now(),
            )
            step_started = monotonic()
            try:
                step_input = map_step_input(binding, run.input_payload, step_outputs)
                validate_payload(
                    binding.skill_version.input_schema,
                    step_input,
                    f"步骤 {binding.code} 输入",
                )
                step.input_payload = step_input
                step.save(update_fields=["input_payload"])
                prompt = render_prompt(binding.skill_version.prompt_template, step_input)
                invocation = _invoke_for_step(run, step, prompt, "initial")
                if invocation.status != "succeeded":
                    raise AgentValidationError(invocation.error_message or "模型调用失败")
                try:
                    output = _parse_json_output(invocation.output_text)
                    validate_payload(binding.skill_version.output_schema, output, "模型输出")
                except SkillValidationError as error:
                    repair = _repair_output(
                        run,
                        step,
                        invocation.output_text,
                        str(error),
                        binding.skill_version.output_schema,
                    )
                    if repair.status != "succeeded":
                        raise AgentValidationError(
                            repair.error_message or "结构修复失败"
                        ) from error
                    output = _parse_json_output(repair.output_text)
                    validate_payload(binding.skill_version.output_schema, output, "修复输出")
                step.output_payload = output
                step.status = "succeeded"
                step_outputs[binding.code] = {"output": output}
            except (AgentValidationError, SkillValidationError) as error:
                step.status = "failed"
                step.error_code = "step_failed"
                step.error_message = str(error)
                run.status = "failed"
                run.error_code = "step_failed"
                run.error_message = f"步骤 {binding.code} 失败: {error}"
            finally:
                attempts = step.attempts.all()
                step.input_tokens = sum(item.input_tokens for item in attempts)
                step.output_tokens = sum(item.output_tokens for item in attempts)
                step.duration_ms = max(int((monotonic() - step_started) * 1000), 0)
                step.finished_at = timezone.now()
                step.save()
            if run.status == "failed":
                break
            run.refresh_from_db(fields=["cancellation_requested_at"])
            if run.cancellation_requested_at:
                run.status = "cancelled"
                break
        if run.status == "running":
            run.status = "succeeded"
            run.final_output = step.output_payload if bindings else None
    except Exception as error:
        run.status = "failed"
        run.error_code = "runtime_error"
        run.error_message = "Agent 运行发生未预期错误"
        OperationAuditLog.objects.create(
            user=run.created_by,
            module="agents",
            action="runtime_error",
            target_type="agent_run",
            target_id=str(run.public_id),
            request_snapshot={"error_type": type(error).__name__},
            result="failed",
        )
    run.input_tokens = sum(item.input_tokens for item in run.steps.all())
    run.output_tokens = sum(item.output_tokens for item in run.steps.all())
    run.duration_ms = max(int((monotonic() - started) * 1000), 0)
    run.finished_at = timezone.now()
    run.save()
    OperationAuditLog.objects.create(
        user=run.created_by,
        module="agents",
        action="run_finished",
        target_type="agent_run",
        target_id=str(run.public_id),
        request_snapshot={"status": run.status, "duration_ms": run.duration_ms},
        result="success" if run.status == "succeeded" else run.status,
    )
    return run


def replace_agent_steps(version: AgentVersion, steps: list[dict]) -> None:
    if len(steps) > 10:
        raise AgentValidationError("Agent 最多允许 10 个步骤")
    with transaction.atomic():
        version.steps.all().delete()
        AgentStepBinding.objects.bulk_create(
            [
                AgentStepBinding(
                    agent_version=version,
                    code=item["code"],
                    order=index + 1,
                    skill_version_id=item["skill_version_id"],
                    input_mapping=item.get("input_mapping", {}),
                )
                for index, item in enumerate(steps)
            ]
        )
        validate_agent_version(version)