"""gate_ladder.py 单元化迁移测试（核心算法保真 #5）。

验证 gate 信息量阶梯从"逐题"迁移到"unit"粒度后，冷启动 2:1 错优先交错、
gamma-EMA、Beta 先验、schema_version 门控等语义"只换键、不改公式/比例/顺序"：

(a) LadderEntry 按 unit_id 键（AR pair 折叠为一个阶梯单元）；
(b) 冷启动"错优先 2:1"以 unit 为单位，unit 错 = P 或 Q 任一错；
(c) update_probs 观测先折叠成 unit 再匹配（防按 qid 匹配失效致 EMA 停摆）；
(d) GatePools.save/load 带 schema_version，存量无版本 json 明确报错（不静默混用）；
(e) BaselineCache 键含 unit_id（pair 的 unit_id 与成员 qid 不同）。
"""

from __future__ import annotations

import json
from typing import TYPE_CHECKING

import pytest

from app.harness.gate_ladder import (
    SCHEMA_VERSION,
    BaselineCache,
    GatePools,
    LadderEntry,
    build_cold_entries,
)
from app.harness.question_units import build_units
from core.types import GeneratedQuestion

if TYPE_CHECKING:
    from pathlib import Path


# ── 构造工具 ──────────────────────────────────────────────────────────


def _single(qid: str, task_type: str = "AR") -> GeneratedQuestion:
    """构造 single 题（unit_id 回填为 question_id）。"""
    return GeneratedQuestion(
        question_id=qid,
        video_id="v1",
        task_type=task_type,
        question="dummy",
        options=("A", "B", "C", "D"),
        answer="A",
        source_nodes=("n1",),
        difficulty="easy",
    )


def _pair(pair_id: str, task_type: str = "AR") -> list[GeneratedQuestion]:
    """构造一条 AR 孪生对（original + mirror），共享 pair_id 即 unit_id。"""
    base = {
        "video_id": "v1",
        "task_type": task_type,
        "question": "dummy",
        "options": ("A", "B", "C", "D"),
        "answer": "A",
        "source_nodes": ("n1",),
        "difficulty": "easy",
        "pair_id": pair_id,
        "flip_axis": "before_after",
    }
    return [
        GeneratedQuestion(question_id=f"{pair_id}_o", question_role="pair_original", **base),
        GeneratedQuestion(question_id=f"{pair_id}_m", question_role="pair_mirror", **base),
    ]


# ── (a) LadderEntry 按 unit_id ─────────────────────────────────────────


class TestLadderEntryKeyedByUnit:
    """LadderEntry 以 unit_id 为键；AR pair 折叠为一个阶梯单元。"""

    def test_entry_has_unit_id(self) -> None:
        """LadderEntry 暴露 unit_id 字段。"""
        e = LadderEntry("u1", 0.5)
        assert e.unit_id == "u1"

    def test_pair_collapses_to_single_entry(self) -> None:
        """一条孪生对（2 题）在阶梯中只产生 1 个 unit 条目（键=pair_id）。"""
        units = build_units(_pair("pr1"))
        correctness = {"pr1_o": True, "pr1_m": True}
        entries = build_cold_entries(units, correctness, probe_quota=0.0, seed=1)
        assert len(entries) == 1
        assert entries[0].unit_id == "pr1"


# ── (b) 冷启动 2:1 错优先 + unit 错 = P 或 Q 任一错 ───────────────────────


