"""节点内容摘要模块 — 两轮 LLM 调用生成 question-conditioned 摘要。

提取轮：带防幻觉 system prompt，提取与问题相关的信息。
验证轮：带核实 system prompt，逐条核实并给置信度。
与 TRM4 core/tree/summarizer.py 保真迁移：
同步 → async、_call_llm → await llm.chat()、ThreadPoolExecutor → asyncio.gather。
"""

from __future__ import annotations

import asyncio
import re
from typing import TYPE_CHECKING, Any

from loguru import logger

if TYPE_CHECKING:
    from collections.abc import Callable
    from pathlib import Path

    from core.protocols import LLMProvider

# ── 正则常量 ──────────────────────────────────────────────────────────

# 行号引注组：括号包裹的 s/c 行号列表，如 (s1) / (c2,s5) / (c70-c73,s196-s200)
# （兼容全角括号与逗号；单元允许范围语法 s3-s5 / s3-5，60-span 实测模型常用）
_ANCHOR_GROUP = re.compile(
    r"[（(]\s*([sc]\d+(?:-[sc]?\d+)?(?:\s*[,，]\s*[sc]\d+(?:-[sc]?\d+)?)*)\s*[)）]"
)
_ANCHOR_RANGE = re.compile(r"([sc])(\d+)-([sc]?)(\d+)")
_RELEVANT_SECTION = re.compile(r"\[相关信息\](.*?)(?=\n\[|\Z)", re.DOTALL)
# 无相关信息声明句：60-span 实测全为"该节点未包含与问题直接相关的信息"类变体
_NO_INFO_STATEMENT = re.compile(r"未包含.*相关.*信息")

# 范围展开条数上限：防 (s1-s9999) 这类爆炸展开
_RANGE_MAX_IDS = 50

# 双封顶参数：上轮 A/B 证明无上限引用膨胀至 8.4 条/span 挤占提取预算（hall +51%）
_EXPAND_MAX_ITEMS = 5
_EXPAND_MAX_CHARS = 800
_EXPAND_LINE_CAP = 200


# ── Prompt 加载 ──────────────────────────────────────────────────────


def _load_prompt(prompts_dir: Path, filename: str) -> str:
    """从 prompts 目录加载 system prompt 文件。

    参数:
        prompts_dir: prompt 文件所在目录。
        filename: prompt 文件名。

    返回:
        文件内容字符串。
    """
    return (prompts_dir / filename).read_text(encoding="utf-8")


# ── Anchor 工具函数 ──────────────────────────────────────────────────


def _expand_anchor_ids(group_text: str) -> list[str]:
    """把引注组文本展开为逐 id 列表（支持范围语法）。

    参数:
        group_text: _ANCHOR_GROUP 捕获的组内文本，如 "s3-s5, c1"。

    返回:
        逐 id 列表。合法范围（同前缀、起点<=终点、展开条数<=50）展开为
        逐 id（"s3-s5"/"s3-5" -> s3,s4,s5）；非法范围（跨前缀如 c3-s5、
        起点>终点、展开条数超限防爆炸）保留原 token——后续查表必然失配，
        整段按 1 个非法锚计罚剔除。
    """
    ids: list[str] = []
    for token in re.split(r"[,，]\s*", group_text):
        token = token.strip()
        m = _ANCHOR_RANGE.fullmatch(token)
        if m is None:
            ids.append(token)
            continue
        prefix, start = m.group(1), int(m.group(2))
        end_prefix, end = m.group(3), int(m.group(4))
        legal_range = (
            (not end_prefix or end_prefix == prefix)
            and start <= end
            and end - start + 1 <= _RANGE_MAX_IDS
        )
        if not legal_range:
            ids.append(token)
            continue
        ids.extend(f"{prefix}{i}" for i in range(start, end + 1))
    return ids


