"""async 推理编排 — 训练循环的 forward()。

从 TRM4 core/harness/inference.py (~560 行) 迁移，重大重构:
- 同步 ThreadPoolExecutor → asyncio.Semaphore + asyncio.gather
- LLMClient.from_env() 每题构造 → llm: LLMProvider 注入共享
- SentenceTransformer/OCR 内部构造 → 调用方通过 tool_dispatch_fn 注入
- run_id 必传，空串 → ValueError
- _aggregate_results 从内存 results 聚合（非 DB 回读）
- record_run 由调用方（Runner）负责
- prompt 构建由调用方注入 prompt_builder
"""

from __future__ import annotations

import asyncio
import json
from collections import defaultdict
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

from loguru import logger

from app.harness.question_units import build_units, unit_correctness
from core.agent.loop import AgentLoop

if TYPE_CHECKING:
    from collections.abc import Callable

    from app.harness.log import HarnessLog
    from core.agent.types import LoopResult
    from core.protocols import LLMProvider
    from core.types import GeneratedQuestion, QuestionUnit


@dataclass(frozen=True)
class InferenceResult:
    """推理聚合结果（正确率按 unit 粒度）。

    属性:
        run_id: 运行标识。
        accuracy: unit 级正确率（correct / total）。
        total: unit 总数（single 数 + pair 数，孤儿 pair 已剔除不计入）。
        correct: 正确 unit 数（single 单题正确；pair 走 original/mirror 双向 AND）。
        per_task_type: 按题型分组的 unit 级指标 {task_type: {accuracy, total, correct}}。
        steps_mean: 平均步数（record 粒度，逐题溯源）。
        token_usage: token 总用量 {prompt_tokens, completion_tokens}（record 粒度）。
        stop_reason_counts: 终止原因计数 {reason: count}（record 粒度）。
    """

    run_id: str
    accuracy: float
    total: int
    correct: int
    per_task_type: dict[str, dict]
    steps_mean: float
    token_usage: dict[str, int]
    stop_reason_counts: dict[str, int]


# ---------------------------------------------------------------------------
# 表 Schema 定义（5 张表，保留 TRM4 全部 schema）
# ---------------------------------------------------------------------------

PREDICTIONS_SCHEMA: dict[str, str] = {
    "video_id": "TEXT",
    "question_id": "TEXT",
    "task_type": "TEXT",
    "prediction": "TEXT",
    "answer": "TEXT",
    "evidence": "TEXT",
    "reasoning": "TEXT",
    "steps_used": "INTEGER",
    "prompt_tokens": "INTEGER",
    "completion_tokens": "INTEGER",
    "stop_reason": "TEXT",
    "steps_json": "JSON",
}

TRACES_SCHEMA: dict[str, str] = {
    "video_id": "TEXT",
    "question_id": "TEXT",
    "step": "INTEGER",
    "tool_name": "TEXT",
    "tool_args": "JSON",
    "tool_output": "TEXT",
    "thought": "TEXT",
}

VALIDATION_FLAGS_SCHEMA: dict[str, str] = {
    "video_id": "TEXT",
    "question_id": "TEXT",
    "has_l3_visit": "INTEGER",
    "l1_count": "INTEGER",
    "l2_count": "INTEGER",
    "l3_count": "INTEGER",
}

ANCHOR_CHECK_SCHEMA: dict[str, str] = {
    "video_id": "TEXT",
    "question_id": "TEXT",
    "step": "INTEGER",
    "n_assertions": "INTEGER",
    "n_anchored": "INTEGER",
    "n_illegal": "INTEGER",
    "n_expanded": "INTEGER",
    "n_trunc": "INTEGER",
    "output_chars": "INTEGER",
}

OF_HEALTH_SCHEMA: dict[str, str] = {
    "video_id": "TEXT",
    "question_id": "TEXT",
    "step": "INTEGER",
    "ocr_injected": "INTEGER",
    "ocr_chars": "INTEGER",
    "ocr_failed": "INTEGER",
    "discrepancy": "INTEGER",
    "abstain": "INTEGER",
}


# ---------------------------------------------------------------------------
# 内部工具
# ---------------------------------------------------------------------------


class _DispatcherAdapter:
    """将裸 async callable 包装为 ToolDispatcher Protocol 实例。

    AgentLoop 要求 ToolDispatcher（有 dispatch 方法），而 run_inference
    接收的 tool_dispatch_fn 是裸 async callable。此适配器桥接两者。

    参数:
        fn: async def (tool_name, args, *, context) -> str。
    """

    def __init__(self, fn: Callable[..., Any]) -> None:
        self._fn = fn

    async def dispatch(
        self, tool_name: str, args: dict[str, Any], *, context: dict[str, Any]
    ) -> str:
        """转发工具调用给被包装的 callable。"""
        return await self._fn(tool_name, args, context=context)


