"""Start, inspect, and cancel persisted batch text test runs."""

from __future__ import annotations

from datetime import UTC, datetime
from uuid import uuid4

from db.models import (
    Assistant,
    ModelResource,
    TestCase,
    TestRun,
    TestRunCase,
    TestSuite,
)
from db.session import get_session
from fastapi import APIRouter, Depends, HTTPException
from models import RuntimeModelResource
from services.auth import require_admin
from services.config_resolver import resolve_runtime_config
from services.runtime_variables import DynamicVariableError, prepare_dynamic_config
from services.realtime.launcher import validate_visual_runtime
from services.test_case_assets import canonicalize_test_assets, image_asset_ids
from services.test_runs.orchestrator import test_run_orchestrator
from services.test_runs.prompt_improvement import (
    PROMPT_IMPROVEMENT_SNAPSHOT_KEY,
    PromptImprovementAgent,
    PromptImprovementError,
    public_test_run_config,
)
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from test_schemas import (
    BatchExecutionError,
    BatchRunCaseOut,
    BatchRunCreate,
    BatchRunSnapshotOut,
    BatchStartupResult,
    BatchTurnResult,
    BatchEvaluationResult,
    PromptImprovementOut,
    TestCaseDefinition,
)


router = APIRouter(
    prefix="/api/test-runs",
    tags=["tests"],
    dependencies=[Depends(require_admin)],
)


def _new_id(prefix: str) -> str:
    return f"{prefix}_{uuid4().hex}"


async def _run_out(
    session: AsyncSession,
    run: TestRun,
) -> BatchRunSnapshotOut:
    rows = (
        await session.execute(
            select(TestRunCase)
            .where(TestRunCase.run_id == run.id)
            .order_by(TestRunCase.position)
        )
    ).scalars().all()
    cases: list[BatchRunCaseOut] = []
    for row in rows:
        result = dict(row.result or {})
        cases.append(
            BatchRunCaseOut(
                id=str(row.test_case_id or row.case_snapshot.get("id") or row.id),
                name=row.test_case_name,
                status=row.status,
                turns=[
                    BatchTurnResult.model_validate(item)
                    for item in result.get("turns") or []
                ],
                startup=BatchStartupResult.model_validate(
                    result.get("startup") or {}
                ),
                overall_criteria=[
                    BatchEvaluationResult.model_validate(item)
                    for item in result.get("overallCriteria")
                    or result.get("overall_criteria")
                    or []
                ],
                attempt_count=row.attempt_count,
                max_attempts=row.max_attempts,
                execution_error=(
                    BatchExecutionError.model_validate(row.execution_error)
                    if row.execution_error
                    else None
                ),
            )
        )
    return BatchRunSnapshotOut(
        id=run.id,
        status=run.status,
        title=run.title,
        assistant_name=run.assistant_name,
        config=public_test_run_config(run.config),
        cases=cases,
        started_at=run.started_at or run.created_at,
        finished_at=run.finished_at,
        stop_reason=run.stop_reason,
    )


@router.post("", response_model=BatchRunSnapshotOut)
async def create_test_run(
    body: BatchRunCreate,
    session: AsyncSession = Depends(get_session),
):
    assistant = await session.get(Assistant, body.assistant_id)
    if assistant is None:
        raise HTTPException(404, "被测助手不存在")
    if assistant.runtime_mode != "pipeline":
        raise HTTPException(422, "第一版批量测试只支持 Pipeline 运行模式")
    if assistant.type not in {"prompt", "workflow"}:
        raise HTTPException(422, f"第一版批量测试暂不支持 {assistant.type} 类型助手")

    evaluator_resource = await session.get(
        ModelResource,
        body.evaluator_model_resource_id,
    )
    if (
        evaluator_resource is None
        or evaluator_resource.capability != "LLM"
        or not evaluator_resource.enabled
    ):
        raise HTTPException(422, "评估模型不存在、已停用或不是 LLM 资源")
    try:
        runtime_config = await resolve_runtime_config(session, assistant.id)
        prepare_dynamic_config(
            runtime_config,
            body.dynamic_variables,
            assistant_id=assistant.id,
        )
    except (ValueError, DynamicVariableError) as exc:
        raise HTTPException(422, str(exc)) from exc

    rows = (
        await session.execute(select(TestCase).where(TestCase.id.in_(body.case_ids)))
    ).scalars().all()
    by_id = {row.id: row for row in rows}
    missing = [case_id for case_id in body.case_ids if case_id not in by_id]
    if missing:
        raise HTTPException(422, "测试用例不存在: " + "、".join(missing))
    ordered_cases = [by_id[case_id] for case_id in body.case_ids]
    tools_by_id = {tool.id: tool for tool in runtime_config.tools}
    definitions_by_id: dict[str, TestCaseDefinition] = {}
    has_image_input = False
    for row in ordered_cases:
        if row.input_mode != "fixed_script_text":
            raise HTTPException(422, f"用例“{row.name}”不是可运行的固定多模态模式")
        try:
            definition = TestCaseDefinition.model_validate(row.definition or {})
            await canonicalize_test_assets(session, definition)
        except ValueError as exc:
            raise HTTPException(422, f"用例“{row.name}”校验失败: {exc}") from exc
        definitions_by_id[row.id] = definition
        has_image_input = has_image_input or bool(image_asset_ids(definition))
        for simulation in definition.tool_simulations:
            tool = tools_by_id.get(simulation.tool_id)
            if tool is None or tool.function_name != simulation.function_name:
                raise HTTPException(
                    422,
                    f"用例“{row.name}”引用的工具 {simulation.function_name} "
                    "不属于当前被测助手",
                )
        for turn in definition.turns:
            for behavior in turn.behaviors:
                if behavior.type != "tool_call":
                    continue
                tool = tools_by_id.get(behavior.tool_id)
                if tool is None or tool.function_name != behavior.function_name:
                    raise HTTPException(
                        422,
                        f"用例“{row.name}”引用的工具 {behavior.function_name} "
                        "不属于当前被测助手",
                    )

    if has_image_input:
        try:
            vision_enabled = validate_visual_runtime(runtime_config)
        except ValueError as exc:
            raise HTTPException(422, str(exc)) from exc
        if not vision_enabled:
            raise HTTPException(422, "测试用例包含图片，但当前助手未开启视觉能力")

    suite_ids = list(dict.fromkeys(row.suite_id for row in ordered_cases))
    suites = (
        await session.execute(select(TestSuite).where(TestSuite.id.in_(suite_ids)))
    ).scalars().all()
    suite_by_id = {row.id: row for row in suites}
    if body.title:
        title = body.title.strip()
    elif len(suite_ids) == 1 and suite_ids[0] in suite_by_id:
        title = f"{suite_by_id[suite_ids[0]].name} · {len(ordered_cases)} 个用例"
    else:
        title = f"批量测试 · {len(ordered_cases)} 个用例"

    now = datetime.now(UTC)
    run = TestRun(
        id=_new_id("run"),
        assistant_id=assistant.id,
        assistant_name=assistant.name,
        title=title,
        status="queued",
        config={
            "suiteCount": len(suite_ids),
            "evaluatorModelResourceId": evaluator_resource.id,
            "evaluatorModelResourceName": evaluator_resource.name,
            "evaluatorModel": str(
                (evaluator_resource.values or {}).get("modelId") or ""
            ),
            "dynamicVariables": body.dynamic_variables,
            **body.config.model_dump(mode="json", by_alias=True),
        },
        started_at=now,
    )
    session.add(run)
    for position, row in enumerate(ordered_cases):
        definition = definitions_by_id[row.id]
        snapshot = {
            "id": row.id,
            "suiteId": row.suite_id,
            "name": row.name,
            "description": row.description,
            "inputMode": row.input_mode,
            **definition.model_dump(mode="json", by_alias=True),
        }
        session.add(
            TestRunCase(
                id=_new_id("rc"),
                run_id=run.id,
                test_case_id=row.id,
                test_case_name=row.name,
                suite_id=row.suite_id,
                position=position,
                case_snapshot=snapshot,
                status="waiting",
                attempt_count=0,
                max_attempts=body.config.error_retry_count + 1,
                result={"turns": [], "overallCriteria": []},
            )
        )
    await session.commit()
    await session.refresh(run)
    test_run_orchestrator.start(run.id)
    return await _run_out(session, run)


