"""崩溃往返等价：在每一个写入边界上崩一次，续跑的结果必须和一口气跑完的一样。

`research-wiki/design/0002-step-level-resume.md` 末尾点名要这类测试，理由是**恢复的 bug 天然
没有失败现场**——它不会抛异常，只表现成「跑出来的结果有点不一样」，而一次运行本来就次次不同，
所以人眼分不出来。这里把每一个写入边界都崩一遍，逐个比对。

**它跑真实文件系统，但仍然算 unit。** 分层判据是「依赖什么」（`CLAUDE.md` §1.9）：模型、
解释器、执行器全是替身，没有连任何外部服务。等 `adapters` 落地、真的连上模型网关时，
`tests/integration/` 与 `tests/e2e/` 那两层才有内容。
"""

import dataclasses
from collections.abc import Mapping
from pathlib import Path

import pytest

from polyloop.ports import Action, ModelCall, ParsedReply, ToolCall
from polyloop.session import AgentDefinition, RunRequest, resume, run
from polyloop.stores import JsonlRunStore
from polyloop.tools import ToolRegistry, ToolSpec
from polyloop.types import (
    ActionOutcome,
    ActionStatus,
    Budget,
    Context,
    Message,
    ModelReply,
    ReplayPolicy,
    Role,
    RunResult,
    StopReason,
    SyntheticObservations,
    TextBlock,
)

pytestmark = pytest.mark.unit

#: 一次两步的运行写十次：运行开始 + 每步四次 + 运行结束。
TOTAL_WRITES = 10

SYNTHETIC = SyntheticObservations(
    action_rejected="被拒绝了", env_failed="环境坏了", model_call_failed="调用失败了"
)


class _CrashError(RuntimeError):
    """模拟进程被杀。它不是存储故障，是「这个进程不存在了」。"""


class _CrashingStore:
    """包住真存储，放过前 N 次写，第 N+1 次直接抛。

    崩在「第 N 次写落了盘、第 N+1 次还没开始」这个位置上——那正是四次写之间的每一个边界。
    """

    def __init__(self, inner: JsonlRunStore, *, crash_after: int) -> None:
        self._inner = inner
        self._crash_after = crash_after
        self.writes = 0
        #: 逐条记下写的是哪种记录，用来验 TOTAL_WRITES 那个常量不是拍脑袋的。
        self.kinds: list[str] = []

    def _tick(self, record: object) -> None:
        self.writes += 1
        self.kinds.append(type(record).__name__)
        if self.writes > self._crash_after:
            raise _CrashError(f"第 {self.writes} 次写之前进程没了")

    async def write_run_started(self, record) -> None:
        self._tick(record)
        await self._inner.write_run_started(record)

    async def write_intent(self, record) -> None:
        self._tick(record)
        await self._inner.write_intent(record)

    async def write_model_call_result(self, record) -> None:
        self._tick(record)
        await self._inner.write_model_call_result(record)

    async def write_step_completed(self, record) -> None:
        self._tick(record)
        await self._inner.write_step_completed(record)

    async def write_run_finished(self, record) -> None:
        self._tick(record)
        await self._inner.write_run_finished(record)

    async def read_log(self, run_id: str):
        return await self._inner.read_log(run_id)

    def parameters(self) -> Mapping[str, str]:
        return self._inner.parameters()


class _ScriptedModel:
    """按**调用序号**回复，不按调用次数。

    按次数的话，一次重放会拿到下一条回复，于是「重放」变成了「往下走一步」，而这条测试要验的
    正是重放拿到的是同一条。
    """

    def __init__(self) -> None:
        self.calls: list[int] = []
        #: 每次调用收到的那份消息序列，摊平成文本。恢复出来的历史对不对就靠比它。
        self.seen: list[list[str]] = []

    async def call(self, call: ModelCall) -> ModelReply:
        self.calls.append(call.call_index)
        self.seen.append([block.text for message in call.messages for block in message.content])
        return ModelReply(
            call_id=f"c{call.call_index}", content=f"go{call.call_index}", thinking=""
        )

    def parameters(self) -> Mapping[str, str]:
        return {"model": "scripted"}