def _to_text_field(value: Any) -> str:
    """把 prediction 的 evidence/reasoning 归一为可入库的文本。

    LLM 有时把这些字段返回成 list 或 dict（而非字符串）。sqlite 无法绑定
    非标量类型，直接入库会抛 ProgrammingError 致该题丢失预测行、进而触发
    rollout 完整性护栏中止整轮。凡非 str 一律 JSON 序列化为文本。

    参数:
        value: evidence/reasoning 原始值（可能是 str/list/dict）。

    返回:
        可直接入库的字符串。
    """
    if isinstance(value, str):
        return value
    return json.dumps(value, ensure_ascii=False)


def _zero_result(run_id: str) -> InferenceResult:
    """空记录时的零值 InferenceResult。

    参数:
        run_id: 运行标识。

    返回:
        全零的 InferenceResult。
    """
    return InferenceResult(
        run_id=run_id,
        accuracy=0.0,
        total=0,
        correct=0,
        per_task_type={},
        steps_mean=0.0,
        token_usage={"prompt_tokens": 0, "completion_tokens": 0},
        stop_reason_counts={},
    )


def _group_by_task_type(graded: list[tuple[QuestionUnit, bool]]) -> dict[str, dict[str, Any]]:
    """按 task_type 分组聚合 unit 级正确率指标。

    pair 单元整体计 1 个 unit，归入其 task_type；single 单元计 1 个 unit。

    参数:
        graded: (单元, 该单元是否整体正确) 元组列表。

    返回:
        {task_type: {accuracy, total, correct}} 映射（unit 粒度）。
    """
    task_groups: dict[str, list[bool]] = defaultdict(list)
    for unit, is_correct in graded:
        task_groups[unit.task_type].append(is_correct)

    per_task_type: dict[str, dict[str, Any]] = {}
    for task_type, verdicts in task_groups.items():
        t_total = len(verdicts)
        t_correct = sum(verdicts)
        per_task_type[task_type] = {
            "accuracy": t_correct / t_total,
            "total": t_total,
            "correct": t_correct,
        }
    return per_task_type


def _is_valid_pair(group: list[GeneratedQuestion]) -> bool:
    """判定同一 pair_id 分组是否为合法孪生对（恰好 1 original + 1 mirror，无多余）。

    要求分组总数恰为 2 且角色齐备唯一；有额外非法 role 记录（total>2）或角色
    缺失/重复均视为非法，交由调用方剔除，防非法记录混入 build_units。

    参数:
        group: 归属同一 pair_id 的题目列表。

    返回:
        合法孪生对为 True，否则 False。
    """
    if len(group) != 2:
        return False
    originals = sum(1 for q in group if q.question_role == "pair_original")
    mirrors = sum(1 for q in group if q.question_role == "pair_mirror")
    return originals == 1 and mirrors == 1


def _drop_orphan_pairs(questions: list[GeneratedQuestion]) -> list[GeneratedQuestion]:
    """剔除收不齐 2 条 / 角色非法的孤儿 pair，告警不静默。

    每条题目均会各答一次并逐题落库；能否合成 pair 单元仅取决于 questions
    是否同时含该 pair_id 的 original + mirror（且无多余非法记录）。非法者告警并
    整对剔除，使后续 build_units 只面对合法孪生对（不触发 fail-fast），孤儿 unit
    不计入 total（对齐设计 §8 聚合入口的"告警 + 剔除"）。

    参数:
        questions: 待聚合的题目列表（可混含 single 与孪生对成员）。

    返回:
        可安全交给 build_units 的题目列表（single 全保留，pair 仅保留合法成对者）。
    """
    by_pair: dict[str, list[GeneratedQuestion]] = defaultdict(list)
    singles: list[GeneratedQuestion] = []
    for q in questions:
        if q.pair_id:
            by_pair[q.pair_id].append(q)
        else:
            singles.append(q)

    kept_pairs: list[GeneratedQuestion] = []
    for pair_id, group in by_pair.items():
        if _is_valid_pair(group):
            kept_pairs.extend(group)
        else:
            logger.warning(
                "孤儿 pair {}：非法配对（total={}），剔除该 unit 不计入 total",
                pair_id,
                len(group),
            )
    return singles + kept_pairs


def _per_question_correctness(records: list[dict[str, Any]]) -> dict[str, bool]:
    """由逐题 record 构造 question_id → 该题作答是否正确 的映射。

    prediction 为 None（作答异常）时与 answer 不相等 → False，天然计错。

    参数:
        records: _run_single_question 返回的逐题 record 列表。

    返回:
        {question_id: prediction == answer} 映射，供 unit_correctness 取值。
    """
    return {r["question_id"]: r["prediction"] == r["answer"] for r in records}


