"""CE-Gate 信息量阶梯与基线缓存。

阶梯（每题型一条）：gate 的出题顺序表。冷启动（FRESH）用种子基线对错
两档粗排（错题高优先 2:1 交错 + 全错题 probe_quota 探针插尾）；
epoch >=1 用非 gate run 观测做 gamma-EMA 更新 p_hat，按信息量 p_hat(1-p_hat) 降序、
剔 p_hat 不在 [p_low, p_high]。防泄露铁律：gate 内 rollout 永不回流 p_hat
（调用方以 run_id 含 "_gate_" 过滤观测源）。

BaselineCache：基线侧逐题对错缓存，键 = (task_type, skill_hash,
prompts_version, qid) 内容寻址、无显式失效。JSON 持久化到 workspace，
供 resume 后合法复用已冻结阶梯上的新鲜 draw。
"""

from __future__ import annotations

import hashlib
import json
import os
import random
from dataclasses import dataclass
from typing import TYPE_CHECKING

from loguru import logger

if TYPE_CHECKING:
    from pathlib import Path

    from core.types import GeneratedQuestion


def skill_hash(content: str) -> str:
    """对 skill 正文取 sha1 摘要，作缓存键的内容维度。

    参数:
        content: skill 文件全文（基线侧为解析后生效文件的正文）。

    返回:
        sha1 十六进制摘要。
    """
    return hashlib.sha1(content.encode("utf-8")).hexdigest()


@dataclass
class LadderEntry:
    """阶梯单元：题目与其估计答对率。

    字段:
        question_id: 题目唯一标识。
        p_hat: 估计答对率。冷启动为 Beta(1,1) 平滑的单次观测后验均值
            （错=1/3、对=2/3），此后经 gamma-EMA 更新。
    """

    question_id: str
    p_hat: float


def build_cold_entries(
    questions: list[GeneratedQuestion],
    correctness: dict[str, bool],
    probe_quota: float,
    seed: int,
) -> list[LadderEntry]:
    """冷启动排序：错题高优先 2:1 交错 + 全错题 probe_quota 探针插尾。

    参数:
        questions: 该题型的全部候选题（已排除 test 池）。
        correctness: question_id -> 种子基线是否答对（900 题全量对错）。
        probe_quota: 从错题中随机抽出插到梯尾的探针比例（防"解锁新能力"盲区）。
        seed: 洗牌种子，保证确定性重建。

    返回:
        排序后的 LadderEntry 列表（p_hat 用 Beta(1,1) 平滑：错=1/3、对=2/3，
        与 warm 阶段 gamma-EMA / 信息量排序自然衔接）。

    关键实现细节:
        错题、对题各自固定种子洗牌 -> 抽探针 -> 剩余按 错错对 2:1 交错
        （一方耗尽后顺排另一方）-> 探针追加尾部。
    """
    rng = random.Random(seed)
    wrong = [q for q in questions if not correctness.get(q.question_id, False)]
    right = [q for q in questions if correctness.get(q.question_id, False)]
    rng.shuffle(wrong)
    rng.shuffle(right)

    n_probe = int(len(wrong) * probe_quota)
    probes, wrong_main = wrong[:n_probe], wrong[n_probe:]

    interleaved: list[GeneratedQuestion] = []
    wi, ri = 0, 0
    while wi < len(wrong_main) or ri < len(right):
        for _ in range(2):
            if wi < len(wrong_main):
                interleaved.append(wrong_main[wi])
                wi += 1
        if ri < len(right):
            interleaved.append(right[ri])
            ri += 1
    interleaved.extend(probes)

    def _p0(q: GeneratedQuestion) -> float:
        return 2 / 3 if correctness.get(q.question_id, False) else 1 / 3

    return [LadderEntry(q.question_id, _p0(q)) for q in interleaved]


def order_ladder(entries: list[LadderEntry], p_low: float, p_high: float) -> list[LadderEntry]:
    """warm 排序：剔 p_hat 不在 [p_low, p_high] 的零信息题，按信息量 p_hat(1-p_hat) 降序。

    参数:
        entries: 待排序的阶梯单元。
        p_low / p_high: p_hat 保留区间。

    返回:
        过滤并排序后的新列表（稳定排序，同信息量保持原相对序）。
    """
    kept = [e for e in entries if p_low <= e.p_hat <= p_high]
    return sorted(kept, key=lambda e: e.p_hat * (1 - e.p_hat), reverse=True)


