"""压测入口驱动的测试。**不打真实模型、不起容器**，模型那一侧全是写死脚本的替身。

分五块：参数校验、预算护栏与并发上限、产物与记分板的对接、错误隔离与取消、GovDoc 的阶段串行。

最有价值的一条是「产物喂给记分板判成全绿」。驱动写四个文件、记分板读四个文件，两边的字段
约定只写在记分板的模块 docstring 里，没有任何机器约束把它们钉在一起——那条测试就是那个约束：
少写一个字段、类型写错一个，它当场红。

GovDoc 那几块用手工搭的 `AuditTask`，不读磁盘上的真实语料：那份数据是另一个项目的工作副本，
换一台机器就没有，而这里要验的是编排（阶段串行、任务并发），跟语料内容无关。
"""

from __future__ import annotations

import asyncio
import json
from typing import TYPE_CHECKING

import pytest

from polyloop.types import ModelReply
from tools.soak import scoreboard
from tools.soak.run_soak import (
    BudgetGuard,
    RunOutcome,
    Unit,
    dispatch,
    execute_run,
    interleave,
    main,
    run_govdoc_task,
    select_appworld_tasks,
)
from tools.soak.scenarios import govdoc as govdoc_scenario

if TYPE_CHECKING:
    from collections.abc import Callable, Sequence
    from pathlib import Path

    from polyloop.ports import ModelCall

# ---------------------------------------------------------------------------
# 替身
# ---------------------------------------------------------------------------


class _ScriptedModel:
    """按 run_id 给回复的模型替身。每次调用都让出一次事件循环，好让并发真的发生。"""

    def __init__(self, *, render: Callable[[ModelCall], str]) -> None:
        self._render = render
        #: 调用到达的顺序，按 run_id 记。阶段串行与任务并发都从它上面判。
        self.seen: list[str] = []

    async def call(self, call: ModelCall) -> ModelReply:
        await asyncio.sleep(0.001)
        self.seen.append(call.run_id)
        return ModelReply(
            call_id=f"call-{len(self.seen)}",
            content=self._render(call),
            thinking="",
        )

    def parameters(self):  # noqa: ANN201 - 替身，签名由 Protocol 定
        return {"kind": "scripted"}


def _guard(*, render: Callable[[ModelCall], str], limit: int = 1000) -> BudgetGuard:
    return BudgetGuard(inner=_ScriptedModel(render=render), limit=limit)


_SUBMIT = json.dumps(
    {
        "tool": "submit_finding",
        "arguments": {
            "verdict": "存疑",
            "evidence": "tender.md 第 1 行：示例甲公司",
            "reasoning": "测试替身给的固定结论",
        },
    },
    ensure_ascii=False,
)

#: 解析不出来的输出。GovDoc 的预算里连续解析失败上限是 3，所以三步就停——plan 与 execute
#: 两个阶段没有结束通路（提交工具只在 summarize 那一阶段），不这样它们会一路走到 50 步。
_GARBAGE = "我先想一想这道题。"

#: summarize 阶段的第一步：先读一段再提交。
#:
#: **不是可有可无的一步。** 记分板的「提示词字符数单调不减」要两步才凑得出相邻一对，
#: 一步就提交的运行在它那里是「无从判起」而不是「通过」。这条测试断言的是整批全绿，
#: 所以替身必须走够两步——否则它测到的是记分板在数据不足时的行为，不是产物合不合约定。
_READ = json.dumps(
    {"tool": "read_document", "arguments": {"path": "tender.md", "start_line": 1, "end_line": 2}},
    ensure_ascii=False,
)


def _audit_task(index: int) -> govdoc_scenario.AuditTask:
    checkpoint = govdoc_scenario.Checkpoint(
        checkpoint_id=f"cp-{index}",
        category="资格条件",
        title="供应商资格要求",
        description="不得以注册地设置差别待遇。",
        legal_basis=("政府采购法第五条",),
        severity="高",
    )
    document = govdoc_scenario.CorpusDocument.from_text(
        logical_name="tender.md",
        text="第一行：示例甲公司参与投标。\n第二行：投标截止时间为示例日期。",
    )
    return govdoc_scenario.AuditTask(index=index, checkpoint=checkpoint, documents=(document,))


