"""async 块序贯验证编排 — CE-Gate 局部验证的唯一独立子编排器。

从 TRM4 core/harness/validate.py (626 行) 迁移，重大重构：
- 同步 → async（run_inference 注入为 async callable）
- _classify_quadrants → core.evolution.classify_quadrants 纯函数
- 配对逻辑 → 复用 core.evolution.pair_block + 本地证据行组装
- _load_run_rows / _candidate_correctness_from_db → 共享 log.query()
- materialize_candidate_skill 保持同步（纯文件操作）

基线与候选在同一阶梯前缀上逐块配对，只数翻转（基线错→候选对 = W，
基线对→候选错 = L），每块结束调 gate_decision 做四出口判定。
基线侧逐题对错走 BaselineCache 内容寻址缓存，miss 才新鲜跑。
判定逻辑全部在 core/evolution/gate，本模块只负责推理编排与证据收集。
"""

from __future__ import annotations

import json
import shutil
import tempfile
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable

from loguru import logger

from app.harness.gate_ladder import BaselineCache, skill_hash
from core.evolution import (
    GateParams,
    GateVerdict,
    RejectedEdit,
    classify_quadrants,
    gate_decision,
    pair_block,
)

if TYPE_CHECKING:
    from app.harness.inference import InferenceResult
    from app.harness.log import HarnessLog
    from core.types import GeneratedQuestion


# gate_decision 的 decision → ValidationOutcome.stop_reason 映射
_STOP_REASON_BY_DECISION: dict[str, str] = {
    "accept_confirmed": "confirmed",
    "reject_directional": "directional",
    "reject_futility": "futility",
    "accept_provisional": "provisional",
    "reject_inertia": "inertia",
}


# ---------------------------------------------------------------------------
# 注入协议
# ---------------------------------------------------------------------------


@runtime_checkable
class RunInferenceFn(Protocol):
    """注入的推理函数协议。

    调用方（runner）负责绑定 llm、tool_dispatch_fn、prompt_builder、
    log、concurrency、max_steps、skill_mode 等共享依赖。
    validate 侧只传 questions、run_id、skills_dir 三个逐块变化的参数。
    """

    async def __call__(
        self,
        questions: list[GeneratedQuestion],
        *,
        run_id: str,
        skills_dir: Path,
    ) -> InferenceResult: ...


# ---------------------------------------------------------------------------
# 数据类型
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class InferenceRunConfig:
    """一次推理运行的配置三元组，把"如何跑推理"内聚成一组。

    字段:
        concurrency: 推理并发度。
        max_steps: 单题最大推理步数。
        skill_mode: 推理 skill 模式（"auto" / "manual" / "none"）。
    """

    concurrency: int
    max_steps: int
    skill_mode: str


@dataclass
class ValidationOutcome:
    """CE-Gate 局部验证结果：三态动作 + e-process 证据 + 已观测题逐题对错。

    correctness 二轨语义：candidate_correctness 只含已观测题（早停后是
    阶梯前缀子集）；accept 时由 runner 按题粒度增量合并进 state.correctness。
    """

    action: str  # accept_confirmed | accept_provisional | reject
    accepted: bool
    stop_reason: str  # confirmed | directional | futility | provisional | inertia
    e_value: float
    w: int
    l: int  # noqa: E741
    n_used: int
    delta_hat: float
    delta_shrunk: float
    baseline_acc: float  # 已观测题上的基线准确率（观测口径）
    candidate_acc: float  # 已观测题上的候选准确率（观测口径）
    improvements: list[str] = field(default_factory=list)
    regressions: list[str] = field(default_factory=list)
    persistent_fails: list[str] = field(default_factory=list)
    stable_successes: list[str] = field(default_factory=list)
    candidate_correctness: dict[str, bool] = field(default_factory=dict)
    evidence_rows: list[dict] = field(default_factory=list)  # gate_evidence 逐题行，runner 落库