def check_anchors(summary: str, anchor_map: dict[str, str]) -> tuple[str, dict[str, int]]:
    """校验行号引注：非法行号删锚不删断言。

    参数:
        summary: 提取轮输出（含行号引注）。
        anchor_map: {锚: 原文行} 查表。

    返回:
        (清理后文本, {"n_assertions", "n_anchored", "n_illegal"})。

    关键实现细节:
        清洗全文、统计限段：非法锚无论出现在哪一段都删除并计入 n_illegal
        （避免未校验段落的编造锚流入装配展开）；断言统计
        （n_assertions/n_anchored）仅数 [相关信息] 段内非空内容行。
        引注组先经 _expand_anchor_ids 把范围语法展开为逐 id 再逐 id 校验
        （合法子集重写为逐 id 列表如 (s3,s4,s5)），组内全非法则整组删除；
        组外文本一律不动（删锚不删断言）。分母口径：匹配"未包含...相关...
        信息"词面的声明句不计入 n_assertions——它们天然无锚，计入会虚压
        遵从率。
    """
    stats: dict[str, int] = {"n_assertions": 0, "n_anchored": 0, "n_illegal": 0}

    def _clean_group(gm: re.Match) -> str:
        ids = _expand_anchor_ids(gm.group(1))
        legal = [i for i in ids if i in anchor_map]
        stats["n_illegal"] += len(ids) - len(legal)
        return f"({','.join(legal)})" if legal else ""

    cleaned = _ANCHOR_GROUP.sub(_clean_group, summary)
    m = _RELEVANT_SECTION.search(cleaned)
    if m is None:
        return cleaned, stats
    for line in m.group(1).splitlines():
        line = line.strip().lstrip("-•*").strip()
        if not line:
            continue
        if _NO_INFO_STATEMENT.search(line):
            continue
        stats["n_assertions"] += 1
        if _ANCHOR_GROUP.search(line):
            stats["n_anchored"] += 1
    return cleaned, stats


def _cited_anchor_ids(summary: str, anchor_map: dict[str, str]) -> list[str]:
    """按引注首次出现顺序收集合法锚 id（去重）。

    参数:
        summary: 含行号引注的文本。
        anchor_map: {锚: 原文行} 查表。

    返回:
        去重后的合法锚 id 列表（保持首次出现顺序）。

    关键实现细节:
        从 assemble_anchored_output 提取以满足圈复杂度门槛；范围语法经
        _expand_anchor_ids 展开后逐 id 收集；只收合法锚（非法锚已由
        check_anchors 清除，此处过滤是防御性双保险）。
    """
    ordered: list[str] = []
    for gm in _ANCHOR_GROUP.finditer(summary):
        for aid in _expand_anchor_ids(gm.group(1)):
            if aid in anchor_map and aid not in ordered:
                ordered.append(aid)
    return ordered


def assemble_anchored_output(
    summary: str, anchor_map: dict[str, str], mode: str
) -> tuple[str, dict[str, int]]:
    """按装配形态生成最终输出：展开引文并施加双封顶。

    参数:
        summary: check_anchors 清理后的文本。
        anchor_map: {锚: 原文行}。
        mode: "ids"（裸行号）| "ids_expand"（行号+展开）| "expand_only"（展开剥行号）。

    返回:
        (最终文本, {"n_expanded", "n_trunc"})。

    关键实现细节:
        展开按引注首次出现顺序取前 5 条；总额帽按 [引文] 条目完整长度
        （含前缀与引号）记账，<=800 字符；单行原文超 200 字符先截断。
        n_expanded/n_trunc 仅计实际输出的条目。expand_only 先对正文剥除
        全部引注 token、再拼接 [引文] 段（judge 探针判定 id token 被计罚
        时的回退形态）——引文行不经过剥离，原文行中的括号文本得以保留。
    """
    assert mode in ("ids", "ids_expand", "expand_only"), f"未知装配形态: {mode}"
    stats: dict[str, int] = {"n_expanded": 0, "n_trunc": 0}
    if mode != "ids":
        ordered = _cited_anchor_ids(summary, anchor_map)
        expansions: list[str] = []
        total = 0
        for aid in ordered[:_EXPAND_MAX_ITEMS]:
            line = anchor_map[aid]
            truncated = len(line) > _EXPAND_LINE_CAP
            if truncated:
                line = line[:_EXPAND_LINE_CAP] + "…"
            entry = f'  ▸ {aid}: "{line}"'
            if total + len(entry) > _EXPAND_MAX_CHARS:
                break
            total += len(entry)
            expansions.append(entry)
            stats["n_expanded"] += 1
            if truncated:
                stats["n_trunc"] += 1
        if mode == "expand_only":
            summary = _ANCHOR_GROUP.sub("", summary)
        if expansions:
            summary = summary + "\n[引文]\n" + "\n".join(expansions)
    return summary, stats


# ── LLM 调用辅助 ─────────────────────────────────────────────────────