def _fake_unit(
    label: str,
    *,
    body: Callable[[], object],
) -> Unit:
    async def start() -> Sequence[RunOutcome]:
        result = body()
        if asyncio.iscoroutine(result):
            await result
        return [
            RunOutcome(
                run_id=label,
                scenario="fake",
                task_id=label,
                phase=None,
                stop_reason="task_completed",
                steps=1,
                model_calls=1,
                wall_ms=0,
            )
        ]

    return Unit(label=label, scenario="fake", run=start)


# ---------------------------------------------------------------------------
# 一、参数校验
# ---------------------------------------------------------------------------


def test_missing_budget_calls_refuses(tmp_path: Path) -> None:
    with pytest.raises(SystemExit) as caught:
        main(["--scenario", "govdoc", "--concurrency", "2", "--runs-dir", str(tmp_path)])
    assert caught.value.code != 0


def test_missing_concurrency_refuses(tmp_path: Path) -> None:
    with pytest.raises(SystemExit) as caught:
        main(["--scenario", "govdoc", "--budget-calls", "10", "--runs-dir", str(tmp_path)])
    assert caught.value.code != 0


def test_non_positive_budget_refuses(tmp_path: Path) -> None:
    with pytest.raises(SystemExit):
        main(
            [
                "--scenario",
                "govdoc",
                "--budget-calls",
                "0",
                "--concurrency",
                "1",
                "--runs-dir",
                str(tmp_path),
            ]
        )


def test_appworld_without_data_root_refuses(tmp_path: Path) -> None:
    with pytest.raises(SystemExit):
        main(
            [
                "--scenario",
                "appworld",
                "--budget-calls",
                "5",
                "--concurrency",
                "1",
                "--runs-dir",
                str(tmp_path),
                "--dry-run",
            ]
        )


# ---------------------------------------------------------------------------
# 二、预算护栏与并发上限
# ---------------------------------------------------------------------------


async def test_budget_counts_are_per_run_and_total() -> None:
    guard = _guard(render=lambda call: "ok")
    calls = [_call("run-a"), _call("run-a"), _call("run-b")]
    for item in calls:
        await guard.call(item)
    assert guard.total == 3
    assert guard.calls_for("run-a") == 2
    assert guard.calls_for("run-b") == 1
    assert guard.exhausted is False


async def test_budget_stops_dispatch_and_lets_running_tasks_finish() -> None:
    guard = _guard(render=lambda call: "ok", limit=2)

    async def slow() -> None:
        await guard.call(_call("slow"))
        await asyncio.sleep(0.05)

    async def quick() -> None:
        await guard.call(_call("quick"))

    units = [
        _fake_unit("unit-1", body=slow),
        _fake_unit("unit-2", body=quick),
        _fake_unit("unit-3", body=quick),
        _fake_unit("unit-4", body=quick),
    ]
    report = await dispatch(units, concurrency=2, guard=guard)

    assert report.dispatched == 2
    # 「第 N 个任务」从 1 数：第 3 个是第一个没派出去的。
    assert report.stopped_at == 3
    # 慢的那个是在预算耗尽时正在跑的，它必须跑完并留下结果，不许被砍断。
    assert {item.run_id for item in report.outcomes} == {"unit-1", "unit-2"}
    assert guard.total == 2


async def test_concurrency_never_exceeds_the_cap() -> None:
    live = 0
    peak = 0

    async def body() -> None:
        nonlocal live, peak
        live += 1
        peak = max(peak, live)
        await asyncio.sleep(0.005)
        live -= 1

    units = [_fake_unit(f"unit-{index}", body=body) for index in range(12)]
    report = await dispatch(units, concurrency=3)

    assert peak == 3
    assert report.dispatched == 12
    assert len(report.outcomes) == 12


async def test_zero_concurrency_is_rejected() -> None:
    with pytest.raises(ValueError, match="并发上限"):
        await dispatch([], concurrency=0)


def test_interleave_alternates_between_scenarios() -> None:
    left = [_fake_unit(f"L{index}", body=lambda: None) for index in range(3)]
    right = [_fake_unit(f"R{index}", body=lambda: None) for index in range(2)]
    assert [unit.label for unit in interleave([left, right])] == ["L0", "R0", "L1", "R1", "L2"]