class _ScriptedParser:
    def parse(self, reply: ModelReply) -> ParsedReply:
        index = int(reply.content.removeprefix("go"))
        return ParsedReply(
            history_text=reply.content,
            decision=Action(
                text=f"a{index}", tool_call=ToolCall(name=f"t{index}", arguments={"i": index})
            ),
        )

    def parameters(self) -> Mapping[str, str]:
        return {"parser": "scripted"}


class _ScriptedExecutor:
    """按动作文本查表，所以重放同一个动作拿到同一个结果。"""

    def __init__(self) -> None:
        self.executed: list[str] = []

    async def execute(self, action: Action) -> ActionOutcome:
        self.executed.append(action.text)
        return ActionOutcome(
            status=ActionStatus.EXECUTED,
            observation=f"{action.text} 的输出",
            observation_is_synthetic=False,
            # 第二步把这次运行做完。
            env_reported_completion=action.text == "a1",
            observation_truncated_chars=0,
        )

    def parameters(self) -> Mapping[str, str]:
        return {"env": "scripted"}


class _Sink:
    async def emit(self, event) -> None: ...

    def parameters(self) -> Mapping[str, str]:
        return {"sink": "none"}


def _registry(*, replay_policy: ReplayPolicy) -> ToolRegistry:
    return ToolRegistry(
        [
            ToolSpec(name=f"t{index}", description="", parameters={}, replay_policy=replay_policy)
            for index in range(2)
        ]
    )


def _definition(store) -> AgentDefinition:
    return AgentDefinition(
        model_client=_ScriptedModel(),
        decision_parser=_ScriptedParser(),
        store=store,
        event_sink=_Sink(),
        synthetic_observations=SYNTHETIC,
    )


def _request(executor, registry: ToolRegistry, *, model_replay: ReplayPolicy) -> RunRequest:
    return RunRequest(
        run_id="round-trip",
        budget=Budget(
            max_steps=5, max_actions=5, max_consecutive_parse_failures=3, max_prompt_chars=100000
        ),
        action_executor=executor,
        tools=registry,
        context=Context(
            run_level=(Message(role=Role.USER, content=(TextBlock(text="你是助手"),)),),
            goal_level=(Message(role=Role.USER, content=(TextBlock(text="数到二"),)),),
        ),
        injections={},
        model_binding={"item": "x"},
        model_replay_policy=model_replay,
        observation_template="观察：{observation}",
        cancel_grace_seconds=1.0,
    )


def _comparable(result: RunResult) -> RunResult:
    """把墙钟时间抹平。

    它是唯一一个次次不同的字段——一次运行的其余每一列都该是确定的，否则「等价」就没法断言。
    """
    return dataclasses.replace(
        result,
        steps=tuple(dataclasses.replace(step, step_wall_ms=0) for step in result.steps),
    )


async def _run_uninterrupted(directory: Path, policy: ReplayPolicy) -> RunResult:
    store = JsonlRunStore(directory=directory)
    registry = _registry(replay_policy=policy)
    return await run(
        _definition(store), _request(_ScriptedExecutor(), registry, model_replay=policy)
    )