class TestColdStartUnit:
    """冷启动排序在 unit 粒度保持 2:1 交错与 Beta 先验；unit 错 = 任一成员错。"""

    def test_unit_wrong_if_any_member_wrong(self) -> None:
        """pair 中任一成员错 → 该 unit 判错（p_hat=1/3）；全对才判对（2/3）。"""
        units = build_units(_pair("wrong") + _pair("right"))
        # wrong: original 对、mirror 错 → unit 错；right: 两题均对 → unit 对
        correctness = {
            "wrong_o": True,
            "wrong_m": False,
            "right_o": True,
            "right_m": True,
        }
        entries = build_cold_entries(units, correctness, probe_quota=0.0, seed=3)
        p_by_unit = {e.unit_id: e.p_hat for e in entries}
        assert p_by_unit["wrong"] == pytest.approx(1 / 3)
        assert p_by_unit["right"] == pytest.approx(2 / 3)

    def test_two_to_one_interleave_over_units(self) -> None:
        """6 错 unit + 3 对 unit（含 pair），probe_quota=0 → 交错序 W W R W W R W W R。

        交错是 unit 粒度：pair 折叠成一个 unit 参与交错，序列长度为 9（unit 数），
        而非 18（题数），证明 2:1 比例语义只换键不改。
        """
        wrong_units_q: list[GeneratedQuestion] = []
        for i in range(6):
            wrong_units_q += _pair(f"w{i}")  # 6 个 pair unit
        right_units_q = [_single(f"r{i}") for i in range(3)]  # 3 个 single unit
        units = build_units(wrong_units_q + right_units_q)

        correctness: dict[str, bool] = {}
        for i in range(6):
            correctness[f"w{i}_o"] = True
            correctness[f"w{i}_m"] = False  # 任一错 → unit 错
        for i in range(3):
            correctness[f"r{i}"] = True

        entries = build_cold_entries(units, correctness, probe_quota=0.0, seed=42)
        assert len(entries) == 9
        # unit 错 = p_hat≈1/3；unit 对 = p_hat≈2/3
        pattern = ["W" if e.p_hat < 0.5 else "R" for e in entries]
        assert pattern == ["W", "W", "R", "W", "W", "R", "W", "W", "R"]

    def test_probe_at_tail_unit(self) -> None:
        """probe_quota>0 时从错 unit 抽探针追加梯尾（按 unit 抽，非按题）。"""
        wrong_q: list[GeneratedQuestion] = []
        for i in range(10):
            wrong_q += _pair(f"w{i}")
        right_q = [_single(f"r{i}") for i in range(2)]
        units = build_units(wrong_q + right_q)
        correctness = {}
        for i in range(10):
            correctness[f"w{i}_o"] = False
            correctness[f"w{i}_m"] = False
        for i in range(2):
            correctness[f"r{i}"] = True

        entries = build_cold_entries(units, correctness, probe_quota=0.3, seed=7)
        # 10 错 unit * 0.3 = 3 个探针 unit 在尾部，均为错 unit
        tail = entries[-3:]
        for e in tail:
            assert e.p_hat < 0.5


# ── (c) update_probs 折叠成 unit 再匹配 ──────────────────────────────────


class TestUpdateProbsFold:
    """gamma-EMA 更新前先把逐题观测折叠成 unit 观测（AR pair 双向 AND）。"""

    def _pools_and_units(self) -> tuple[GatePools, dict]:
        """构造含一个 pair unit + 一个 single unit 的池与 unit 索引。"""
        units = build_units(_pair("pr1") + [_single("s1")])
        units_by_id = {u.unit_id: u for u in units}
        entries = {"AR": [LadderEntry("pr1", 0.5), LadderEntry("s1", 0.5)]}
        return GatePools(entries=entries, seed=0, fingerprint="x"), units_by_id

    def test_pair_updates_after_fold(self) -> None:
        """pair 两成员均观测 → 折叠成 unit 观测 → EMA 更新（不停摆）。"""
        pools, units_by_id = self._pools_and_units()
        # pair 两题均对 → unit 对(1.0)；single 错(0.0)
        per_q = {"pr1_o": True, "pr1_m": True, "s1": False}
        pools.update_probs(per_q, units_by_id, gamma=0.8)
        p = {e.unit_id: e.p_hat for e in pools.entries["AR"]}
        assert p["pr1"] == pytest.approx(0.8 * 0.5 + 0.2 * 1.0)  # 0.6
        assert p["s1"] == pytest.approx(0.8 * 0.5 + 0.2 * 0.0)  # 0.4

    def test_pair_wrong_if_any_member_wrong(self) -> None:
        """pair 任一成员错 → unit 观测为错(0.0)，EMA 向下。"""
        pools, units_by_id = self._pools_and_units()
        per_q = {"pr1_o": True, "pr1_m": False, "s1": True}
        pools.update_probs(per_q, units_by_id, gamma=0.8)
        p = {e.unit_id: e.p_hat for e in pools.entries["AR"]}
        assert p["pr1"] == pytest.approx(0.8 * 0.5 + 0.2 * 0.0)  # 0.4

    def test_partial_pair_observation_skips_update(self) -> None:
        """pair 只观测到半个成员 → 无法折叠 → 该 unit p_hat 不变（不半 pair 污染）。"""
        pools, units_by_id = self._pools_and_units()
        per_q = {"pr1_o": True}  # 缺 mirror
        pools.update_probs(per_q, units_by_id, gamma=0.8)
        p = {e.unit_id: e.p_hat for e in pools.entries["AR"]}
        assert p["pr1"] == pytest.approx(0.5)  # 未更新

    def test_qid_keyed_observation_does_not_match_pair(self) -> None:
        """若观测里错用 pair_id 之外的裸 qid 键、且未提供 unit 折叠，pair 不应被误更新。

        证明"必须折叠"：单 single unit 用其自身 qid 可更新，pair 需 unit 折叠。
        """
        pools, units_by_id = self._pools_and_units()
        # 只给 single 的观测，pair 两成员均无观测
        per_q = {"s1": True}
        pools.update_probs(per_q, units_by_id, gamma=0.8)
        p = {e.unit_id: e.p_hat for e in pools.entries["AR"]}
        assert p["pr1"] == pytest.approx(0.5)  # pair 未观测 → 不变
        assert p["s1"] == pytest.approx(0.8 * 0.5 + 0.2 * 1.0)  # single 更新