def test_select_appworld_tasks_dedupes_across_splits() -> None:
    class _Pool:
        def list_task_ids(self, split: str) -> list[str]:
            return {"train": ["a", "b"], "dev": ["b", "c", "d"]}[split]

    assert select_appworld_tasks(_Pool(), splits=["train", "dev"], limit=None) == [
        "a",
        "b",
        "c",
        "d",
    ]
    assert select_appworld_tasks(_Pool(), splits=["train", "dev"], limit=3) == ["a", "b", "c"]


# ---------------------------------------------------------------------------
# 三、错误隔离与取消
# ---------------------------------------------------------------------------


async def test_one_failing_unit_does_not_stop_the_batch() -> None:
    def boom() -> None:
        raise RuntimeError("这一道题炸了")

    units = [
        _fake_unit("ok-1", body=lambda: None),
        _fake_unit("boom", body=boom),
        _fake_unit("ok-2", body=lambda: None),
    ]
    report = await dispatch(units, concurrency=1)

    assert {item.run_id for item in report.outcomes} == {"ok-1", "ok-2"}
    assert len(report.failures) == 1
    label, message = report.failures[0]
    assert label == "boom"
    assert "RuntimeError" in message and "这一道题炸了" in message


async def test_cancellation_passes_through_and_cleans_up() -> None:
    running = asyncio.Event()
    cleaned = False

    async def body() -> None:
        nonlocal cleaned
        running.set()
        try:
            await asyncio.sleep(10)
        finally:
            cleaned = True

    units = [_fake_unit("slow", body=body)]
    task = asyncio.create_task(dispatch(units, concurrency=1))
    await running.wait()
    task.cancel()

    with pytest.raises(asyncio.CancelledError):
        await task
    assert cleaned is True


# ---------------------------------------------------------------------------
# 四、产物：字段齐全，且记分板判得动
# ---------------------------------------------------------------------------


async def _run_summarize(tmp_path: Path, *, index: int = 0) -> None:
    """跑一次 summarize 阶段：模型第一步就提交结论，运行以 task_completed 结束。"""
    guard = _guard(render=lambda call: _SUBMIT)
    await run_govdoc_task(
        task=_audit_task(index),
        runs_dir=tmp_path / "runs",
        workspace_root=tmp_path / "workspaces",
        guard=guard,
        phases=("summarize",),
    )


async def test_sidecar_fields_are_complete_and_typed(tmp_path: Path) -> None:
    await _run_summarize(tmp_path)
    runs = tmp_path / "runs"
    run_id = govdoc_scenario.make_run_id(task_index=0, phase="summarize")

    assert (runs / f"{run_id}.jsonl").is_file()
    assert (runs / f"{run_id}.result.json").is_file()
    assert (runs / f"{run_id}.events.jsonl").is_file()

    meta = json.loads((runs / f"{run_id}.meta.json").read_text(encoding="utf-8"))
    assert set(meta) == {
        "scenario",
        "task_id",
        "phase",
        "wall_ms",
        "model_calls",
        "sink_failures",
        "env_executions",
        "fault",
        "success",
        "resumed_from_step",
    }
    assert meta["scenario"] == "govdoc"
    assert meta["task_id"] == "cp-0"
    assert meta["phase"] == "summarize"
    assert isinstance(meta["wall_ms"], int)
    assert meta["model_calls"] == 1
    assert meta["sink_failures"] == 0
    # submit_finding 往审计账里追加一行，那是环境自己记的账。
    assert meta["env_executions"] == 1
    assert meta["fault"] is None
    assert meta["success"] is None
    assert meta["resumed_from_step"] == 0

    events = [
        json.loads(line)
        for line in (runs / f"{run_id}.events.jsonl").read_text(encoding="utf-8").splitlines()
    ]
    assert events == [{"kind": "step_finished", "run_id": run_id, "step_idx": 0}]

    result = json.loads((runs / f"{run_id}.result.json").read_text(encoding="utf-8"))
    assert result["run_id"] == run_id
    assert result["stop_reason"] == "task_completed"
    assert result["event_delivery_failures"] == 0