@dataclass
class Probation:
    """一个题型的在途试用账本（每题型至多一个）。

    字段:
        task_type: 题型。
        anchor_skills_version: 锚版本名（最近一个 CONFIRMED 的 skills 版本）——
            回滚时恢复该版本中本题型 skill 文件的内容。
        target_file: 该题型解析后的 skill 文件名。
        correctness_snapshot: 开账时该题型 val 题的对错快照（回滚时恢复）。
        opened_step: 开账时的 global_step（观测用）。
        pending_edits: 试用链上全部候选 edit 的黑名单素材（回滚时整链入黑名单）。
    """

    task_type: str
    anchor_skills_version: str
    target_file: str
    correctness_snapshot: dict[str, bool]
    opened_step: int
    pending_edits: list[RejectedEdit] = field(default_factory=list)


# ---------------------------------------------------------------------------
# 同步辅助函数
# ---------------------------------------------------------------------------


def materialize_candidate_skill(
    workspace_dir: Path,
    base_skills_version: str,
    target_file: str,
    content: str,
) -> Path:
    """将候选 skill 正文物化为 workspace 专用临时目录下唯一命名的候选 skills 目录。

    复制基线 skills 目录到 .cand_tmp/ 下的唯一命名临时目录，然后覆写 target_file。
    构建失败时尽力清理已建临时目录再重抛原始异常。

    参数:
        workspace_dir: Workspace 根目录。基线 skills 从 workspace_dir/skills/<base>
            复制，临时候选落 workspace_dir/.cand_tmp/。
        base_skills_version: 基线 skills 版本名。
        target_file: 被替换的 skill 文件名。
        content: 候选 skill 文件全文。

    返回:
        新建的临时候选目录绝对路径。

    契约:
        构建失败（OSError）时尽力清理已建临时目录再重抛原始异常；
        清理本身失败记 warning。
    """
    cand_tmp_root = workspace_dir / ".cand_tmp"
    cand_tmp_root.mkdir(parents=True, exist_ok=True)
    cand_dir = Path(tempfile.mkdtemp(prefix=f"{base_skills_version}_cand_", dir=cand_tmp_root))
    try:
        base_dir = workspace_dir / "skills" / base_skills_version
        shutil.copytree(base_dir, cand_dir, dirs_exist_ok=True)
        (cand_dir / target_file).write_text(content, encoding="utf-8")
    except OSError:
        try:
            shutil.rmtree(cand_dir)
        except OSError as cleanup_err:
            logger.warning("候选物化失败后清理临时目录也失败 {}: {}", cand_dir, cleanup_err)
        raise
    return cand_dir


def _load_run_rows(
    log: HarnessLog,
    run_id: str,
) -> dict[str, dict[str, Any]]:
    """读取单个 run 的逐题预测行并规范化轨迹字段。

    从 predictions 表读取指定 run 的题目级记录，补充 _correct
    与规范化后的 steps 字段。保持同步（log.query）——仅在推理完成后调用。

    参数:
        log: HarnessLog 共享实例（用 query 方法做只读 SELECT）。
        run_id: 待读取的预测 run_id。

    返回:
        以 question_id 为键的行字典。每行至少包含 prediction、answer、
        _correct、steps 等字段。
    """
    rows = log.query(
        "SELECT question_id, prediction, answer, steps_json FROM predictions WHERE run_id=?",
        (run_id,),
    )
    normalized: dict[str, dict[str, Any]] = {}
    for row in rows:
        raw_steps = row.get("steps_json")
        parsed_steps: Any = raw_steps
        if isinstance(raw_steps, str):
            try:
                parsed_steps = json.loads(raw_steps)
            except json.JSONDecodeError:
                parsed_steps = []
        steps = parsed_steps if isinstance(parsed_steps, list) else []
        normalized[row["question_id"]] = {
            **row,
            "_correct": row.get("prediction") == row.get("answer"),
            "steps": steps,
        }
    return normalized


