"""断点续跑 checkpoint 的 unit_id 粒度契约集成测试（Task 10）。

覆盖三条铁律：
- checkpoint 存 **unit_id 序列**而非逐题 question_id（孪生对折叠为单个 unit_id）；
- 续跑恢复以 **完整 unit** 展开（pair 两成员同进同出，绝不被劈开）；
- 真实经 write_checkpoint → load_checkpoint 落盘往返后，unit_id 序列无损存活、
  重建的 batch 与原 batch 逐字节一致（core 算法保真 #3 断点续跑）。
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING

from app.harness.batching import build_batches
from app.harness.checkpoint import load_checkpoint, write_checkpoint
from app.harness.pools import Pools
from app.harness.question_units import build_units, validate_units
from app.harness.runner import _batch_from_ids, _batch_unit_ids
from core.types import GeneratedQuestion

if TYPE_CHECKING:
    from pathlib import Path

# ---------------------------------------------------------------------------
# 辅助构造
# ---------------------------------------------------------------------------


def _make_single(qid: str, task_type: str = "RETRIEVAL", video_id: str = "v1") -> GeneratedQuestion:
    """构造最小 single 题目。"""
    return GeneratedQuestion(
        question_id=qid,
        video_id=video_id,
        task_type=task_type,
        question=f"q_{qid}",
        options=("A. a", "B. b", "C. c", "D. d"),
        answer="A",
        source_nodes=("n1",),
        difficulty="medium",
        question_role="single",
    )


def _make_pair(
    pair_id: str, task_type: str = "AR", video_id: str = "v1"
) -> tuple[GeneratedQuestion, GeneratedQuestion]:
    """构造一个孪生对（original + mirror），共享 pair_id / flip_axis。"""
    common = {
        "video_id": video_id,
        "task_type": task_type,
        "question": "?",
        "options": ("A. a", "B. b", "C. c", "D. d"),
        "answer": "A",
        "source_nodes": ("n1",),
        "difficulty": "medium",
        "pair_id": pair_id,
        "flip_axis": "before_after",
    }
    original = GeneratedQuestion(
        question_id=f"{pair_id}_o", question_role="pair_original", **common
    )
    mirror = GeneratedQuestion(question_id=f"{pair_id}_m", question_role="pair_mirror", **common)
    return original, mirror


def _pools_with(diagnosis: list[GeneratedQuestion]) -> Pools:
    """用给定诊断池构造最小 Pools（其余池置空，_batch_from_ids 只读 diagnosis）。"""
    return Pools(
        diagnosis=diagnosis,
        validation=[],
        test=[],
        baseline_run_id="baseline",
        baseline_val_accuracy=0.0,
        correctness={q.question_id: False for q in diagnosis},
    )


# ---------------------------------------------------------------------------
# checkpoint 落盘所需的最小 _TrainState / RunConfig 替身
# ---------------------------------------------------------------------------


@dataclass
class _FakeState:
    """serialize_state 只读的可持久化字段（本测试用空累加包即可）。"""

    correctness: dict[str, bool] = field(default_factory=dict)
    eval_prev_acc: float = 0.0
    eval_prev_run_id: str = "run-0"
    baseline_skills_version: str = "v1"
    baseline_prompts_version: str = "v1"
    steps_since_best_improved: int = 0
    epoch_start_skills: str = "v1"
    changed_task_types_this_epoch: set[str] = field(default_factory=set)
    rejected_buffer: dict = field(default_factory=dict)
    system_packs: list = field(default_factory=list)
    tool_packs: list = field(default_factory=list)
    probations: dict = field(default_factory=dict)
    gate_cooldown: dict = field(default_factory=dict)
    gate_epoch_observed: dict = field(default_factory=dict)


@dataclass(frozen=True)
class _FakeConfig:
    """compute_fingerprint 读取的结构性 + 决策性字段。"""

    batch_size: int = 4
    min_class_per_batch: int = 2
    epochs: int = 3
    diag_size: int = 30
    val_size: int = 50
    batch_correct_ratio: float = 0.0
    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


# ---------------------------------------------------------------------------
# checkpoint 存 unit_id 序列（孪生对折叠为单个 unit_id）
# ---------------------------------------------------------------------------


class TestCheckpointStoresUnitIds:
    """一个含孪生对的 batch，checkpoint 序列应折叠为 unit_id（pair 只记一次）。"""

    def test_pair_folds_to_single_unit_id(self) -> None:
        po, pm = _make_pair("pA")
        s1, s2 = _make_single("s1"), _make_single("s2")
        batch = [s1, po, pm, s2]  # pair 两成员相邻

        unit_ids = _batch_unit_ids(batch)

        # 4 题折叠为 3 个 unit：s1 / pA(pair) / s2
        assert unit_ids == ["s1", "pA", "s2"]
        # 逐题 question_id 序列（旧行为）会含 4 项且把 pair 拆成两条
        assert [q.question_id for q in batch] == ["s1", "pA_o", "pA_m", "s2"]

    def test_pure_single_unit_ids_equal_question_ids(self) -> None:
        """纯 single batch 的 unit_id 序列与逐题 question_id 序列逐字节一致。"""
        batch = [_make_single(f"s{i}") for i in range(5)]
        assert _batch_unit_ids(batch) == [q.question_id for q in batch]


# ---------------------------------------------------------------------------
# 恢复以完整 unit 展开：pair 不拆
# ---------------------------------------------------------------------------


class TestResumeExpandsFullUnit:
    """按 unit_id 序列恢复时，孪生对整体展开、绝不被劈开。"""

    def test_pair_not_split_on_resume(self) -> None:
        po, pm = _make_pair("pA")
        s1, s2 = _make_single("s1"), _make_single("s2")
        batch = [s1, po, pm, s2]
        # 诊断池顺序故意与 batch 不同，验证恢复以池折叠的 unit 为准
        pools = _pools_with([s2, pm, s1, po])

        unit_ids = _batch_unit_ids(batch)
        restored = _batch_from_ids(pools, unit_ids)

        # 逐字节一致（顺序保留）
        assert [q.question_id for q in restored] == ["s1", "pA_o", "pA_m", "s2"]
        # 恢复后可无损重建为完整孪生对单元
        units = validate_units(build_units(restored))
        pair_units = [u for u in units if u.kind == "pair"]
        assert len(pair_units) == 1
        assert pair_units[0].unit_id == "pA"
        assert pair_units[0].size == 2

    def test_roundtrip_preserves_order_mixed(self) -> None:
        """混格 batch 经 unit_id 折叠 → 展开 round-trip 逐字节还原原顺序。"""
        po, pm = _make_pair("pA")
        qo, qm = _make_pair("pB")
        singles = [_make_single(f"s{i}") for i in range(3)]
        batch = [singles[0], po, pm, singles[1], qo, qm, singles[2]]
        pools = _pools_with(list(batch))

        restored = _batch_from_ids(pools, _batch_unit_ids(batch))
        assert [q.question_id for q in restored] == [q.question_id for q in batch]

    def test_missing_unit_id_raises(self) -> None:
        """checkpoint 引用了诊断池不存在的 unit_id 时硬失败（防静默兜底）。"""
        pools = _pools_with([_make_single("s1")])
        try:
            _batch_from_ids(pools, ["s1", "ghost"])
        except KeyError:
            return
        raise AssertionError("引用缺失 unit_id 应触发 KeyError")


# ---------------------------------------------------------------------------
# 真实落盘往返：write_checkpoint → load_checkpoint → 重建 batch
# ---------------------------------------------------------------------------


class TestCheckpointPersistenceRoundtrip:
    """unit_id 序列经真实 checkpoint.json 原子写/读回后无损、重建逐字节一致。"""

    def test_persisted_unit_ids_rebuild_batches(self, tmp_path: Path) -> None:
        po, pm = _make_pair("pA")
        singles = [_make_single(f"s{i}", task_type="RETRIEVAL") for i in range(6)]
        items = [*singles, po, pm]
        correctness = {q.question_id: False for q in items}
        batches, _ = build_batches(items, correctness, batch_size=4, min_class_per_batch=2, seed=7)
        pools = _pools_with(list(items))

        # 存 unit_id 序列（train() 落盘路径）
        epoch_batches = [_batch_unit_ids(b) for b in batches]
        write_checkpoint(
            tmp_path,
            state=_FakeState(correctness=correctness),
            epoch=1,
            step_completed=0,
            phase="in_epoch",
            global_step=1,
            total_steps=len(batches),
            version_snapshot={"skills": "skills/v1", "prompts": "prompts/v1"},
            epoch_batches=epoch_batches,
            config=_FakeConfig(),
        )

        ckpt = load_checkpoint(tmp_path)
        assert ckpt is not None
        rebuilt = [_batch_from_ids(pools, ids) for ids in ckpt["epoch_batches"]]

        # 落盘往返后逐批逐字节还原
        assert [[q.question_id for q in b] for b in rebuilt] == [
            [q.question_id for q in b] for b in batches
        ]
        # 含孪生对的批次恢复后 pair 仍成对
        for b in rebuilt:
            for u in validate_units(build_units(b)):
                if u.kind == "pair":
                    assert u.size == 2