async def _call_llm(
    llm: LLMProvider,
    system_prompt: str,
    user_text: str,
    *,
    session_id: str | None = None,
    parent_call_id: str | None = None,
) -> str:
    """调用 LLM 并返回响应文本。

    参数:
        llm: LLMProvider 端口实例。
        system_prompt: 系统提示词。
        user_text: 用户消息文本。
        session_id: 会话 ID（透传遥测）。
        parent_call_id: 父调用 ID（透传遥测）。

    返回:
        模型回答文本。
    """
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_text},
    ]
    response = await llm.chat(messages, session_id=session_id, parent_call_id=parent_call_id)
    return response.content


# ── 摘要函数 ─────────────────────────────────────────────────────────


async def summarize_node(
    llm: LLMProvider,
    raw_text: str,
    question: str,
    prompts_dir: Path,
    *,
    anchor_map: dict[str, str] | None,
    assemble_mode: str,
    stats_sink: Callable[[dict[str, Any]], None] | None = None,
    session_id: str | None = None,
    parent_call_id: str | None = None,
) -> str:
    """对单个节点做 question-conditioned 两轮摘要（可选行号锚模式）。

    参数:
        llm: LLMProvider 端口实例。
        raw_text: 节点文本（锚模式下为带 [c1]/[s1] 行号的素材）。
        question: Agent 当前关注的具体问题。
        prompts_dir: prompt 文件目录。
        anchor_map: {锚: 原文行}；None 表示 v1 行为（无校验无装配无统计）。
        assemble_mode: 装配形态（"ids"/"ids_expand"/"expand_only"），
                       anchor_map 为 None 时忽略。
        stats_sink: 统计回调（None 不收集）；统计严禁写入输出文本。
        session_id: 会话 ID（透传遥测）。
        parent_call_id: 父调用 ID（透传遥测）。

    返回:
        "[内容摘要] {结果}\\n[核实] {验证结果}" 或错误信息。

    关键实现细节:
        锚模式流程：提取 -> check_anchors 清洗 -> 核实轮（见清洗后未装配文本）
        -> assemble_anchored_output 装配 -> sink 上报。sink dict 完整键名：
        n_assertions/n_anchored/n_illegal（check_anchors）、
        n_expanded/n_trunc（装配）、output_chars（最终输出字符数）、
        pre_assembly（清洗后未装配文本快照）、anchor_map（原样透传）。
    """
    extract_input = f"问题: {question}\n\n以下是视频片段的描述和字幕:\n{raw_text}"
    try:
        raw_summary = await _call_llm(
            llm,
            _load_prompt(prompts_dir, "view_node_extract.md"),
            extract_input,
            session_id=session_id,
            parent_call_id=parent_call_id,
        )
    except Exception as e:
        return f"[摘要错误] {e}"

    anchor_stats: dict[str, int] = {}
    if anchor_map is not None:
        raw_summary, anchor_stats = check_anchors(raw_summary, anchor_map)
    pre_assembly = raw_summary

    verify_input = (
        f"问题: {question}\n\n"
        f"原始内容:\n{raw_text}\n\n"
        f"以下是另一个模型基于上述内容生成的摘要，请核实:\n{raw_summary}"
    )
    try:
        verify_result = await _call_llm(
            llm,
            _load_prompt(prompts_dir, "view_node_verify.md"),
            verify_input,
            session_id=session_id,
            parent_call_id=parent_call_id,
        )
    except Exception as e:
        logger.warning("验证轮调用失败，跳过: {}", e)
        verify_result = "跳过（调用失败）"

    if anchor_map is not None:
        raw_summary, asm_stats = assemble_anchored_output(raw_summary, anchor_map, assemble_mode)
        anchor_stats.update(asm_stats)

    result = f"[内容摘要] {raw_summary}\n[核实] {verify_result}"
    if anchor_map is not None and stats_sink is not None:
        stats_sink(
            {
                **anchor_stats,
                "output_chars": len(result),
                "pre_assembly": pre_assembly,
                "anchor_map": anchor_map,
            }
        )
    return result