def _grade_unit(unit: QuestionUnit, per_q: dict[str, bool]) -> bool:
    """判定单元整体正确性，缺 prediction 时 fail-loud（带上下文）。

    _drop_orphan_pairs 已剔除孤儿/非法配对，正常情况下 unit 内每题都应有对应
    record；若仍缺失说明聚合不变量被破坏（如 records 与 questions 不同源）。此处
    显式抛带上下文的 ValueError（fail-loud，不 catch/不跳过/不兜底），而非放任
    unit_correctness 抛裸 KeyError 丢失定位信息。

    参数:
        unit: 待判定单元。
        per_q: question_id → 该题是否作答正确 的映射。

    返回:
        单元整体是否正确（single 即单题正确；pair 走双向 AND）。

    异常:
        ValueError: unit 内某 question_id 不在 per_q 中（聚合不变量被破坏）。
    """
    missing = [q.question_id for q in unit.questions if q.question_id not in per_q]
    if missing:
        raise ValueError(
            f"unit {unit.unit_id} 的 question {missing} 缺 prediction"
            "（_drop_orphan_pairs 后不应发生，聚合不变量被破坏）"
        )
    return unit_correctness(unit, per_q)


def _aggregate_results(
    records: list[dict[str, Any]],
    questions: list[GeneratedQuestion],
    run_id: str,
) -> InferenceResult:
    """从内存 records + 题目列表按 unit 粒度聚合推理指标。

    逐题 record 保留逐题溯源（token/steps/stop_reason 诊断仍按 record 汇总）；
    正确率则按 unit 粒度计：single 计 1，AR pair 经 build_units 收齐 original +
    mirror 后走 unit_correctness 的双向 AND 判定，整对计 1 个 unit。孤儿 pair
    在 _drop_orphan_pairs 中告警 + 剔除，不计入 total。

    参数:
        records: _run_single_question 返回的逐题 record 列表。
        questions: 与 records 对应的题目列表（提供 pair_id/question_role 元数据）。
        run_id: 当前运行标识。

    返回:
        InferenceResult 冻结实例（total/correct/per_task_type 为 unit 粒度）。
    """
    if not records:
        return _zero_result(run_id)

    per_q = _per_question_correctness(records)
    units = build_units(_drop_orphan_pairs(questions))
    graded = [(unit, _grade_unit(unit, per_q)) for unit in units]

    total = len(graded)
    correct = sum(1 for _, is_correct in graded if is_correct)

    stop_counts: dict[str, int] = defaultdict(int)
    for r in records:
        stop_counts[r["stop_reason"]] += 1

    n_records = len(records)
    return InferenceResult(
        run_id=run_id,
        accuracy=correct / total if total else 0.0,
        total=total,
        correct=correct,
        per_task_type=_group_by_task_type(graded),
        steps_mean=sum(r["steps_used"] for r in records) / n_records,
        token_usage={
            "prompt_tokens": sum(r["prompt_tokens"] for r in records),
            "completion_tokens": sum(r["completion_tokens"] for r in records),
        },
        stop_reason_counts=dict(stop_counts),
    )


# ---------------------------------------------------------------------------
# 单题推理
# ---------------------------------------------------------------------------


async def _run_single_question(
    qa: GeneratedQuestion,
    *,
    llm: LLMProvider,
    tool_dispatch_fn: Callable[..., Any],
    prompt_builder: Callable[[GeneratedQuestion], tuple[str, str]],
    log: HarnessLog,
    max_steps: int,
    plugins: list[object],
) -> dict[str, Any]:
    """执行单道题目的 Agent 推理。

    悲观默认值：record 初始 stop_reason="error"，成功后覆盖。
    prediction 必落库：log.insert 在 try/except 之后（无论成败）。

    参数:
        qa: 待推理的题目。
        llm: LLMProvider 共享实例。
        tool_dispatch_fn: async 工具调度函数 (tool_name, args, *, context) -> str。
        prompt_builder: (GeneratedQuestion) -> (system_prompt, user_prompt)。
        log: HarnessLog 实例（线程安全）。
        max_steps: AgentLoop 最大步数。
        plugins: pluggy 插件列表。

    返回:
        预测结果字典（含 video_id, question_id, prediction, answer 等）。
    """
    record: dict[str, Any] = {
        "video_id": qa.video_id,
        "question_id": qa.question_id,
        "task_type": qa.task_type,
        "prediction": None,
        "answer": qa.answer,
        "evidence": "",
        "reasoning": "",
        "steps_used": 0,
        "prompt_tokens": 0,
        "completion_tokens": 0,
        "stop_reason": "error",  # 悲观默认
        "steps_json": "[]",
    }

    try:
        system_prompt, user_prompt = prompt_builder(qa)
        dispatcher = _DispatcherAdapter(tool_dispatch_fn)
        loop = AgentLoop(llm, max_steps=max_steps)
        loop_result: LoopResult = await loop.run(
            system_prompt,
            user_prompt,
            dispatcher,
            plugins=plugins,
            session_id=qa.question_id,
        )

        result_dict = loop_result.result if isinstance(loop_result.result, dict) else {}
        evidence = _to_text_field(result_dict.get("evidence", ""))
        reasoning = _to_text_field(result_dict.get("reasoning", ""))
        record.update(
            {
                "prediction": result_dict.get("answer"),
                "evidence": evidence,
                "reasoning": reasoning,
                "steps_used": loop_result.steps_used,
                "prompt_tokens": loop_result.token_usage["prompt_tokens"],
                "completion_tokens": loop_result.token_usage["completion_tokens"],
                "stop_reason": loop_result.stop_reason,
                "steps_json": json.dumps(
                    [
                        {
                            "thought": s.thought,
                            "tool_call": s.tool_call,
                            "tool_output": s.tool_output,
                        }
                        for s in loop_result.steps
                    ],
                    ensure_ascii=False,
                ),
            }
        )
    except Exception:
        logger.exception("[{}] QA {} 执行异常", qa.video_id, qa.question_id)

    # prediction 必落库（try 外，无论成败）
    await asyncio.to_thread(log.insert, "predictions", record)
    return record