def _candidate_correctness_from_db(
    log: HarnessLog,
    run_id: str,
    chunk: list[GeneratedQuestion],
) -> dict[str, bool]:
    """从 db 读取候选/基线 run 在指定题目上的逐题对错。

    参数:
        log: HarnessLog 共享实例。
        run_id: 推理 run_id。
        chunk: 题目列表。

    返回:
        question_id -> 是否答对的映射。缺行的题目记为 False。
    """
    rows = _load_run_rows(log, run_id)
    return {q.question_id: rows.get(q.question_id, {}).get("_correct", False) for q in chunk}


# ---------------------------------------------------------------------------
# 块级 async 函数
# ---------------------------------------------------------------------------


async def _resolve_baseline_block(
    chunk: list[GeneratedQuestion],
    task_type: str,
    s_hash: str,
    prompts_version: str,
    baseline_cache: BaselineCache,
    base_skills_dir: Path,
    run_inference: RunInferenceFn,
    log: HarnessLog,
    run_id: str,
) -> tuple[dict[str, bool], int, int]:
    """基线侧处理一个块：缓存优先，miss 的题新鲜跑基线版本并回写缓存。

    参数:
        chunk: 当前块的题目列表。
        task_type: 当前验证题型（缓存键成分）。
        s_hash: 基线侧生效 skill 的内容哈希（缓存键成分）。
        prompts_version: 当前 prompts 版本（缓存键成分）。
        baseline_cache: 基线侧逐题对错缓存。
        base_skills_dir: 基线 skills 版本目录。
        run_inference: 注入的 async 推理函数。
        log: HarnessLog 共享实例（推理后读预测）。
        run_id: 本块基线 run_id。

    返回:
        (b_map, errors_inc, denom_inc)：块内 question_id -> 基线对错、
        本块新增的 INFRA error 计数与推理题次分母增量（全命中时为 0, 0）。
    """
    misses = [
        q
        for q in chunk
        if baseline_cache.get(task_type, s_hash, prompts_version, q.question_id) is None
    ]
    errors_inc = 0
    denom_inc = 0
    if misses:
        r_b = await run_inference(misses, run_id=run_id, skills_dir=base_skills_dir)
        errors_inc = r_b.stop_reason_counts.get("error", 0)
        denom_inc = r_b.total
        fresh = _candidate_correctness_from_db(log, r_b.run_id, misses)
        for qid, correct in fresh.items():
            baseline_cache.put(task_type, s_hash, prompts_version, qid, correct)

    b_map: dict[str, bool] = {}
    for q in chunk:
        val = baseline_cache.get(task_type, s_hash, prompts_version, q.question_id)
        assert val is not None, f"基线缓存补齐后仍有 miss: {q.question_id} run_id={run_id}"
        b_map[q.question_id] = val
    return b_map, errors_inc, denom_inc


async def _run_candidate_block(
    chunk: list[GeneratedQuestion],
    cand_dir: Path,
    run_inference: RunInferenceFn,
    log: HarnessLog,
    run_id: str,
) -> tuple[dict[str, bool], int, int]:
    """候选侧处理一个块：全块新鲜跑候选版本并从 db 读逐题对错。

    参数:
        chunk: 当前块的题目列表。
        cand_dir: 已物化的候选 skills 目录。
        run_inference: 注入的 async 推理函数。
        log: HarnessLog 共享实例（推理后读预测）。
        run_id: 本块候选 run_id。

    返回:
        (c_map, errors_inc, denom_inc)。
    """
    r_c = await run_inference(chunk, run_id=run_id, skills_dir=cand_dir)
    c_map = _candidate_correctness_from_db(log, r_c.run_id, chunk)
    return c_map, r_c.stop_reason_counts.get("error", 0), r_c.total