async def summarize_children(
    llm: LLMProvider,
    children_info: list[dict[str, Any]],
    question: str,
    prompts_dir: Path,
    *,
    session_id: str | None = None,
    parent_call_id: str | None = None,
) -> str:
    """对子节点列表做 question-conditioned 相关性标注（两轮）。

    参数:
        llm: LLMProvider 端口实例。
        children_info: 子节点信息列表，每项含 id, time_range, summary。
        question: Agent 当前关注的具体问题。
        prompts_dir: prompt 文件目录。
        session_id: 会话 ID（透传遥测）。
        parent_call_id: 父调用 ID（透传遥测）。

    返回:
        带相关性标注的子节点概览文本。失败时降级返回原始列表。
    """
    lines = []
    for child in children_info:
        t_start, t_end = child["time_range"]
        lines.append(f"- {child['id']} ({t_start:.0f}-{t_end:.0f}s): {child['summary']}")
    children_text = "\n".join(lines)

    extract_input = f"问题: {question}\n\n{children_text}"
    try:
        raw_ranking = await _call_llm(
            llm,
            _load_prompt(prompts_dir, "view_node_children_extract.md"),
            extract_input,
            session_id=session_id,
            parent_call_id=parent_call_id,
        )
    except Exception as e:
        logger.warning("子节点标注失败，回退原始列表: {}", e)
        return children_text

    verify_input = (
        f"问题: {question}\n\n"
        f"原始子节点列表:\n{children_text}\n\n"
        f"以下是另一个模型基于上述信息生成的相关性标注，请核实:\n{raw_ranking}"
    )
    try:
        verify_result = await _call_llm(
            llm,
            _load_prompt(prompts_dir, "view_node_children_verify.md"),
            verify_input,
            session_id=session_id,
            parent_call_id=parent_call_id,
        )
        return f"{raw_ranking}\n[核实] {verify_result}"
    except Exception as e:
        logger.warning("子节点标注验证轮失败，跳过: {}", e)
        return raw_ranking


async def _summarize_search_result(
    llm: LLMProvider,
    raw_text: str,
    question: str,
    prompts_dir: Path,
    *,
    session_id: str | None = None,
    parent_call_id: str | None = None,
) -> str:
    """对搜索结果做两轮摘要（search_similar 专用）。

    参数:
        llm: LLMProvider 端口实例。
        raw_text: 节点原始文本。
        question: Agent 当前关注的具体问题。
        prompts_dir: prompt 文件目录。
        session_id: 会话 ID（透传遥测）。
        parent_call_id: 父调用 ID（透传遥测）。

    返回:
        "[内容摘要] {提取结果}\\n[核实] {验证结果}" 或错误信息。
    """
    extract_input = f"问题: {question}\n\n以下是语义搜索命中的视频节点描述和字幕:\n{raw_text}"
    try:
        raw_summary = await _call_llm(
            llm,
            _load_prompt(prompts_dir, "search_similar_extract.md"),
            extract_input,
            session_id=session_id,
            parent_call_id=parent_call_id,
        )
    except Exception as e:
        return f"[摘要错误] {e}"

    verify_input = (
        f"问题: {question}\n\n"
        f"原始内容:\n{raw_text}\n\n"
        f"以下是另一个模型基于上述内容生成的摘要，请核实:\n{raw_summary}"
    )
    try:
        verify_result = await _call_llm(
            llm,
            _load_prompt(prompts_dir, "search_similar_verify.md"),
            verify_input,
            session_id=session_id,
            parent_call_id=parent_call_id,
        )
        return f"[内容摘要] {raw_summary}\n[核实] {verify_result}"
    except Exception as e:
        logger.warning("搜索结果验证轮失败，跳过: {}", e)
        return f"[内容摘要] {raw_summary}\n[核实] 跳过（调用失败）"


async def summarize_nodes_batch(
    llm: LLMProvider,
    items: list[tuple[str, str, str]],
    question: str,
    prompts_dir: Path,
    *,
    session_id: str | None = None,
    parent_call_id: str | None = None,
) -> list[tuple[str, str]]:
    """并发对多个搜索结果做两轮摘要。

    参数:
        llm: LLMProvider 端口实例。
        items: [(node_id, raw_text, extra_info), ...] 列表。
        question: Agent 当前关注的具体问题。
        prompts_dir: prompt 文件目录。
        session_id: 会话 ID（透传遥测）。
        parent_call_id: 父调用 ID（透传遥测）。

    返回:
        [(node_id, summary_text), ...] 列表，顺序与输入一致。
    """
    if not items:
        return []

    async def _worker(idx: int, node_id: str, raw_text: str) -> tuple[int, str, str]:
        """单个节点的摘要工作协程。"""
        summary = await _summarize_search_result(
            llm,
            raw_text,
            question,
            prompts_dir,
            session_id=session_id,
            parent_call_id=parent_call_id,
        )
        return idx, node_id, summary

    tasks = [_worker(i, nid, text) for i, (nid, text, _) in enumerate(items)]
    results_raw = await asyncio.gather(*tasks)

    results: dict[int, tuple[str, str]] = {}
    for idx, node_id, summary in results_raw:
        results[idx] = (node_id, summary)

    return [results[i] for i in range(len(items))]
