"""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 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


@dataclass(frozen=True)
class InferenceResult:
    """推理聚合结果。

    属性:
        run_id: 运行标识。
        accuracy: 总正确率。
        total: 总题数。
        correct: 正确题数。
        per_task_type: 按题型分组的指标 {task_type: {accuracy, total, correct}}。
        steps_mean: 平均步数。
        token_usage: token 总用量 {prompt_tokens, completion_tokens}。
        stop_reason_counts: 终止原因计数 {reason: count}。
    """

    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(records: list[dict[str, Any]]) -> dict[str, dict[str, Any]]:
    """按 task_type 分组聚合正确率指标。

    参数:
        records: 预测记录列表。

    返回:
        {task_type: {accuracy, total, correct}} 映射。
    """
    task_groups: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for r in records:
        task_groups[r["task_type"]].append(r)

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


def _aggregate_results(records: list[dict[str, Any]], run_id: str) -> InferenceResult:
    """从内存 records 聚合推理指标。

    TRM4 从 DB 回读 predictions 表聚合；TRM5 改为从内存直接聚合，
    避免 DB 回读的同步开销和额外依赖。

    参数:
        records: _run_single_question 返回的 record 列表。
        run_id: 当前运行标识。

    返回:
        InferenceResult 冻结实例。
    """
    total = len(records)
    if total == 0:
        return _zero_result(run_id)

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

    return InferenceResult(
        run_id=run_id,
        accuracy=correct / total,
        total=total,
        correct=correct,
        per_task_type=_group_by_task_type(records),
        steps_mean=sum(r["steps_used"] for r in records) / total,
        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), run_id)
    logger.info(
        "推理完成: accuracy={:.2%} ({}/{})",
        inference_result.accuracy,
        inference_result.correct,
        inference_result.total,
    )
    return inference_result