"""Read-only diagnosis for a run stuck on one execution."""

import argparse
import asyncio
import json
import time

from app.core.config import get_settings
from app.db.repositories.repositories import ResultRepository, TestRunRepository
from app.db.session import SessionLocal
from app.providers.model_provider import ProviderFactory
from app.services.dataset_service import DatasetGateway
from app.services.test_execution_service import TestExecutionService


def arguments():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--run-id", type=int, default=7)
    parser.add_argument("--execution-id", default="R0229")
    parser.add_argument(
        "--timeout", type=float, default=60, help="Each model probe timeout in seconds"
    )
    return parser.parse_args()


async def timed(label, awaitable, timeout):
    started = time.monotonic()
    try:
        result = await asyncio.wait_for(awaitable, timeout)
    except Exception as exc:
        elapsed = time.monotonic() - started
        print(f"{label}: FAIL after {elapsed:.2f}s: {type(exc).__name__}: {exc}")
        return None
    print(f"{label}: OK in {time.monotonic() - started:.2f}s")
    return result


async def main():
    args = arguments()
    settings = get_settings()
    case = next(
        (
            item
            for item in DatasetGateway(settings).load()
            if item["execution_id"] == args.execution_id
        ),
        None,
    )
    if not case:
        raise SystemExit(f"Unknown execution_id: {args.execution_id}")

    with SessionLocal() as db:
        run = TestRunRepository(db).get(args.run_id)
        if not run:
            raise SystemExit(f"Unknown run_id: {args.run_id}")
        rows = ResultRepository(db).list_by_run(args.run_id)
        print(
            "DATABASE:",
            json.dumps(
                {
                    "status": run.status,
                    "selected": run.selected_count,
                    "stored_results": len(rows),
                    "execution_stored": any(
                        row.execution_id == args.execution_id for row in rows
                    ),
                    "summary": json.loads(run.summary_json or "{}"),
                },
                ensure_ascii=False,
            ),
        )

        model_input = case["model_input"]
        messages = []
        system_parts = [model_input["system"]] if model_input.get("system") else []
        system_parts.extend(
            "测试资料：" + json.dumps(item, ensure_ascii=False)
            for item in model_input.get("attachments", [])
            if item["type"] != "image"
        )
        if system_parts:
            messages.append({"role": "system", "content": "\n\n".join(system_parts)})
        messages.extend(model_input["messages"])
        tools = [model_input["tool_schema"]] if model_input.get("tool_schema") else None

        factory = ProviderFactory(settings, db)
        target = factory.create("target")
        judge = factory.create("judge")
        print(f"TARGET: endpoint={target.endpoint} model={target.model}")
        target_reply = await timed(
            "TARGET PROBE", target.chat(messages, tools=tools), args.timeout
        )
        if target_reply is None:
            print("DIAGNOSIS: 目标模型请求链失败或超时。")
            return
        model_response = target_reply.content or (
            json.dumps({"tool_calls": target_reply.tool_calls}, ensure_ascii=False)
            if target_reply.tool_calls
            else ""
        )
        print(
            "TARGET OUTPUT:",
            json.dumps(
                {
                    "content_length": len(target_reply.content),
                    "tool_call_count": len(target_reply.tool_calls),
                    "preview": model_response[:500],
                },
                ensure_ascii=False,
            ),
        )

        print(f"JUDGE: endpoint={judge.endpoint} model={judge.model}")
        service = TestExecutionService(
            db=db, target_provider=target, judge_provider=judge
        )
        judge_result = await timed(
            "JUDGE PROBE", service._judge(case, model_response), args.timeout
        )
        if judge_result is None:
            print("DIAGNOSIS: 目标模型正常，仲裁模型请求链失败或超时。")
            return
        print("JUDGE OUTPUT:", json.dumps(judge_result, ensure_ascii=False))
        print(
            "DIAGNOSIS: 两个模型直连均正常；若运行仍卡住，"
            "优先检查后端进程重启、后台任务丢失或间歇性上游故障。"
        )


if __name__ == "__main__":
    asyncio.run(main())