"""app/harness/checkpoint.py 单元测试。

覆盖序列化/反序列化往返、嵌套 dataclass 复活、缺键硬失败、
配置指纹结构性 vs 决策性判定、原子写与 load 缺失场景。
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

import pytest

if TYPE_CHECKING:
    from pathlib import Path

from app.harness.checkpoint import (
    CHECKPOINT_SCHEMA_VERSION,
    Probation,
    check_fingerprint,
    compute_fingerprint,
    deserialize_state_fields,
    load_checkpoint,
    serialize_state,
    write_checkpoint,
)
from core.evolution.types import (
    CaseSample,
    RejectedEdit,
    SystemCasePack,
    ToolCasePack,
)

# ---------------------------------------------------------------------------
# fixtures: 模拟 _TrainState 与 RunConfig
# ---------------------------------------------------------------------------


def _make_case_sample(**overrides: Any) -> CaseSample:
    """构造一个最小可用 CaseSample。"""
    defaults: dict[str, Any] = {
        "question_id": "q001",
        "video_id": "v001",
        "task_type": "temporal",
        "question": "What happened?",
        "options": ["A", "B", "C"],
        "answer": "A",
        "prediction": "B",
        "correct": False,
        "error_type": "reasoning",
        "selection_reason": "worst",
        "metrics": {"acc": 0.5},
        "trace": [{"step": 1, "action": "search"}],
    }
    defaults.update(overrides)
    return CaseSample(**defaults)


def _make_rejected_edit(**overrides: Any) -> RejectedEdit:
    """构造一个最小可用 RejectedEdit。"""
    defaults: dict[str, Any] = {
        "target_file": "temporal-reasoning.md",
        "target_type": "skill",
        "change_summary": "added step",
        "delta": -0.05,
        "source_version": "v2",
        "epoch": 1,
        "gate_w": 3,
        "gate_l": 5,
        "gate_e_value": 0.8,
        "gate_delta_shrunk": -0.02,
    }
    defaults.update(overrides)
    return RejectedEdit(**defaults)


def _make_system_pack() -> SystemCasePack:
    """构造包含嵌套 CaseSample 的 SystemCasePack。"""
    return SystemCasePack(
        stats={"pattern": "repeat_visit", "count": 3},
        failure_cases=[_make_case_sample(question_id="q010")],
        success_cases=[_make_case_sample(question_id="q011", correct=True, error_type=None)],
    )


def _make_tool_pack() -> ToolCasePack:
    """构造 ToolCasePack。"""
    return ToolCasePack(
        tool_name="search_subtree",
        target_files=["search_subtree_extract.md"],
        stats={"completeness": 0.8},
        failure_spans=[{"step": 2, "issue": "missing"}],
        success_spans=[{"step": 3, "quality": "good"}],
    )


def _make_probation() -> Probation:
    """构造包含嵌套 RejectedEdit 的 Probation。"""
    return Probation(
        task_type="temporal",
        anchor_skills_version="v1",
        target_file="temporal-reasoning.md",
        correctness_snapshot={"q001": True, "q002": False},
        opened_step=5,
        pending_edits=[_make_rejected_edit()],
    )


@dataclass
class _FakeState:
    """模拟 _TrainState 全部可持久化字段。"""

    correctness: dict[str, bool]
    eval_prev_acc: float
    eval_prev_run_id: str
    baseline_skills_version: str
    baseline_prompts_version: str
    steps_since_best_improved: int
    epoch_start_skills: str
    changed_task_types_this_epoch: set[str]
    rejected_buffer: dict[str, list[RejectedEdit]]
    system_packs: list[SystemCasePack]
    tool_packs: list[ToolCasePack]
    probations: dict[str, Probation]
    gate_cooldown: dict[str, int]
    gate_epoch_observed: dict[str, bool]


def _make_state() -> _FakeState:
    """构造一个填满全部字段的 _FakeState。"""
    return _FakeState(
        correctness={"q001": True, "q002": False},
        eval_prev_acc=0.65,
        eval_prev_run_id="run-abc",
        baseline_skills_version="v1",
        baseline_prompts_version="v1",
        steps_since_best_improved=2,
        epoch_start_skills="v1",
        changed_task_types_this_epoch={"temporal", "causal"},
        rejected_buffer={"temporal": [_make_rejected_edit()]},
        system_packs=[_make_system_pack()],
        tool_packs=[_make_tool_pack()],
        probations={"temporal": _make_probation()},
        gate_cooldown={"temporal": 3},
        gate_epoch_observed={"temporal": True},
    )


@dataclass(frozen=True)
class _FakeConfig:
    """模拟 RunConfig 的指纹相关字段。"""

    batch_size: int = 8
    min_class_per_batch: int = 2
    epochs: int = 5
    diag_size: int = 30
    val_size: int = 50
    batch_correct_ratio: float = 0.5
    edit_budget_start: int = 6
    edit_budget_end: int = 3
    early_stop_patience: int = 3
    use_slow_momentum: bool = True
    skill_update_mode: str = "patch"
    appendix_consolidate_threshold: int = 10
    momentum_samples: int = 20
    gate_e_confirm: float = 20.0
    gate_e_provisional: float = 6.0
    gate_w_net_min: int = 2
    gate_delta_min: float = 0.02
    gate_lambda_dir: float = -3.0
    gate_e_rollback: float = 10.0
    gate_block: int = 4
    gate_n_max: int = 40
    gate_p_low: float = 0.1
    gate_p_high: float = 0.9
    gate_probe_quota: float = 0.2
    gate_gamma_decay: float = 0.9
    gate_cooldown_steps: int = 2
    gate_guard_err: float = 0.3


# =========================================================================
# 测试用例
# =========================================================================


class TestSerializeDeserializeRoundtrip:
    """序列化 → JSON 往返 → 反序列化应完全复原。"""

    def test_serialize_deserialize_roundtrip(self) -> None:
        state = _make_state()
        serialized = serialize_state(state)
        # JSON 往返（模拟实际落盘-读回）
        json_str = json.dumps(serialized, ensure_ascii=False)
        loaded = json.loads(json_str)
        restored = deserialize_state_fields(loaded)

        assert restored["correctness"] == state.correctness
        assert restored["eval_prev_acc"] == state.eval_prev_acc
        assert restored["eval_prev_run_id"] == state.eval_prev_run_id
        assert restored["baseline_skills_version"] == state.baseline_skills_version
        assert restored["baseline_prompts_version"] == state.baseline_prompts_version
        assert restored["steps_since_best_improved"] == state.steps_since_best_improved
        assert restored["epoch_start_skills"] == state.epoch_start_skills
        assert restored["changed_task_types_this_epoch"] == state.changed_task_types_this_epoch
        assert restored["gate_cooldown"] == state.gate_cooldown
        assert restored["gate_epoch_observed"] == state.gate_epoch_observed


class TestSerializeSetToSortedList:
    """set 字段序列化为排序列表。"""

    def test_serialize_set_to_sorted_list(self) -> None:
        state = _make_state()
        state.changed_task_types_this_epoch = {"z_type", "a_type", "m_type"}
        serialized = serialize_state(state)
        assert serialized["changed_task_types_this_epoch"] == ["a_type", "m_type", "z_type"]


class TestDeserializeNestedSystemPack:
    """SystemCasePack 内嵌套的 CaseSample 正确复活。"""

    def test_deserialize_nested_system_pack(self) -> None:
        state = _make_state()
        serialized = serialize_state(state)
        json_str = json.dumps(serialized, ensure_ascii=False)
        loaded = json.loads(json_str)
        restored = deserialize_state_fields(loaded)

        packs = restored["system_packs"]
        assert len(packs) == 1
        pack = packs[0]
        assert isinstance(pack, SystemCasePack)
        assert len(pack.failure_cases) == 1
        assert isinstance(pack.failure_cases[0], CaseSample)
        assert pack.failure_cases[0].question_id == "q010"
        assert len(pack.success_cases) == 1
        assert isinstance(pack.success_cases[0], CaseSample)
        assert pack.success_cases[0].question_id == "q011"


class TestDeserializeNestedProbation:
    """Probation 内嵌套的 RejectedEdit 正确复活。"""

    def test_deserialize_nested_probation(self) -> None:
        state = _make_state()
        serialized = serialize_state(state)
        json_str = json.dumps(serialized, ensure_ascii=False)
        loaded = json.loads(json_str)
        restored = deserialize_state_fields(loaded)

        probations = restored["probations"]
        assert "temporal" in probations
        prob = probations["temporal"]
        assert isinstance(prob, Probation)
        assert prob.task_type == "temporal"
        assert prob.anchor_skills_version == "v1"
        assert prob.correctness_snapshot == {"q001": True, "q002": False}
        assert len(prob.pending_edits) == 1
        edit = prob.pending_edits[0]
        assert isinstance(edit, RejectedEdit)
        assert edit.target_file == "temporal-reasoning.md"
        assert edit.delta == -0.05


class TestDeserializeMissingKeyRaises:
    """缺键即 checkpoint 损坏，应硬失败。"""

    def test_deserialize_missing_key_raises(self) -> None:
        state = _make_state()
        serialized = serialize_state(state)
        del serialized["gate_epoch_observed"]
        with pytest.raises(KeyError):
            deserialize_state_fields(serialized)


class TestFingerprintStructuralVsDecision:
    """compute_fingerprint 包含全部结构性 + 决策性键。"""

    def test_fingerprint_structural_vs_decision(self) -> None:
        config = _FakeConfig()
        fp = compute_fingerprint(config)

        structural = {
            "batch_size",
            "min_class_per_batch",
            "epochs",
            "diag_size",
            "val_size",
            "batch_correct_ratio",
        }
        decision = {
            "edit_budget_start",
            "edit_budget_end",
            "early_stop_patience",
            "use_slow_momentum",
            "skill_update_mode",
            "appendix_consolidate_threshold",
            "momentum_samples",
            "gate_e_confirm",
            "gate_e_provisional",
            "gate_w_net_min",
            "gate_delta_min",
            "gate_lambda_dir",
            "gate_e_rollback",
            "gate_block",
            "gate_n_max",
            "gate_p_low",
            "gate_p_high",
            "gate_probe_quota",
            "gate_gamma_decay",
            "gate_cooldown_steps",
            "gate_guard_err",
        }
        assert structural | decision == set(fp.keys())
        assert fp["batch_size"] == 8
        assert fp["gate_e_confirm"] == 20.0


class TestCheckFingerprintStructuralReject:
    """结构性键变化应出现在 structural 列表中。"""

    def test_check_fingerprint_structural_reject(self) -> None:
        config_old = _FakeConfig()
        saved = compute_fingerprint(config_old)
        # 修改结构性参数
        config_new = _FakeConfig(batch_size=16, epochs=10)
        structural, decision = check_fingerprint(saved, config_new)
        assert "batch_size" in structural
        assert "epochs" in structural
        assert len(decision) == 0


class TestCheckFingerprintDecisionWarn:
    """决策性键变化应出现在 decision 列表中，structural 为空。"""

    def test_check_fingerprint_decision_warn(self) -> None:
        config_old = _FakeConfig()
        saved = compute_fingerprint(config_old)
        config_new = _FakeConfig(early_stop_patience=10, gate_e_confirm=50.0)
        structural, decision = check_fingerprint(saved, config_new)
        assert len(structural) == 0
        assert "early_stop_patience" in decision
        assert "gate_e_confirm" in decision


class TestWriteCheckpointAtomic:
    """原子写：先 .tmp 再 os.replace。"""

    def test_write_checkpoint_atomic(self, tmp_path: Path) -> None:
        state = _make_state()
        config = _FakeConfig()
        write_checkpoint(
            tmp_path,
            state=state,
            epoch=2,
            step_completed=5,
            phase="in_epoch",
            global_step=15,
            total_steps=40,
            version_snapshot={"skills": "v3", "prompts": "v2"},
            epoch_batches=[["q001", "q002"], ["q003"]],
            config=config,
        )
        ckpt_path = tmp_path / "checkpoint.json"
        assert ckpt_path.exists()
        # .tmp 应已被 os.replace 移除
        assert not (tmp_path / "checkpoint.json.tmp").exists()

        payload = json.loads(ckpt_path.read_text())
        assert payload["schema_version"] == CHECKPOINT_SCHEMA_VERSION
        assert payload["progress"]["epoch"] == 2
        assert payload["progress"]["step_completed"] == 5
        assert payload["progress"]["phase"] == "in_epoch"
        assert payload["progress"]["global_step"] == 15
        assert payload["progress"]["total_steps"] == 40
        assert payload["version_snapshot"] == {"skills": "v3", "prompts": "v2"}
        assert payload["epoch_batches"] == [["q001", "q002"], ["q003"]]
        assert "config_fingerprint" in payload
        assert "state" in payload


class TestLoadCheckpointMissing:
    """checkpoint.json 不存在时返回 None。"""

    def test_load_checkpoint_missing(self, tmp_path: Path) -> None:
        result = load_checkpoint(tmp_path)
        assert result is None

    def test_load_checkpoint_exists(self, tmp_path: Path) -> None:
        """checkpoint.json 存在时正确读回。"""
        state = _make_state()
        config = _FakeConfig()
        write_checkpoint(
            tmp_path,
            state=state,
            epoch=1,
            step_completed=3,
            phase="post_evolve",
            global_step=8,
            total_steps=20,
            version_snapshot={"skills": "v2", "prompts": "v1"},
            epoch_batches=[["q001"]],
            config=config,
        )
        loaded = load_checkpoint(tmp_path)
        assert loaded is not None
        assert loaded["schema_version"] == CHECKPOINT_SCHEMA_VERSION
        assert loaded["progress"]["epoch"] == 1
        # 完整往返测试：state 可 deserialize
        restored = deserialize_state_fields(loaded["state"])
        assert restored["eval_prev_acc"] == 0.65