async def test_artifacts_pass_the_scoreboard(tmp_path: Path) -> None:
    """把驱动写出来的产物直接喂给记分板，要它报全绿。

    两边的字段约定没有任何机器约束把它们钉在一起，这条测试就是那个约束。
    """

    def render(call: object) -> str:
        if "summarize" not in call.run_id:  # type: ignore[attr-defined]
            return _GARBAGE
        return _READ if call.call_index == 0 else _SUBMIT  # type: ignore[attr-defined]

    guard = _guard(render=render)
    for index in range(2):
        await run_govdoc_task(
            task=_audit_task(index),
            runs_dir=tmp_path / "runs",
            workspace_root=tmp_path / "workspaces",
            guard=guard,
        )

    board = scoreboard.evaluate(
        tmp_path / "runs",
        completing_tools=frozenset({"submit_finding"}),
    )
    trouble = [
        f"{item.name}: {[evidence.describe() for evidence in (*item.breaches, *item.undetermined)]}"
        for item in (*board.breached, *board.undetermined)
    ]
    assert board.verdict is scoreboard.Verdict.PASSED, trouble
    assert len(board.summaries) == 6
    assert {item.scenario for item in board.summaries} == {"govdoc"}


async def test_sink_failures_match_what_the_library_counted(tmp_path: Path) -> None:
    """事件文件写不出去时，出口自己数的失败次数与结果里那份必须一致。

    这条不变量是记分板的一条判定，而两边的计数由不同的代码写：出口在自己的 `except` 里加一，
    库在接住异常之后加一。只要出口漏加或多加，记分板当场报击穿。
    """
    runs = tmp_path / "runs"
    runs.mkdir()
    run_id = govdoc_scenario.make_run_id(task_index=0, phase="summarize")
    # 在事件文件该在的位置放一个目录，追加写就必然失败。
    (runs / f"{run_id}.events.jsonl").mkdir()

    guard = _guard(render=lambda call: _SUBMIT)
    await run_govdoc_task(
        task=_audit_task(0),
        runs_dir=runs,
        workspace_root=tmp_path / "workspaces",
        guard=guard,
        phases=("summarize",),
    )

    meta = json.loads((runs / f"{run_id}.meta.json").read_text(encoding="utf-8"))
    result = json.loads((runs / f"{run_id}.result.json").read_text(encoding="utf-8"))
    assert meta["sink_failures"] == 1
    assert result["event_delivery_failures"] == meta["sink_failures"]


async def test_a_failing_run_still_leaves_meta(tmp_path: Path) -> None:
    """`run` 抛异常也要留下 `.meta.json`——缺文件在记分板那边只是「无法判定」。

    这里用同一个运行标识跑第二次来触发失败：库对已经有日志的标识直接报 `RunIdentityError`，
    而那是压测里最可能真的撞上的一种失败（任务集去重漏了一处）。
    """
    runs = tmp_path / "runs"
    await _run_summarize(tmp_path)
    run_id = govdoc_scenario.make_run_id(task_index=0, phase="summarize")
    (runs / f"{run_id}.result.json").unlink()

    guard = _guard(render=lambda call: _SUBMIT)
    task = _audit_task(0)
    request = govdoc_scenario.build_run_request(
        task=task,
        phase="summarize",
        run_id=run_id,
        workspace=tmp_path / "ws",
        model_binding={},
    )
    outcome = await execute_run(
        runs_dir=runs,
        request=request,
        model_client=guard,
        decision_parser=govdoc_scenario.GovDocParser(),
        synthetic_observations=govdoc_scenario.SYNTHETIC_OBSERVATIONS,
        scenario="govdoc",
        task_id=task.checkpoint.checkpoint_id,
        phase="summarize",
        count_model_calls=lambda: guard.calls_for(run_id),
    )

    assert outcome.error is not None
    assert "RunIdentityError" in outcome.error
    assert outcome.stop_reason is None
    meta = json.loads((runs / f"{run_id}.meta.json").read_text(encoding="utf-8"))
    assert meta["model_calls"] == 0
    assert meta["success"] is None
    # 这一次没有结果，所以没有 `.result.json` 可写——记分板会把它记成一条要人看见的观察。
    assert not (runs / f"{run_id}.result.json").exists()