# ── (d) schema_version 门控 ─────────────────────────────────────────────


class TestSchemaVersion:
    """GatePools.save/load 带 schema_version；存量无版本 json 明确报错。"""

    def test_save_writes_schema_version_and_unit_id(self, tmp_path: Path) -> None:
        """save 落盘含 schema_version 且 entries 用 unit_id 键。"""
        entries = {"AR": [LadderEntry("pr1", 0.33), LadderEntry("s1", 0.67)]}
        pools = GatePools(entries=entries, seed=1, fingerprint="fp")
        path = tmp_path / "gate_pools.json"
        pools.save(path)

        raw = json.loads(path.read_text(encoding="utf-8"))
        assert raw["schema_version"] == SCHEMA_VERSION
        assert raw["entries"]["AR"][0]["unit_id"] == "pr1"
        assert "question_id" not in raw["entries"]["AR"][0]

    def test_save_load_roundtrip(self, tmp_path: Path) -> None:
        """save → load 往返保真（unit_id + p_hat + seed + fingerprint）。"""
        entries = {"AR": [LadderEntry("pr1", 0.4)]}
        pools = GatePools(entries=entries, seed=9, fingerprint="fp2")
        path = tmp_path / "gate_pools.json"
        pools.save(path)
        loaded = GatePools.load(path)
        assert loaded.seed == 9
        assert loaded.fingerprint == "fp2"
        assert loaded.entries["AR"][0].unit_id == "pr1"
        assert loaded.entries["AR"][0].p_hat == pytest.approx(0.4)

    def test_load_legacy_without_schema_version_raises(self, tmp_path: Path) -> None:
        """存量无 schema_version（旧 qid 键）→ 明确报错，不静默混用。"""
        legacy = {
            "seed": 1,
            "fingerprint": "fp",
            "entries": {"AR": [{"question_id": "q1", "p_hat": 0.5}]},
        }
        path = tmp_path / "gate_pools.json"
        path.write_text(json.dumps(legacy), encoding="utf-8")
        with pytest.raises(RuntimeError, match="schema_version"):
            GatePools.load(path)

    def test_load_wrong_schema_version_raises(self, tmp_path: Path) -> None:
        """schema_version 不匹配 → 明确报错。"""
        bad = {
            "schema_version": SCHEMA_VERSION + 99,
            "seed": 1,
            "fingerprint": "fp",
            "entries": {"AR": [{"unit_id": "pr1", "p_hat": 0.5}]},
        }
        path = tmp_path / "gate_pools.json"
        path.write_text(json.dumps(bad), encoding="utf-8")
        with pytest.raises(RuntimeError, match="schema_version"):
            GatePools.load(path)


# ── (e) BaselineCache 键含 unit_id ─────────────────────────────────────


class TestBaselineCacheUnitKey:
    """BaselineCache 以 unit_id 为第四维；pair 的 unit_id 与成员 qid 区分。"""

    def test_unit_id_key_distinct_from_member_qid(self, tmp_path: Path) -> None:
        """pair unit_id（=pair_id）与其成员 qid 是不同缓存键。"""
        cache = BaselineCache(tmp_path / "baseline_cache.json")
        cache.put("AR", "h1", "v1", "pr1", True)  # unit_id=pair_id
        assert cache.get("AR", "h1", "v1", "pr1") is True
        # 成员 qid 不是同一键 → miss
        assert cache.get("AR", "h1", "v1", "pr1_o") is None
        assert cache.get("AR", "h1", "v1", "pr1_m") is None


# ── ladder_for 排除按 unit ─────────────────────────────────────────────


class TestLadderForExcludeUnit:
    """ladder_for 返回 unit_id 序，按 unit 排除（防半 pair 灌入）。"""

    def test_exclude_units_filters_whole_unit(self) -> None:
        """exclude_units 命中的 unit 被整体排除，返回 unit_id 列表。"""
        entries = {"AR": [LadderEntry("pr1", 0.5), LadderEntry("s1", 0.4), LadderEntry("s2", 0.6)]}
        pools = GatePools(entries=entries, seed=0, fingerprint="x")
        result = pools.ladder_for("AR", exclude_units={"pr1"}, p_low=0.0, p_high=1.0, cold=True)
        assert "pr1" not in result
        assert "s1" in result
        assert "s2" in result