def _build_evidence_rows(
    chunk: list[GeneratedQuestion],
    b_map: dict[str, bool],
    c_map: dict[str, bool],
    task_type: str,
    block_idx: int,
) -> list[dict]:
    """组装一个块的 gate_evidence 逐题证据行。

    e_value 留 None 待块判定后回填，stop_reason 留空串待终态回填。

    参数:
        chunk: 当前块的题目列表。
        b_map: 块内 question_id -> 基线对错。
        c_map: 块内 question_id -> 候选对错。
        task_type: 当前验证题型。
        block_idx: 当前块序号。

    返回:
        逐题证据行列表。
    """
    return [
        {
            "question_id": q.question_id,
            "task_type": task_type,
            "block_idx": block_idx,
            "baseline_correct": b_map[q.question_id],
            "candidate_correct": c_map[q.question_id],
            "e_value": None,
            "stop_reason": "",
        }
        for q in chunk
    ]


# ---------------------------------------------------------------------------
# INFRA 护栏
# ---------------------------------------------------------------------------


def _check_infra_guard(errors: int, infra_denom: int, gate_guard_err: float) -> None:
    """跨块累计 INFRA 错误率护栏：分母 >=10 且超阈值时 raise。

    参数:
        errors: 两侧累计 error 计数。
        infra_denom: 两侧累计推理题次分母。
        gate_guard_err: 错误率阈值。

    异常:
        RuntimeError: 错误率超阈值。
    """
    if infra_denom >= 10 and errors / infra_denom > gate_guard_err:
        raise RuntimeError(f"gate 推理累计错误率过高 {errors / infra_denom:.0%}，中止本轮")


# ---------------------------------------------------------------------------
# 终态组装
# ---------------------------------------------------------------------------


def _finalize_outcome(
    verdict: GateVerdict,
    w: int,
    l: int,  # noqa: E741
    n_used: int,
    n_plan: int,
    base_obs: dict[str, bool],
    cand_obs: dict[str, bool],
    evidence_rows: list[dict],
    task_type: str,
) -> ValidationOutcome:
    """将块循环终态判定组装为 ValidationOutcome。

    参数:
        verdict: 最后一块的 gate 判定结果。
        w: 累计 W（基线错→候选对翻转）。
        l: 累计 L（基线对→候选错翻转）。
        n_used: 已消费的阶梯题数。
        n_plan: 阶梯总题数。
        base_obs: 累计基线已观测对错。
        cand_obs: 累计候选已观测对错。
        evidence_rows: 逐题证据行。
        task_type: 验证题型（日志用）。

    返回:
        ValidationOutcome。
    """
    action = {
        "accept_confirmed": "accept_confirmed",
        "accept_provisional": "accept_provisional",
    }.get(verdict.decision, "reject")
    stop_reason = _STOP_REASON_BY_DECISION[verdict.decision]
    # 只有终态题的证据行才携带 stop_reason
    evidence_rows[-1]["stop_reason"] = stop_reason

    quadrants = classify_quadrants({qid: (base_obs[qid], cand_obs[qid]) for qid in base_obs})
    baseline_acc = sum(base_obs.values()) / len(base_obs)
    candidate_acc = sum(cand_obs.values()) / len(cand_obs)
    accepted = action != "reject"

    logger.info(
        "gate 局部验证[{}]: 基线{:.1%} → 候选{:.1%} (W={} L={} E={:.2f} n={}/{}) {}",
        task_type,
        baseline_acc,
        candidate_acc,
        w,
        l,
        verdict.e_value,
        n_used,
        n_plan,
        "接受" if accepted else "回滚",
    )

    return ValidationOutcome(
        action=action,
        accepted=accepted,
        stop_reason=stop_reason,
        e_value=verdict.e_value,
        w=w,
        l=l,
        n_used=n_used,
        delta_hat=verdict.delta_hat,
        delta_shrunk=verdict.delta_shrunk,
        baseline_acc=baseline_acc,
        candidate_acc=candidate_acc,
        improvements=quadrants.improvements,
        regressions=quadrants.regressions,
        persistent_fails=quadrants.persistent_fails,
        stable_successes=quadrants.stable_successes,
        candidate_correctness=cand_obs,
        evidence_rows=evidence_rows,
    )


# ---------------------------------------------------------------------------
# 主编排
# ---------------------------------------------------------------------------