async def test_a_model_error_is_a_stop_reason_not_a_crash(tmp_path: Path) -> None:
    """模型调用失败不会把异常抛出循环，它是一次以 `llm_error` 结束的正常运行。

    压测的报告里这两件事必须分得开：`error` 是驱动这边出的事，`llm_error` 是库判定的停止原因。
    """

    class _Exploding:
        async def call(self, call: ModelCall) -> ModelReply:
            raise RuntimeError("网关炸了")

        def parameters(self):  # noqa: ANN202 - 替身
            return {"kind": "exploding"}

    guard = BudgetGuard(inner=_Exploding(), limit=5)
    task = _audit_task(0)
    request = govdoc_scenario.build_run_request(
        task=task,
        phase="summarize",
        run_id="govdoc-0-summarize",
        workspace=tmp_path / "ws",
        model_binding={},
    )
    outcome = await execute_run(
        runs_dir=tmp_path / "runs",
        request=request,
        model_client=guard,
        decision_parser=govdoc_scenario.GovDocParser(),
        synthetic_observations=govdoc_scenario.SYNTHETIC_OBSERVATIONS,
        scenario="govdoc",
        task_id=task.checkpoint.checkpoint_id,
        phase="summarize",
        count_model_calls=lambda: guard.calls_for("govdoc-0-summarize"),
    )

    assert outcome.error is None
    assert outcome.stop_reason == "llm_error"
    meta = json.loads(
        (tmp_path / "runs" / "govdoc-0-summarize.meta.json").read_text(encoding="utf-8")
    )
    assert meta["model_calls"] == 1
    assert (tmp_path / "runs" / "govdoc-0-summarize.result.json").is_file()


# ---------------------------------------------------------------------------
# 五、GovDoc：阶段串行、任务并发
# ---------------------------------------------------------------------------


async def test_govdoc_phases_run_in_order_within_a_task(tmp_path: Path) -> None:
    model = _ScriptedModel(render=lambda call: _GARBAGE)
    guard = BudgetGuard(inner=model, limit=1000)
    await run_govdoc_task(
        task=_audit_task(0),
        runs_dir=tmp_path / "runs",
        workspace_root=tmp_path / "workspaces",
        guard=guard,
    )

    phases = [run_id.rsplit("-", 1)[1] for run_id in model.seen]
    # 每个阶段的调用连成一段，段与段之间不交错。
    assert phases == sorted(phases, key=["plan", "execute", "summarize"].index)
    assert set(phases) == {"plan", "execute", "summarize"}


async def test_govdoc_tasks_overlap_while_phases_do_not(tmp_path: Path) -> None:
    model = _ScriptedModel(render=lambda call: _GARBAGE)
    guard = BudgetGuard(inner=model, limit=1000)
    units = [
        Unit(
            label=f"govdoc/{index}",
            scenario="govdoc",
            run=_govdoc_runner(
                index=index,
                runs_dir=tmp_path / "runs",
                workspace_root=tmp_path / "workspaces",
                guard=guard,
            ),
        )
        for index in range(2)
    ]
    report = await dispatch(units, concurrency=2, guard=guard)
    assert len(report.outcomes) == 6
    assert report.failures == ()

    # 阶段串行：同一个任务里，后一阶段的第一次调用晚于前一阶段的最后一次调用。
    for index in range(2):
        own = [
            position
            for position, run_id in enumerate(model.seen)
            if run_id.startswith(f"govdoc-{index}-")
        ]
        for earlier, later in (("plan", "execute"), ("execute", "summarize")):
            last_earlier = max(
                position for position in own if model.seen[position].endswith(earlier)
            )
            first_later = min(position for position in own if model.seen[position].endswith(later))
            assert last_earlier < first_later

    # 任务并发：两个任务的调用在时间上交错。
    owners = [run_id.split("-")[1] for run_id in model.seen]
    assert any(left != right for left, right in zip(owners, owners[1:], strict=False))


def _govdoc_runner(*, index: int, runs_dir: Path, workspace_root: Path, guard: BudgetGuard):  # noqa: ANN202
    async def start() -> Sequence[RunOutcome]:
        return await run_govdoc_task(
            task=_audit_task(index),
            runs_dir=runs_dir,
            workspace_root=workspace_root,
            guard=guard,
        )

    return start


# ---------------------------------------------------------------------------
# 小工具
# ---------------------------------------------------------------------------


def _call(run_id: str) -> ModelCall:
    from polyloop.ports import ModelCall as _ModelCall

    return _ModelCall(
        messages=(),
        call_index=0,
        run_id=run_id,
        result_id=f"{run_id}-0",
        binding={},
    )