# ---------------------------------------------------------------------------
# 建表
# ---------------------------------------------------------------------------


def _ensure_tables(log: HarnessLog) -> None:
    """创建推理所需的 5 张表。

    参数:
        log: HarnessLog 实例。
    """
    log.create_table("predictions", PREDICTIONS_SCHEMA)
    log.create_table("traces", TRACES_SCHEMA)
    log.create_table("validation_flags", VALIDATION_FLAGS_SCHEMA)
    log.create_table("anchor_check", ANCHOR_CHECK_SCHEMA)
    log.create_table("observe_frame_health", OF_HEALTH_SCHEMA)


# ---------------------------------------------------------------------------
# 公共入口
# ---------------------------------------------------------------------------


async def run_inference(
    questions: list[GeneratedQuestion],
    *,
    llm: LLMProvider,
    tool_dispatch_fn: Callable[..., Any],
    prompt_builder: Callable[[GeneratedQuestion], tuple[str, str]],
    log: HarnessLog,
    run_id: str,
    concurrency: int,
    max_steps: int,
    skill_mode: str,
    plugins_factory: Callable[[str, str], list[object]] | None = None,
) -> InferenceResult:
    """在视频树上执行 Agent 推理，对应训练循环的 forward()。

    参数:
        questions: 待推理的题目列表。
        llm: LLMProvider 共享实例（依赖注入）。
        tool_dispatch_fn: async 工具调度函数 (tool_name, args, *, context) -> str。
        prompt_builder: prompt 构建函数 (GeneratedQuestion) -> (system_prompt, user_prompt)。
        log: HarnessLog 实例（由调用方管理生命周期）。
        run_id: 运行标识（必传，空串 → ValueError）。
        concurrency: 最大并发数（asyncio.Semaphore 控制）。
        max_steps: AgentLoop 单题最大步数。
        skill_mode: "auto" / "manual" / "none"（传递给调用方的 prompt/plugin 构建逻辑）。
        plugins_factory: 可选的插件工厂 (video_id, question_id) -> plugins 列表。

    返回:
        InferenceResult（含 accuracy、per_task_type 等聚合指标）。

    异常:
        ValueError: run_id 为空串或纯空白。
    """
    if not run_id or not run_id.strip():
        raise ValueError("run_id 不得为空串或纯空白")

    _ensure_tables(log)

    if not questions:
        logger.info("题目列表为空，返回零值 InferenceResult")
        return _aggregate_results([], [], run_id)

    sem = asyncio.Semaphore(concurrency)
    total_count = len(questions)

    async def _bounded(index: int, qa: GeneratedQuestion) -> dict[str, Any]:
        """信号量限流的单题推理包装。"""
        async with sem:
            plugins = (
                plugins_factory(qa.video_id, qa.question_id) if plugins_factory is not None else []
            )
            result = await _run_single_question(
                qa,
                llm=llm,
                tool_dispatch_fn=tool_dispatch_fn,
                prompt_builder=prompt_builder,
                log=log,
                max_steps=max_steps,
                plugins=plugins,
            )
            logger.info(
                "[{}/{}] {} QA {} 完成 (stop={})",
                index + 1,
                total_count,
                qa.video_id,
                qa.question_id,
                result["stop_reason"],
            )
            return result

    results = await asyncio.gather(*[_bounded(i, qa) for i, qa in enumerate(questions)])

    inference_result = _aggregate_results(list(results), questions, run_id)
    logger.info(
        "推理完成: accuracy={:.2%} ({}/{})",
        inference_result.accuracy,
        inference_result.correct,
        inference_result.total,
    )
    return inference_result
