"""In-process MVP scheduler for persisted batch test runs."""

from __future__ import annotations

import asyncio
from dataclasses import dataclass, field
from datetime import UTC, datetime

from loguru import logger
from sqlalchemy import select

from db.models import ModelResource, TestCase, TestRun, TestRunCase
from db.session import SessionLocal
from models import RuntimeModelResource
from services.config_resolver import resolve_runtime_config
from services.test_runs.errors import TestExecutionError
from services.test_runs.evaluator import EvaluationEngine
from services.test_runs.text_runner import TextPipelineRunner
from test_schemas import TestCaseDefinition


@dataclass
class _RunControl:
    stop_event: asyncio.Event = field(default_factory=asyncio.Event)
    reason: str | None = None
    case_tasks: set[asyncio.Task[None]] = field(default_factory=set)


class TestRunOrchestrator:
    """Run cases concurrently while keeping PostgreSQL as the source of truth."""

    def __init__(self) -> None:
        self._run_tasks: dict[str, asyncio.Task[None]] = {}
        self._controls: dict[str, _RunControl] = {}

    def start(self, run_id: str) -> None:
        current = self._run_tasks.get(run_id)
        if current is not None and not current.done():
            return
        control = _RunControl()
        self._controls[run_id] = control
        task = asyncio.create_task(self._execute_run(run_id, control), name=f"test-run-{run_id}")
        self._run_tasks[run_id] = task
        task.add_done_callback(lambda _task: self._forget(run_id))

    async def cancel(self, run_id: str) -> None:
        control = self._controls.get(run_id)
        if control is None:
            return
        control.reason = "manual"
        control.stop_event.set()
        for task in tuple(control.case_tasks):
            task.cancel()

    async def shutdown(self) -> None:
        tasks = [task for task in self._run_tasks.values() if not task.done()]
        for task in tasks:
            task.cancel()
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)
        self._run_tasks.clear()
        self._controls.clear()

    def _forget(self, run_id: str) -> None:
        self._run_tasks.pop(run_id, None)
        self._controls.pop(run_id, None)

    async def _execute_run(self, run_id: str, control: _RunControl) -> None:
        try:
            async with SessionLocal() as session:
                run = await session.get(TestRun, run_id)
                if run is None:
                    return
                run.status = "running"
                run.started_at = run.started_at or datetime.now(UTC)
                await session.commit()
                try:
                    cfg = await resolve_runtime_config(session, str(run.assistant_id or ""))
                except Exception as exc:
                    await self._fail_run_setup(run_id, str(exc))
                    return
                config = dict(run.config or {})
                evaluator_resource_id = str(
                    config.get("evaluatorModelResourceId") or ""
                )
                evaluator_row = await session.get(
                    ModelResource,
                    evaluator_resource_id,
                )
                if (
                    evaluator_row is None
                    or evaluator_row.capability != "LLM"
                    or not evaluator_row.enabled
                ):
                    await self._fail_run_setup(
                        run_id,
                        "评估模型不存在、已停用或不是 LLM 资源",
                    )
                    return
                evaluator_resource = RuntimeModelResource(
                    id=evaluator_row.id,
                    name=evaluator_row.name,
                    capability=evaluator_row.capability,
                    interface_type=evaluator_row.interface_type,
                    values=evaluator_row.values or {},
                    secrets=evaluator_row.secrets or {},
                    support_image_input=bool(evaluator_row.support_image_input),
                )
                rows = (
                    await session.execute(
                        select(TestRunCase)
                        .where(TestRunCase.run_id == run_id)
                        .order_by(TestRunCase.position)
                    )
                ).scalars().all()

            semaphore = asyncio.Semaphore(int(config.get("concurrency") or 3))
            tasks = {
                asyncio.create_task(
                    self._execute_case(
                        run_id,
                        row.id,
                        cfg,
                        evaluator_resource,
                        config,
                        semaphore,
                        control,
                    ),
                    name=f"test-run-case-{row.id}",
                )
                for row in rows
            }
            control.case_tasks = tasks
            if tasks:
                await asyncio.gather(*tasks, return_exceptions=True)
            await self._finish_run(run_id, control)
        except asyncio.CancelledError:
            await self._interrupt_run(run_id)
            raise
        except Exception as exc:  # noqa: BLE001 - background task must persist failure
            logger.exception(f"批量测试运行失败: run_id={run_id}: {exc}")
            await self._fail_run_setup(run_id, str(exc))

    async def _execute_case(
        self,
        run_id: str,
        run_case_id: str,
        cfg,
        evaluator_resource: RuntimeModelResource,
        config: dict,
        semaphore: asyncio.Semaphore,
        control: _RunControl,
    ) -> None:
        async with semaphore:
            if control.stop_event.is_set():
                await self._mark_skipped(run_case_id)
                return

            async with SessionLocal() as session:
                row = await session.get(TestRunCase, run_case_id)
                if row is None:
                    return
                snapshot = dict(row.case_snapshot or {})
                definition = TestCaseDefinition.model_validate(snapshot)
                max_attempts = row.max_attempts

            for attempt in range(1, max_attempts + 1):
                if control.stop_event.is_set():
                    await self._mark_skipped(run_case_id)
                    return
                await self._mark_running(run_case_id, attempt)
                try:
                    timeout_seconds = int(
                        config.get("timeoutSecs")
                        or config.get("timeout_secs")
                        or 60
                    )
                    async with asyncio.timeout(timeout_seconds):
                        execution = await TextPipelineRunner().run(
                            cfg.model_copy(deep=True),
                            definition,
                            timeout_seconds=timeout_seconds,
                        )
                        evaluated = await EvaluationEngine(evaluator_resource).evaluate(
                            definition,
                            execution,
                        )
                except asyncio.CancelledError:
                    await self._mark_skipped(run_case_id)
                    raise
                except TimeoutError:
                    error = TestExecutionError(
                        code="CASE_TIMEOUT",
                        message=f"测试用例执行与评估超过 {timeout_seconds} 秒",
                        stage="pipeline",
                        retryable=True,
                    )
                    if attempt < max_attempts and not control.stop_event.is_set():
                        continue
                    await self._mark_error(run_case_id, error)
                    if str(
                        config.get("errorStrategy")
                        or config.get("error_strategy")
                    ) == "stop_on_error":
                        await self._request_stop(run_id, control, "execution_error")
                    return
                except TestExecutionError as exc:
                    if exc.retryable and attempt < max_attempts and not control.stop_event.is_set():
                        continue
                    await self._mark_error(run_case_id, exc)
                    if str(config.get("errorStrategy") or config.get("error_strategy")) == "stop_on_error":
                        await self._request_stop(run_id, control, "execution_error")
                    return
                except Exception as exc:  # noqa: BLE001 - normalize unknown runtime errors
                    error = TestExecutionError(
                        code="UNEXPECTED_EXECUTION_ERROR",
                        message=str(exc) or type(exc).__name__,
                        stage="pipeline",
                        retryable=False,
                    )
                    await self._mark_error(run_case_id, error)
                    if str(config.get("errorStrategy") or config.get("error_strategy")) == "stop_on_error":
                        await self._request_stop(run_id, control, "execution_error")
                    return

                status = "pass" if evaluated.passed else "fail"
                await self._mark_evaluated(run_case_id, status, evaluated)
                if status == "fail" and str(
                    config.get("failureStrategy") or config.get("failure_strategy")
                ) == "stop_on_fail":
                    await self._request_stop(run_id, control, "assertion_failure")
                return

    async def _request_stop(
        self,
        run_id: str,
        control: _RunControl,
        reason: str,
    ) -> None:
        if control.reason is None:
            control.reason = reason
        control.stop_event.set()
        async with SessionLocal() as session:
            run = await session.get(TestRun, run_id)
            if run is not None and run.stop_reason is None:
                run.stop_reason = reason
                await session.commit()

    async def _mark_running(self, run_case_id: str, attempt: int) -> None:
        async with SessionLocal() as session:
            row = await session.get(TestRunCase, run_case_id)
            if row is None:
                return
            row.status = "running"
            row.attempt_count = attempt
            row.started_at = row.started_at or datetime.now(UTC)
            row.finished_at = None
            row.execution_error = None
            await session.commit()

    async def _mark_evaluated(self, run_case_id: str, status: str, evaluated) -> None:
        async with SessionLocal() as session:
            row = await session.get(TestRunCase, run_case_id)
            if row is None:
                return
            row.status = status
            row.result = {
                "turns": [
                    item.model_dump(mode="json", by_alias=True)
                    for item in evaluated.turns
                ],
                "overallCriteria": [
                    item.model_dump(mode="json", by_alias=True)
                    for item in evaluated.overall_criteria
                ],
            }
            row.execution_error = None
            row.finished_at = datetime.now(UTC)
            if row.test_case_id:
                test_case = await session.get(TestCase, row.test_case_id)
                if test_case is not None:
                    test_case.last_result = status
            await session.commit()

    async def _mark_error(
        self,
        run_case_id: str,
        error: TestExecutionError,
    ) -> None:
        async with SessionLocal() as session:
            row = await session.get(TestRunCase, run_case_id)
            if row is None:
                return
            row.status = "error"
            row.execution_error = error.as_dict()
            row.finished_at = datetime.now(UTC)
            if row.test_case_id:
                test_case = await session.get(TestCase, row.test_case_id)
                if test_case is not None:
                    test_case.last_result = "fail"
            await session.commit()

    async def _mark_skipped(self, run_case_id: str) -> None:
        async with SessionLocal() as session:
            row = await session.get(TestRunCase, run_case_id)
            if row is None or row.status in {"pass", "fail", "error", "skipped"}:
                return
            row.status = "skipped"
            row.finished_at = datetime.now(UTC)
            await session.commit()

    async def _finish_run(self, run_id: str, control: _RunControl) -> None:
        async with SessionLocal() as session:
            run = await session.get(TestRun, run_id)
            if run is None:
                return
            waiting = (
                await session.execute(
                    select(TestRunCase).where(
                        TestRunCase.run_id == run_id,
                        TestRunCase.status.in_(["waiting", "running"]),
                    )
                )
            ).scalars().all()
            for row in waiting:
                row.status = "skipped"
                row.finished_at = datetime.now(UTC)
            manual = bool(run.cancel_requested or control.reason == "manual")
            run.status = "cancelled" if manual else "completed"
            run.stop_reason = "manual" if manual else control.reason or run.stop_reason
            run.finished_at = datetime.now(UTC)
            await session.commit()

    async def _fail_run_setup(self, run_id: str, message: str) -> None:
        error = TestExecutionError(
            code="RUN_SETUP_FAILED",
            message=message or "批量测试初始化失败",
            stage="pipeline",
            retryable=False,
        )
        async with SessionLocal() as session:
            run = await session.get(TestRun, run_id)
            if run is None:
                return
            rows = (
                await session.execute(
                    select(TestRunCase).where(
                        TestRunCase.run_id == run_id,
                        TestRunCase.status.in_(["waiting", "running"]),
                    )
                )
            ).scalars().all()
            now = datetime.now(UTC)
            for row in rows:
                row.status = "error"
                row.execution_error = error.as_dict()
                row.finished_at = now
            run.status = "completed"
            run.stop_reason = "execution_error"
            run.finished_at = now
            await session.commit()

    async def _interrupt_run(self, run_id: str) -> None:
        await _mark_run_interrupted(run_id, "后端服务停止，运行已中断")


async def _mark_run_interrupted(run_id: str, message: str) -> None:
    async with SessionLocal() as session:
        run = await session.get(TestRun, run_id)
        if run is None or run.status not in {"queued", "running"}:
            return
        rows = (
            await session.execute(
                select(TestRunCase).where(
                    TestRunCase.run_id == run_id,
                    TestRunCase.status.in_(["waiting", "running"]),
                )
            )
        ).scalars().all()
        now = datetime.now(UTC)
        for row in rows:
            if row.status == "running":
                row.status = "error"
                row.execution_error = {
                    "code": "SERVER_RESTARTED",
                    "message": message,
                    "stage": "pipeline",
                    "retryable": True,
                }
            else:
                row.status = "skipped"
            row.finished_at = now
        run.status = "completed"
        run.stop_reason = "execution_error"
        run.finished_at = now
        await session.commit()


async def recover_interrupted_test_runs() -> None:
    async with SessionLocal() as session:
        run_ids = (
            await session.execute(
                select(TestRun.id).where(TestRun.status.in_(["queued", "running"]))
            )
        ).scalars().all()
    for run_id in run_ids:
        await _mark_run_interrupted(str(run_id), "后端服务重启，运行已中断")


test_run_orchestrator = TestRunOrchestrator()