@dataclass
class GatePools:
    """全部题型的阶梯容器，含构建种子与数据指纹（确定性重建凭据）。

    字段:
        entries: task_type -> 冷启动序 LadderEntry 列表（warm 排序在取用时做，
            保持存储序稳定、避免每次更新重写全表顺序）。
        seed: 冷启动洗牌种子。
        fingerprint: 构建输入指纹（基线 run_id + 题集 hash 等），resume 校验用。
    """

    entries: dict[str, list[LadderEntry]]
    seed: int
    fingerprint: str

    def ladder_for(
        self,
        task_type: str,
        exclude_qids: set[str],
        p_low: float,
        p_high: float,
        cold: bool,
    ) -> list[str]:
        """取该题型的 gate 出题序（qid 列表），排除本 step 进化案例包题。

        参数:
            task_type: 目标题型。
            exclude_qids: 本 step 案例包（failure/success cases）的题目 id，
                防止在"刚学的那道题"上自测。
            p_low / p_high: warm 阶段的 p_hat 保留区间。
            cold: True 表示尚无 epoch 级观测（epoch 1），用冷启动存储序；
                False 走 order_ladder 信息量排序。

        返回:
            排除后的有序 question_id 列表。

        异常:
            ValueError: 该题型无阶梯（冷启动构建缺失），或该题型阶梯为空。
        """
        if task_type not in self.entries:
            raise ValueError(f"task_type={task_type} 无阶梯，冷启动构建缺失该题型")
        pool = self.entries[task_type]
        if not pool:
            raise ValueError(f"task_type={task_type} 阶梯为空，无可出题目")
        ordered = pool if cold else order_ladder(pool, p_low, p_high)
        return [e.question_id for e in ordered if e.question_id not in exclude_qids]

    def update_probs(self, observations: dict[str, bool], gamma: float) -> None:
        """gamma-EMA 更新 p_hat：p_hat <- gamma * p_hat + (1-gamma) * obs。只更新有新观测的题。

        参数:
            observations: question_id -> 本 epoch 非 gate run 的最新对错。
                调用方必须已按 run_id 过滤掉 gate 内 rollout（防泄露铁律）。
            gamma: EMA 衰减系数。
        """
        for entries in self.entries.values():
            for e in entries:
                if e.question_id in observations:
                    obs = 1.0 if observations[e.question_id] else 0.0
                    e.p_hat = gamma * e.p_hat + (1 - gamma) * obs

    def save(self, path: Path) -> None:
        """原子写 gate_pools.json（.tmp 再 replace）。

        参数:
            path: 目标 JSON 路径。
        """
        payload = {
            "seed": self.seed,
            "fingerprint": self.fingerprint,
            "entries": {
                t: [{"question_id": e.question_id, "p_hat": e.p_hat} for e in es]
                for t, es in self.entries.items()
            },
        }
        tmp = path.with_suffix(".json.tmp")
        tmp.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
        os.replace(tmp, path)

    @classmethod
    def load(cls, path: Path) -> GatePools:
        """从 gate_pools.json 恢复。

        参数:
            path: gate_pools.json 路径。

        返回:
            复活的 GatePools。
        """
        d = json.loads(path.read_text(encoding="utf-8"))
        return cls(
            entries={
                t: [LadderEntry(x["question_id"], x["p_hat"]) for x in es]
                for t, es in d["entries"].items()
            },
            seed=d["seed"],
            fingerprint=d["fingerprint"],
        )