@pytest.mark.parametrize("crash_after", range(1, TOTAL_WRITES + 1))
async def test_crashing_at_any_write_boundary_resumes_to_the_same_result(
    tmp_path: Path, crash_after: int
) -> None:
    """一次两步的运行写十次，在每一次之后崩一遍，续跑的结果都必须和一口气跑完的一样。

    工具与模型调用都声明可重放，所以每一个断点都该续得上——声明不可重放的那些断点会正当地停在
    「状态未知」，那是另一条测试。
    """
    expected = _comparable(await _run_uninterrupted(tmp_path / "clean", ReplayPolicy.SAFE))
    directory = tmp_path / "crashed"
    registry = _registry(replay_policy=ReplayPolicy.SAFE)

    crashing = _CrashingStore(JsonlRunStore(directory=directory), crash_after=crash_after)
    try:
        await run(
            _definition(crashing),
            _request(_ScriptedExecutor(), registry, model_replay=ReplayPolicy.SAFE),
        )
    except _CrashError:
        pass
    else:
        # 崩点排在最后一次写之后，这次运行其实跑完了。
        assert crash_after == TOTAL_WRITES

    resumed = await resume(
        _definition(JsonlRunStore(directory=directory)),
        _request(_ScriptedExecutor(), registry, model_replay=ReplayPolicy.SAFE),
    )

    assert _comparable(resumed) == expected


async def test_the_history_the_model_sees_after_resuming_matches_the_uninterrupted_one(
    tmp_path: Path,
) -> None:
    """恢复出来的消息历史不能比不中断跑完时多一轮或少一轮。

    多一轮少一轮都不会报错，只会让模型在续跑之后看见的东西和原本不一样，然后每一步都跟着偏。
    这里比的是「第二步那次模型调用收到的消息序列」。
    """
    registry = _registry(replay_policy=ReplayPolicy.SAFE)

    clean_model = _ScriptedModel()
    clean = AgentDefinition(
        model_client=clean_model,
        decision_parser=_ScriptedParser(),
        store=JsonlRunStore(directory=tmp_path / "clean"),
        event_sink=_Sink(),
        synthetic_observations=SYNTHETIC,
    )
    await run(clean, _request(_ScriptedExecutor(), registry, model_replay=ReplayPolicy.SAFE))

    directory = tmp_path / "crashed"
    # 崩在第一步的步记录落盘之后（第 5 次写），第二步一个字都没写。
    crashing = _CrashingStore(JsonlRunStore(directory=directory), crash_after=5)
    with pytest.raises(_CrashError):
        await run(
            _definition(crashing),
            _request(_ScriptedExecutor(), registry, model_replay=ReplayPolicy.SAFE),
        )

    resumed_model = _ScriptedModel()
    resumed = AgentDefinition(
        model_client=resumed_model,
        decision_parser=_ScriptedParser(),
        store=JsonlRunStore(directory=directory),
        event_sink=_Sink(),
        synthetic_observations=SYNTHETIC,
    )
    await resume(resumed, _request(_ScriptedExecutor(), registry, model_replay=ReplayPolicy.SAFE))

    assert resumed_model.calls == [1]  # 续跑只补第二步
    # 不中断那次的第二次调用，和续跑那次的唯一一次调用，收到的必须是同一份消息序列。
    assert resumed_model.seen[0] == clean_model.seen[1]


async def test_an_unreplayable_action_stops_at_unknown_instead_of_guessing(
    tmp_path: Path,
) -> None:
    """动作声明「绝不重放」而恰好崩在动作意图与步记录之间，续跑停在「状态未知」。

    这不是缺陷，是设计的目的：库不替谁决定要不要把一次可能已经发生的副作用再做一遍。
    「跑到第几步、已经花了多少、前面那些步的轨迹」都在结果里，项目拿它决定重跑还是人工介入。
    """
    directory = tmp_path / "never"
    registry = _registry(replay_policy=ReplayPolicy.NEVER)
    # 第 4 次写是第一步的动作意图；崩在它之后，步记录就没写。
    crashing = _CrashingStore(JsonlRunStore(directory=directory), crash_after=4)
    with pytest.raises(_CrashError):
        await run(
            _definition(crashing),
            _request(_ScriptedExecutor(), registry, model_replay=ReplayPolicy.NEVER),
        )

    result = await resume(
        _definition(JsonlRunStore(directory=directory)),
        _request(_ScriptedExecutor(), registry, model_replay=ReplayPolicy.NEVER),
    )

    assert result.stop_reason is StopReason.RESUME_STATE_UNKNOWN
    assert result.steps == ()


