"""慢更新动量生成 — epoch 末为单个 skill 产出新的动量指导。

对标 SkillOpt 的 slow_update 机制：拿上一 epoch 末与当前 epoch 末两版 skill，
在固定样本上各跑一遍得到纵向对比（comparison_pairs），反思上一轮动量指导是否奏效、
本轮正文改动是改善还是漂移，据此重写动量指导。新指导经 patch 引擎的 replace_momentum
写回 skill 的 momentum 受保护区，作为下一轮进化的方向锚。

从 TRM4 core/harness/momentum.py（156 行）迁移 + async 化。
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from loguru import logger

from core.evolution.diagnose import extract_json_from_response

if TYPE_CHECKING:
    from pathlib import Path

    from core.protocols import LLMProvider


# =========================================================================
# 四类纵向对比类别名（单一真源）
# =========================================================================

IMPROVED = "improved"  # 错→对
REGRESSED = "regressed"  # 对→错
PERSISTENT_FAIL = "persistent_fail"  # 错→错
STABLE_SUCCESS = "stable_success"  # 对→对

# 类别名 → 展示标题，列表顺序即展示顺序。
# 回退（REGRESSED）刻意排在改善（IMPROVED）之前——它是最该警惕的伤害信号。
_CATEGORY_LABELS: tuple[tuple[str, str], ...] = (
    (REGRESSED, "从对变错（回退，最高优先级）"),
    (PERSISTENT_FAIL, "始终答错（持续失败）"),
    (IMPROVED, "从错变对（改善）"),
    (STABLE_SUCCESS, "始终答对（稳定成功）"),
)


# =========================================================================
# 辅助函数
# =========================================================================


def _categorize_pair(pair: dict[str, Any]) -> str:
    """按两版正误派生纵向对比类别。

    用键值的真值（bool(...)）表示该题在两版上各自的正误：缺 correct_prev/
    correct_curr 键时直接抛 KeyError 向上传播——这是上游数据损坏（不是裁判语义
    歧义），静默当 False 会伪造 persistent_fail 证据、污染动量指导，故不掩盖。

    参数:
        pair: 单个纵向对比对，须含 correct_prev/correct_curr 两键。

    返回:
        四个类别命名常量之一：IMPROVED（错→对）/REGRESSED（对→错）/
        PERSISTENT_FAIL（错→错）/STABLE_SUCCESS（对→对）。

    异常:
        KeyError: 缺 correct_prev 或 correct_curr 键时。
    """
    correct_prev = bool(pair["correct_prev"])
    correct_curr = bool(pair["correct_curr"])
    if not correct_prev and correct_curr:
        return IMPROVED
    if correct_prev and not correct_curr:
        return REGRESSED
    if not correct_prev and not correct_curr:
        return PERSISTENT_FAIL
    return STABLE_SUCCESS


def _format_comparison_pairs(comparison_pairs: list[dict[str, Any]]) -> str:
    """将纵向对比对格式化为裁判可读文本，按 _CATEGORY_LABELS 分组与排序。

    参数:
        comparison_pairs: 每个 dict 含 question/prev_prediction/curr_prediction/
            correct_prev/correct_curr 字段，描述一道固定样本上两版的成对结果。

    返回:
        可读的纵向对比文本；空列表返回占位说明。

    异常:
        KeyError: 任一 pair 缺 correct_prev/correct_curr 键时；不掩盖的理由见
            _categorize_pair docstring。
    """
    if not comparison_pairs:
        return "（本轮无可用纵向对比样本）"

    grouped: dict[str, list[dict[str, Any]]] = {key: [] for key, _ in _CATEGORY_LABELS}
    for pair in comparison_pairs:
        grouped[_categorize_pair(pair)].append(pair)

    lines: list[str] = [f"固定样本总数：{len(comparison_pairs)}"]
    for key, label in _CATEGORY_LABELS:
        entries = grouped[key]
        lines.append(f"\n### {label}（{len(entries)} 题）")
        if not entries:
            lines.append("（无）")
            continue
        for pair in entries:
            lines.append(
                f"- 题目：{pair.get('question', '')}\n"
                f"  上版预测：{pair.get('prev_prediction', '')} | "
                f"当前版预测：{pair.get('curr_prediction', '')}"
            )
    return "\n".join(lines)


# =========================================================================
# 入口
# =========================================================================


async def run_slow_momentum(
    llm: LLMProvider,
    diagnose_prompts_dir: Path,
    skill_content: str,
    prev_skill: str,
    prev_guidance: str,
    comparison_pairs: list[dict[str, Any]],
) -> str:
    """为单个 skill 生成新的慢更新动量指导。

    参数:
        llm: LLM 端口（async chat）。
        diagnose_prompts_dir: 诊断 prompt 目录（根 prompts/，slow_momentum.md 在此）。
        skill_content: 当前版 skill 正文。
        prev_skill: 上一版 skill 正文。
        prev_guidance: 上一轮写下的动量指导。
        comparison_pairs: 固定样本上两版 rollout 的成对结果（含 question/
            prev_prediction/curr_prediction/correct_prev/correct_curr）。

    返回:
        新的动量指导文本；解析失败时保留 prev_guidance。

    关键实现细节:
        - _format_comparison_pairs 刻意置于 try 块之外（prompt 构造阶段）：它对每个
          pair 取 correct_prev/correct_curr，缺键抛 KeyError 直接向上传播，不被下方
          针对裁判语义歧义的 except ValueError 吞掉。
        - 解析失败保留上轮指导：extract_json_from_response 抛 ValueError、缺
          slow_update_content 字段、或该字段非 str，均视为语义解析失败，返回
          prev_guidance（判不准时保守保留上轮指导，对标 diagnose 的保护性 fallback）。
        - P5 边界：仅捕 ValueError 这一语义歧义；llm.chat 的基础设施失败
          （网络/API 异常）刻意不捕，向上传播，绝不用默认值掩盖。
    """
    system_prompt = (diagnose_prompts_dir / "slow_momentum.md").read_text(encoding="utf-8")
    # _format_comparison_pairs 刻意置于下方 try 块之外（prompt 构造阶段）：它对每个
    # pair 取 correct_prev/correct_curr，缺键抛 KeyError 直接向上传播，不被下方针对
    # 裁判语义歧义的 except ValueError 吞掉。异常类型选 KeyError（非 ValueError），
    # 即便位置疏忽落入 try 也不会被误吞。
    user_prompt = (
        f"## 上一版 skill 正文\n{prev_skill}\n\n"
        f"## 当前版 skill 正文\n{skill_content}\n\n"
        f"## 上一轮的动量指导\n{prev_guidance}\n\n"
        f"## 固定样本纵向对比（上版 vs 当前版）\n"
        f"{_format_comparison_pairs(comparison_pairs)}"
    )
    response = await llm.chat(
        [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]
    )
    raw = response.content
    try:
        parsed = extract_json_from_response(raw)
        new_guidance = parsed.get("slow_update_content")
        if not isinstance(new_guidance, str):
            raise ValueError("slow_update_content 字段缺失或非字符串")
    except ValueError:
        logger.warning("慢更新动量解析失败，保留上轮动量指导")
        return prev_guidance
    return new_guidance