import asyncio
import json
import runpy
from datetime import UTC, datetime
from pathlib import Path
from types import SimpleNamespace
from uuid import UUID, uuid4

import pytest
from langchain_core.messages import HumanMessage, message_to_dict
from test_after_sales import scenario  # noqa: F401
from test_agent_runtime import tool_run

from zhifu.evaluation import AgentManifest, digest, write_report
from zhifu.modules.agent_runtime.config import GRAPH_VERSION, SYSTEM_PROMPT, AgentSettings
from zhifu.modules.agent_runtime.models import AgentRun
from zhifu.modules.agent_runtime.runtime import build_tool_graph
from zhifu.modules.agent_runtime.worker import claim

runner = runpy.run_path("scripts/evaluate_agent_runs.py")


@pytest.mark.integration
async def test_persistent_call_limit_survives_restart_and_preserves_output_budget(scenario):  # noqa: F811
    path = Path(".work") / f"allowance-validation-{uuid4().hex}.json"
    review = {"reviewer": "字段校验", "reviewed_at": datetime.now(UTC),
              "evidence_path": "字段校验不作为人工审核记录", "evidence_sha256": "a" * 64}
    manifest = AgentManifest(dataset_sha256="a" * 64, source_sha256="b" * 64,
        split="development", release_sha256="c" * 64, graph_version=GRAPH_VERSION,
        model_name="validation", max_calls=3, max_output_tokens=128,
        max_total_output_tokens=256, total_timeout_seconds=30,
        cases=[{"case_id": "validation", "tenant_id": uuid4(), "run_id": uuid4(),
                "review": review, "environment_evidence": review}])
    state = scenario["app"].state.state
    args = SimpleNamespace(output=path, execute=True)
    dataset = SimpleNamespace(version="validation")
    identity = digest(manifest.model_dump(mode="json"))
    budget_key = "evaluation-budget:" + identity

    async def reserve_one():
        lock = state.lock("evaluation:" + identity, timeout=30, blocking_timeout=5)
        async with lock:
            report = await runner["initialize_report"](args, dataset, manifest, {}, state)
            await asyncio.to_thread(write_report, path, report)
            await runner["PersistentAllowance"](manifest, report, path, state).take()
            return report["reserved_calls"]

    try:
        args.execute = False
        initial = await runner["initialize_report"](args, dataset, manifest, {}, state)
        assert initial["deadline"] == 0
        await asyncio.to_thread(write_report, path, initial)
        args.execute = True
        assert sorted(await asyncio.gather(reserve_one(), reserve_one())) == [1, 2]
        assert int(await state.hget(budget_key, "calls")) == 2
        with pytest.raises(ValueError, match="限额"):
            await reserve_one()
        args.output = path.with_suffix(".other.json")
        with pytest.raises(ValueError, match="绑定其他"):
            await runner["initialize_report"](args, dataset, manifest, {}, state)
        args.output = path
        await asyncio.to_thread(path.unlink)
        with pytest.raises(ValueError, match="报告文件缺失"):
            await runner["initialize_report"](args, dataset, manifest, {}, state)
    finally:
        await state.delete(budget_key, "evaluation:" + identity)
        await asyncio.to_thread(path.unlink, missing_ok=True)


@pytest.mark.integration
async def test_targeted_claim_and_real_tool_observation(scenario):  # noqa: F811
    sessions = scenario["app"].state.sessions
    first = await tool_run(scenario, "get_order", {"order_id": str(scenario["order_id"])})
    target = await tool_run(scenario, "get_order", {"order_id": str(scenario["order_id"])})
    settings = AgentSettings()
    snapshot = {"model_name": settings.model_name, "model_url": settings.model_url,
                "graph_version": GRAPH_VERSION, "system_prompt": SYSTEM_PROMPT,
                "max_output_tokens": 512}
    prompt = "查询当前订单"
    async with sessions.begin() as session:
        run = await session.get(AgentRun, target.id)
        run.model_name, run.graph_version = settings.model_name, GRAPH_VERSION
        run.release_snapshot = snapshot
        run.messages = [message_to_dict(HumanMessage(content=prompt)), *run.messages]
    assert await claim(sessions, 2, run_ids=[]) is None
    run_id, token = await claim(sessions, 2, run_ids=[target.id])
    assert run_id == target.id
    async with asyncio.timeout(10):
        await build_tool_graph(sessions).ainvoke({"run_id": run_id, "token": token})
    binding = SimpleNamespace(run_id=UUID(bytes=target.id), tenant_id=scenario["tenant"])
    manifest = SimpleNamespace(graph_version=GRAPH_VERSION, model_name=settings.model_name,
                               release_sha256=digest(snapshot), max_output_tokens=512)
    observation = await runner["observe"](sessions, binding, manifest, {"prompt": prompt})
    assert observation["attempts"] == []
    assert observation["messages"][-1]["type"] == "tool"
    assert json.loads(observation["messages"][-1]["data"]["content"])["order_id"] == str(
        scenario["order_id"])
    async with sessions() as session:
        assert (await session.get(AgentRun, first.id)).status == "QUEUED"
    manifest.release_sha256 = "0" * 64
    with pytest.raises(ValueError, match="发布参数"):
        await runner["observe"](sessions, binding, manifest, {"prompt": prompt})