async def _run_local_validation(
    workspace_dir: Path,
    cand_dir: Path,
    base_skills_version: str,
    task_type: str,
    base_skill_content: str,
    plan: list[GeneratedQuestion],
    gate_params: GateParams,
    gate_block: int,
    gate_guard_err: float,
    baseline_cache: BaselineCache,
    prompts_version: str,
    run_inference: RunInferenceFn,
    log: HarnessLog,
    gate_run_prefix: str,
) -> ValidationOutcome:
    """块序贯循环主体：逐块基线（缓存优先）/候选配对推理，块间 e-process 判定。

    按 gate_block 切阶梯前缀，每块先补齐基线侧缓存 miss（新鲜跑基线版本
    并逐题写 BaselineCache），再全块跑候选，配对累计 W/L 后调 gate_decision；
    非 continue 即早停。题尽时最后一块的判定即终态（n_remaining=0 走
    provisional/inertia 分支），无循环外补判。

    参数:
        workspace_dir: Workspace 根目录。
        cand_dir: 已物化的候选 skills 目录。
        base_skills_version: 基线 skills 版本名。
        task_type: 当前验证题型。
        base_skill_content: 基线侧生效 skill 全文（skill_hash 作缓存键成分）。
        plan: 已截断到 gate_n_max 的阶梯出题序。
        gate_params: e-process 判据阈值组。
        gate_block: 块大小。
        gate_guard_err: 跨块累计 INFRA 错误率护栏（分母 >=10 才触发）。
        baseline_cache: 基线侧逐题对错缓存。
        prompts_version: 当前 prompts 版本（缓存键成分）。
        run_inference: 注入的 async 推理函数。
        log: HarnessLog 共享实例。
        gate_run_prefix: 块 run_id 前缀（含 "_gate_" 标记）。

    返回:
        ValidationOutcome。

    关键实现:
        INFRA 护栏跨块累计基线+候选两侧的 error 计数，分母（总推理题次）>=10
        且错误率超 gate_guard_err 时直接 raise，避免坏批次污染判定。
    """
    w = 0
    l = 0  # noqa: E741
    n_used = 0
    errors = 0
    infra_denom = 0
    evidence_rows: list[dict] = []
    base_obs: dict[str, bool] = {}
    cand_obs: dict[str, bool] = {}
    s_hash = skill_hash(base_skill_content)
    base_skills_dir = workspace_dir / "skills" / base_skills_version
    chunks = [plan[i : i + gate_block] for i in range(0, len(plan), gate_block)]
    verdict: GateVerdict | None = None

    for block_idx, chunk in enumerate(chunks):
        # Phase 1: 基线侧（缓存优先，miss 新鲜跑）+ 候选侧（全块新鲜跑）
        b_map, err_b, den_b = await _resolve_baseline_block(
            chunk=chunk,
            task_type=task_type,
            s_hash=s_hash,
            prompts_version=prompts_version,
            baseline_cache=baseline_cache,
            base_skills_dir=base_skills_dir,
            run_inference=run_inference,
            log=log,
            run_id=f"{gate_run_prefix}_b{block_idx}_base",
        )
        c_map, err_c, den_c = await _run_candidate_block(
            chunk=chunk,
            cand_dir=cand_dir,
            run_inference=run_inference,
            log=log,
            run_id=f"{gate_run_prefix}_b{block_idx}_cand",
        )

        # Phase 2: INFRA 护栏（跨块累计，分母 >=10 才触发）
        errors += err_b + err_c
        infra_denom += den_b + den_c
        _check_infra_guard(errors, infra_denom, gate_guard_err)

        # Phase 3: 配对 + 证据行 + 块间判定
        qids = [q.question_id for q in chunk]
        pair_result = pair_block(b_map, c_map, qids)
        for qid, (b, c) in pair_result.observed.items():
            base_obs[qid] = b
            cand_obs[qid] = c

        block_rows = _build_evidence_rows(chunk, b_map, c_map, task_type, block_idx)

        w += pair_result.w
        l += pair_result.l  # noqa: E741
        n_used += len(chunk)
        verdict = gate_decision(w, l, n_used, len(plan) - n_used, params=gate_params)

        for row in block_rows:
            row["e_value"] = verdict.e_value
        evidence_rows.extend(block_rows)

        if verdict.decision != "continue":
            break

    # 最后一块判定即终态（n_remaining=0 → provisional/inertia）
    assert verdict is not None, "空阶梯应已在 validate_skill_local 入口拒绝"
    return _finalize_outcome(
        verdict=verdict,
        w=w,
        l=l,
        n_used=n_used,
        n_plan=len(plan),
        base_obs=base_obs,
        cand_obs=cand_obs,
        evidence_rows=evidence_rows,
        task_type=task_type,
    )