@router.post(
    "/{run_id}/prompt-improvement",
    response_model=PromptImprovementOut,
)
async def create_prompt_improvement(
    run_id: str,
    session: AsyncSession = Depends(get_session),
):
    run = await session.get(TestRun, run_id)
    if run is None:
        raise HTTPException(404, "测试运行不存在")
    rows = (
        await session.execute(
            select(TestRunCase)
            .where(TestRunCase.run_id == run.id)
            .order_by(TestRunCase.position)
        )
    ).scalars().all()
    if len(rows) != 1:
        raise HTTPException(409, "Prompt 改进仅支持单用例测试运行")
    run_case = rows[0]
    if run.status != "completed" or run_case.status != "fail":
        raise HTTPException(409, "只有已完成且断言失败的用例可以生成 Prompt 改进")

    config = dict(run.config or {})
    snapshot = config.get(PROMPT_IMPROVEMENT_SNAPSHOT_KEY)
    if not isinstance(snapshot, dict):
        raise HTTPException(409, "该运行缺少 Prompt 配置快照，请重新测试")
    if snapshot.get("assistantType") != "prompt" or snapshot.get("runtimeMode") != "pipeline":
        raise HTTPException(409, "Prompt 改进仅支持 Pipeline Prompt 助手")

    evaluator_id = str(config.get("evaluatorModelResourceId") or "")
    evaluator = await session.get(ModelResource, evaluator_id)
    if evaluator is None or evaluator.capability != "LLM" or not evaluator.enabled:
        raise HTTPException(422, "该运行使用的评估模型已不存在或不可用")
    resource = RuntimeModelResource(
        id=evaluator.id,
        name=evaluator.name,
        capability=evaluator.capability,
        interface_type=evaluator.interface_type,
        values=evaluator.values or {},
        secrets=evaluator.secrets or {},
        support_image_input=bool(evaluator.support_image_input),
    )
    try:
        return await PromptImprovementAgent(resource).improve(
            snapshot=snapshot,
            case_snapshot=dict(run_case.case_snapshot or {}),
            result=dict(run_case.result or {}),
        )
    except PromptImprovementError as exc:
        raise HTTPException(exc.status_code, exc.message) from exc


@router.get("/{run_id}", response_model=BatchRunSnapshotOut)
async def get_test_run(
    run_id: str,
    session: AsyncSession = Depends(get_session),
):
    run = await session.get(TestRun, run_id)
    if run is None:
        raise HTTPException(404, "批量测试运行不存在")
    return await _run_out(session, run)


@router.post("/{run_id}/cancel", response_model=BatchRunSnapshotOut)
async def cancel_test_run(
    run_id: str,
    session: AsyncSession = Depends(get_session),
):
    run = await session.get(TestRun, run_id)
    if run is None:
        raise HTTPException(404, "批量测试运行不存在")
    if run.status in {"completed", "cancelled"}:
        return await _run_out(session, run)
    run.cancel_requested = True
    run.stop_reason = "manual"
    await session.commit()
    await test_run_orchestrator.cancel(run_id)
    await session.refresh(run)
    return await _run_out(session, run)