def build_or_load_gate_pools(
    workspace_dir: Path,
    questions: list[GeneratedQuestion],
    test_qids: set[str],
    baseline_correctness: dict[str, bool],
    task_types: list[str],
    probe_quota: float,
    seed: int,
    baseline_run_id: str,
) -> GatePools:
    """gate 阶梯获取入口：gate_pools.json 存在且指纹一致则加载，否则冷启动构建。

    参数:
        workspace_dir: workspace 根目录（gate_pools.json 落其下）。
        questions: benchmark 全量题（900 题）。
        test_qids: held-out test 池题目 id（阶梯题源必须排除）。
        baseline_correctness: 种子基线 900 题全量对错（从基线 run 的 db 读）。
        task_types: 参与进化的题型列表。
        probe_quota: 冷启动探针比例。
        seed: 冷启动洗牌种子。
        baseline_run_id: 指纹成分。

    返回:
        GatePools。

    关键实现细节:
        指纹 = sha1(baseline_run_id|全 qid|seed|probe_quota|task_types|test_qids)。
        指纹不一致（题集/基线/参数变了）直接报错——FRESH 语义下不该发生，
        防御性拒绝而非静默重建。
    """
    joined = ",".join(sorted(q.question_id for q in questions))
    fp_src = (
        f"{baseline_run_id}|{joined}|{seed}|{probe_quota}"
        f"|{','.join(sorted(task_types))}|{','.join(sorted(test_qids))}"
    )
    fingerprint = hashlib.sha1(fp_src.encode()).hexdigest()
    path = workspace_dir / "gate_pools.json"
    if path.exists():
        pools = GatePools.load(path)
        if pools.fingerprint != fingerprint:
            raise RuntimeError(
                f"gate_pools.json 指纹不一致（题集或基线变更），拒绝静默重建: {path}"
            )
        return pools

    entries: dict[str, list[LadderEntry]] = {}
    for t in task_types:
        pool = [q for q in questions if q.task_type == t and q.question_id not in test_qids]
        if not pool:
            raise ValueError(f"task_type={t} 无非 test 题，无法建阶梯")
        entries[t] = build_cold_entries(pool, baseline_correctness, probe_quota, seed)
        logger.info("gate 阶梯[{}]: {} 题（冷启动）", t, len(entries[t]))
    pools = GatePools(entries=entries, seed=seed, fingerprint=fingerprint)
    pools.save(path)
    return pools


class BaselineCache:
    """基线侧逐题对错缓存（内容寻址，JSON 持久化）。

    键 = (task_type, skill_hash, prompts_version, qid)：任何影响该题型
    有效 skill 的变化（含共享 default-strategy.md 被他类 accept 改写）
    都使 skill_hash 变化、缓存自然 miss；prompts 版本变化同理。
    """

    def __init__(self, path: Path) -> None:
        """加载或初始化缓存文件。

        参数:
            path: 缓存 JSON 路径（workspace/baseline_cache.json）。
        """
        self._path = path
        self._store: dict[str, bool] = {}
        if path.exists():
            self._store = json.loads(path.read_text(encoding="utf-8"))

    @staticmethod
    def _key(task_type: str, s_hash: str, prompts_version: str, qid: str) -> str:
        """拼缓存键（四维内容寻址）。"""
        return f"{task_type}|{s_hash}|{prompts_version}|{qid}"

    def get(self, task_type: str, s_hash: str, prompts_version: str, qid: str) -> bool | None:
        """读缓存；未命中返回 None。

        参数:
            task_type: 题型。
            s_hash: 基线侧生效 skill 文件的内容哈希。
            prompts_version: 当前 prompts 版本。
            qid: 题目 id。

        返回:
            缓存的对错；未命中 None。
        """
        return self._store.get(self._key(task_type, s_hash, prompts_version, qid))

    def put(
        self, task_type: str, s_hash: str, prompts_version: str, qid: str, correct: bool
    ) -> None:
        """写缓存并落盘（原子写，gate 频度低、全量重写成本可忽略）。

        参数:
            task_type / s_hash / prompts_version / qid: 缓存键四维。
            correct: 基线侧该题对错。

        关键实现细节:
            先盘后存：新条目先原子落盘（tmp 写 + os.replace）成功后才更新
            内存，磁盘写失败时内存与磁盘一致（均无新条目），无分裂窗口。
        """
        updated = {
            **self._store,
            self._key(task_type, s_hash, prompts_version, qid): correct,
        }
        tmp = self._path.with_suffix(".json.tmp")
        tmp.write_text(json.dumps(updated, ensure_ascii=False), encoding="utf-8")
        os.replace(tmp, self._path)
        self._store = updated