async def validate_skill_local(
    workspace_dir: Path,
    base_skills_version: str,
    task_type: str,
    target_file: str,
    candidate_content: str,
    base_skill_content: str,
    ladder_items: list[GeneratedQuestion],
    gate_params: GateParams,
    gate_block: int,
    gate_n_max: int,
    gate_guard_err: float,
    baseline_cache: BaselineCache,
    prompts_version: str,
    run_inference: RunInferenceFn,
    log: HarnessLog,
    gate_run_prefix: str,
) -> ValidationOutcome:
    """块序贯配对验证：阶梯出题，基线/候选逐块配对，e-process 四出口早停。

    参数:
        workspace_dir: workspace 根目录。
        base_skills_version: 基线 skills 版本名（候选物化复制源）。
        task_type: 待验证题型。
        target_file: fallback 解析后该题型的真实生效 skill 文件名
            （record.target_file，可能是共享 default-strategy.md）；
            候选物化写此文件，与 accept 路径同源。
        candidate_content: 候选 skill 全文。
        base_skill_content: 基线侧该题型解析后生效 skill 文件全文
            （skill_hash(base_skill_content) 作 BaselineCache 键成分）。
        ladder_items: 阶梯序题目列表（已排除本 step 案例包题）。
        gate_params: e-process 判据阈值组。
        gate_block: 块大小。
        gate_n_max: 单 gate 题数上限。
        gate_guard_err: 跨块累计 INFRA 错误率护栏（分母 >=10 才触发）。
        baseline_cache: 基线侧逐题对错缓存。
        prompts_version: 当前 prompts 版本（缓存键成分）。
        run_inference: 注入的 async 推理函数（RunInferenceFn 协议）。
        log: HarnessLog 共享实例（供 DB 回读逐题对错）。
        gate_run_prefix: gate 内推理 run_id 前缀，必须含 "_gate_"
            （防泄露过滤靠它识别）。块 run_id = f"{prefix}_b{block_idx}_{arm}"。

    返回:
        ValidationOutcome。逐题证据记入 outcome.evidence_rows 随结果返回，
        gate_evidence 落库由调用方（runner）负责。
    """
    if "_gate_" not in gate_run_prefix:
        raise ValueError(f"gate_run_prefix 必须含 '_gate_'（防泄露过滤依赖）: {gate_run_prefix!r}")
    if not ladder_items:
        raise ValueError(f"task_type={task_type} 阶梯为空，无法验证")

    plan = ladder_items[:gate_n_max]
    cand_dir = materialize_candidate_skill(
        workspace_dir, base_skills_version, target_file, candidate_content
    )
    try:
        return await _run_local_validation(
            workspace_dir=workspace_dir,
            cand_dir=cand_dir,
            base_skills_version=base_skills_version,
            task_type=task_type,
            base_skill_content=base_skill_content,
            plan=plan,
            gate_params=gate_params,
            gate_block=gate_block,
            gate_guard_err=gate_guard_err,
            baseline_cache=baseline_cache,
            prompts_version=prompts_version,
            run_inference=run_inference,
            log=log,
            gate_run_prefix=gate_run_prefix,
        )
    finally:
        try:
            shutil.rmtree(cand_dir)
        except OSError as e:
            logger.warning("候选临时目录清理失败 {}: {}", cand_dir, e)