async def test_an_unreplayable_model_call_stops_at_unknown(tmp_path: Path) -> None:
    """模型调用意图写了、结果没写，而调用方声明这次调用不能重来。

    调用可能已经发出去、也可能没有。库不猜——发出去了就已经花了钱、已经在网关那边记了账。
    """
    directory = tmp_path / "never-model"
    registry = _registry(replay_policy=ReplayPolicy.NEVER)
    # 第 2 次写是第一步的模型调用意图。
    crashing = _CrashingStore(JsonlRunStore(directory=directory), crash_after=2)
    with pytest.raises(_CrashError):
        await run(
            _definition(crashing),
            _request(_ScriptedExecutor(), registry, model_replay=ReplayPolicy.NEVER),
        )

    result = await resume(
        _definition(JsonlRunStore(directory=directory)),
        _request(_ScriptedExecutor(), registry, model_replay=ReplayPolicy.NEVER),
    )

    assert result.stop_reason is StopReason.RESUME_STATE_UNKNOWN


async def test_a_resumed_log_survives_a_second_resume(tmp_path: Path) -> None:
    """续跑写下去的东西不能让下一次续跑读不动。

    重放要是又写一条同种意图，恢复会判成「日志被并发写过」，于是一次成功的重放反倒把日志弄坏，
    而那件事要到下一次崩溃之后才被发现。
    """
    directory = tmp_path / "twice"
    registry = _registry(replay_policy=ReplayPolicy.SAFE)
    crashing = _CrashingStore(JsonlRunStore(directory=directory), crash_after=2)
    with pytest.raises(_CrashError):
        await run(
            _definition(crashing),
            _request(_ScriptedExecutor(), registry, model_replay=ReplayPolicy.SAFE),
        )

    # 第一次续跑：重放那次模型调用，再崩在下一步的动作意图之后。
    second = _CrashingStore(JsonlRunStore(directory=directory), crash_after=3)
    with pytest.raises(_CrashError):
        await resume(
            _definition(second),
            _request(_ScriptedExecutor(), registry, model_replay=ReplayPolicy.SAFE),
        )

    # 第二次续跑：日志必须还读得动，而且跑得完。
    result = await resume(
        _definition(JsonlRunStore(directory=directory)),
        _request(_ScriptedExecutor(), registry, model_replay=ReplayPolicy.SAFE),
    )

    assert result.stop_reason is StopReason.TASK_COMPLETED
    assert [step.step_idx for step in result.steps] == [0, 1]


async def test_the_write_sequence_is_exactly_what_the_crash_matrix_assumes(tmp_path: Path) -> None:
    """把上面那个崩溃矩阵依赖的常量验一遍，顺便钉住四次写的顺序与收尾。

    `TOTAL_WRITES` 要是和实际写入次数对不上，矩阵就会漏掉最后几个边界而没有任何人看得见。
    更要紧的是最后那一条：**结束标记必须在把结果交给调用方之前写下**。漏写它的话，上面每一条
    往返测试照样会绿——恢复会把最后一步之后那次停止判定重演一遍，得出同样的结果——所以那件事
    只能在这里单独钉。
    """
    counting = _CrashingStore(JsonlRunStore(directory=tmp_path), crash_after=10**6)

    await run(
        _definition(counting),
        _request(
            _ScriptedExecutor(),
            _registry(replay_policy=ReplayPolicy.SAFE),
            model_replay=ReplayPolicy.SAFE,
        ),
    )

    assert counting.writes == TOTAL_WRITES
    assert counting.kinds == [
        "RunStarted",
        # 第一步：模型意图 → 模型结果 → 动作意图 → 步记录（后两者一次原子落地）
        "Intent",
        "ModelCallResult",
        "Intent",
        "StepCompleted",
        # 第二步同上
        "Intent",
        "ModelCallResult",
        "Intent",
        "StepCompleted",
        "RunFinished",
    ]