"""TreeEnvironment：单棵视频树的运行时环境。

提供节点查询、字幕获取、帧路径解析和语义检索能力。
纯数据访问层——不涉及 LLM 调用，LLM 摘要逻辑属于 app/search/。

算法 #12 变更：分块 embedding → 单节点 embedding。
祖先去重 + 锚定验证逻辑保留自 TRM4。
"""

from __future__ import annotations

import re
from pathlib import Path
from typing import TYPE_CHECKING, Any

import numpy as np
from loguru import logger

from app.tree.index import L1Node, L2Node, L3Node, TreeIndex

if TYPE_CHECKING:
    from collections.abc import Callable

# 节点联合类型（内部使用）
AnyNode = L1Node | L2Node | L3Node

# 各层级节点对应的主描述字段名
_LEVEL_LABEL = {
    "L1": "场景层",
    "L2": "事件层",
    "L3": "关键帧层",
}


def _node_level(node: AnyNode) -> str:
    """判断节点层级标签。

    参数:
        node: 树节点实例。

    返回:
        "L1" / "L2" / "L3"。
    """
    if isinstance(node, L1Node):
        return "L1"
    if isinstance(node, L2Node):
        return "L2"
    return "L3"


def _node_description(node: AnyNode) -> str:
    """提取节点的主描述文本。

    参数:
        node: 树节点实例。

    返回:
        描述文本字符串。
    """
    if isinstance(node, L1Node):
        return node.card.scene_summary
    if isinstance(node, L2Node):
        return node.card.event_description
    return node.card.frame_summary


def _collect_card_strings(node: AnyNode) -> list[str]:
    """从节点 card 中递归收集所有非空字符串字段。

    参数:
        node: 树节点实例。

    返回:
        字符串列表（每个非空字段值一项，含内嵌换行的按行拆分）。
    """
    result: list[str] = []
    _collect_from_obj(node.card, result)
    return result


def _collect_from_obj(obj: object, out: list[str]) -> None:
    """递归收集任意嵌套结构中的非空字符串。

    参数:
        obj: dict / list / str / 其他。
        out: 收集结果列表（原地修改）。
    """
    if isinstance(obj, str):
        stripped = obj.strip()
        if stripped:
            out.append(stripped)
    elif isinstance(obj, dict):
        for v in obj.values():
            _collect_from_obj(v, out)
    elif isinstance(obj, (list, tuple)):
        for item in obj:
            _collect_from_obj(item, out)
    elif hasattr(obj, "__dataclass_fields__"):
        # frozen dataclass（Card 类型）
        for field_name in obj.__dataclass_fields__:
            _collect_from_obj(getattr(obj, field_name), out)


class TreeEnvironment:
    """单棵视频树的运行时环境，提供节点查询和语义检索。

    纯数据访问层，不涉及 LLM 调用。

    参数:
        index: 已加载的 TreeIndex 实例。
        frames_dir: 帧文件目录路径（可选；未提供时使用节点自带的 frame_path）。
    """

    def __init__(
        self,
        index: TreeIndex,
        frames_dir: Path | None = None,
    ) -> None:
        self._index = index
        self._frames_dir = frames_dir

        # O(1) 查找表：node_id → 节点实例
        self._id_to_node: dict[str, AnyNode] = {}
        # 父节点映射：node_id → parent_id（根节点为 None）
        self._id_to_parent: dict[str, str | None] = {}

        self._build_lookup_tables()
        logger.debug(
            "TreeEnvironment 初始化完成，节点数={}",
            len(self._id_to_node),
        )

    # ------------------------------------------------------------------
    # 初始化辅助
    # ------------------------------------------------------------------

    def _build_lookup_tables(self) -> None:
        """遍历 TreeIndex 构建 _id_to_node 和 _id_to_parent 映射表。"""
        for l1 in self._index.roots:
            self._id_to_node[l1.id] = l1
            self._id_to_parent[l1.id] = None
            for l2 in l1.children:
                self._id_to_node[l2.id] = l2
                self._id_to_parent[l2.id] = l1.id
                for l3 in l2.children:
                    self._id_to_node[l3.id] = l3
                    self._id_to_parent[l3.id] = l2.id

    # ------------------------------------------------------------------
    # 公开方法
    # ------------------------------------------------------------------

    def view_node(self, node_id: str, *, anchor: bool = False) -> str:
        """返回节点卡片内容 + 子节点概览。

        参数:
            node_id: 节点 ID。
            anchor: 为卡片字段添加行锚标 [c1] [s1] 供引用验证。

        返回:
            格式化文本。

        异常:
            KeyError: 节点不存在。
        """
        node = self._id_to_node.get(node_id)
        if node is None:
            raise KeyError(f"节点不存在: {node_id}")

        level = _node_level(node)
        level_label = _LEVEL_LABEL[level]

        # 时间范围
        time_range_str = self._format_time_range(node)

        # 节点内容
        content = self._node_anchored_text(node) if anchor else self._node_full_text(node)

        parts = [
            f"[节点] {node_id} | {level_label} | {time_range_str}",
            "",
            content,
        ]

        # 子节点概览
        children = self._get_children(node)
        if children:
            parts.append("")
            parts.append(f"[子节点概览] {len(children)} 个子节点")
            for child in children:
                child_desc = _node_description(child)
                child_time = self._format_time_range(child)
                # 截断描述到 120 字符
                if len(child_desc) > 120:
                    child_desc = child_desc[:120] + "..."
                parts.append(f"  - {child.id} | {child_time} | {child_desc}")

        return "\n".join(parts)

    def search_similar(
        self,
        query: str,
        top_k: int = 5,
        *,
        embed_fn: Callable[[str | list[str]], np.ndarray] | None = None,
    ) -> list[tuple[str, float]]:
        """语义搜索 + 祖先去重。

        算法 #12 变更：单节点 embedding（非分块），祖先去重 + 锚定验证保留。

        参数:
            query: 搜索文本。
            top_k: 返回数量。
            embed_fn: 嵌入函数（未提供时使用 TreeIndex 已有 embedding）。

        返回:
            [(node_id, score), ...] 按相似度降序。

        异常:
            ValueError: 节点未 embed 且未提供 embed_fn。
        """
        if embed_fn is None:
            raise ValueError(
                "embed_fn 为必需参数：搜索 query 需要 embed_fn 来编码。请传入 embed_fn 参数。"
            )

        # 收集所有节点的 embedding（优先使用 TreeIndex 已有 embedding）
        node_ids: list[str] = []
        embeddings: list[np.ndarray] = []

        if self._index.is_embedded:
            # 使用已有 embedding
            for nid, node in self._id_to_node.items():
                if node.embedding is not None:
                    node_ids.append(nid)
                    embeddings.append(node.embedding)
        else:
            # 使用 embed_fn 为所有节点生成 embedding
            all_ids = list(self._id_to_node.keys())
            all_texts = [_node_description(self._id_to_node[nid]) for nid in all_ids]
            all_embs = embed_fn(all_texts)  # [N, D]
            for i, nid in enumerate(all_ids):
                node_ids.append(nid)
                embeddings.append(all_embs[i])

        if not embeddings:
            return []

        node_embeddings = np.stack(embeddings, axis=0)  # [N, D]
        # 归一化（确保余弦相似度正确）
        norms = np.linalg.norm(node_embeddings, axis=1, keepdims=True)
        norms = np.where(norms == 0, 1.0, norms)
        node_embeddings = node_embeddings / norms

        # 编码 query
        query_emb = embed_fn(query)  # [1, D]

        if query_emb.ndim == 1:
            query_emb = query_emb.reshape(1, -1)
        # 归一化 query
        q_norm = np.linalg.norm(query_emb)
        if q_norm > 0:
            query_emb = query_emb / q_norm

        # 余弦相似度
        scores = (node_embeddings @ query_emb.T).squeeze()  # [N]
        if scores.ndim == 0:
            scores = scores.reshape(1)

        # 按分数排序
        scored_pairs = sorted(
            zip(node_ids, scores.tolist(), strict=True),
            key=lambda x: x[1],
            reverse=True,
        )

        # 祖先去重：如果更细粒度的子节点已入选，跳过其祖先
        deduped: list[tuple[str, float]] = []
        seen_prefixes: set[str] = set()
        for nid, score in scored_pairs:
            is_ancestor_of_seen = any(s.startswith(nid + "_") for s in seen_prefixes)
            if is_ancestor_of_seen:
                continue
            deduped.append((nid, score))
            seen_prefixes.add(nid)
            if len(deduped) >= top_k:
                break

        return deduped

    def get_node_text(
        self,
        node_id: str,
        *,
        anchor: bool = False,
    ) -> tuple[str, dict[str, str] | None]:
        """返回节点原始文本及可选的锚映射表。

        供 SearchToolDispatcher 使用：将原始文本和锚映射传给
        summarizer.summarize_node()，实现引用验证。

        参数:
            node_id: 节点 ID。
            anchor: 若 True，返回带 [cN]/[sN] 锚标的文本并构建 anchor_map。

        返回:
            (text, anchor_map) 元组。anchor=False 时 anchor_map 为 None；
            anchor=True 时 anchor_map 为 {"c1": "行文本", "s1": "字幕行", ...}。

        异常:
            KeyError: 节点不存在。
        """
        node = self._id_to_node.get(node_id)
        if node is None:
            raise KeyError(f"节点不存在: {node_id}")

        if not anchor:
            return self._node_full_text(node), None

        anchored_text = self._node_anchored_text(node)
        # 解析锚标行 "[c1] xxx" / "[s2] yyy" 构建映射
        anchor_map: dict[str, str] = {}
        anchor_pattern = re.compile(r"^\[([cs]\d+)\]\s(.+)$")
        for line in anchored_text.splitlines():
            m = anchor_pattern.match(line)
            if m:
                anchor_map[m.group(1)] = m.group(2)

        return anchored_text, anchor_map

    def get_children_info(self, node_id: str) -> list[dict[str, Any]]:
        """返回节点的直接子节点结构化信息。

        供 SearchToolDispatcher 使用：将子节点列表传给
        summarizer.summarize_children()，用于层级摘要。

        参数:
            node_id: 节点 ID。

        返回:
            子节点信息列表，每项包含 {"id", "time_range", "summary"}。
            time_range 为 (start, end) 数值元组（L3 节点退化为 (ts, ts)）。
            L3 叶子节点返回空列表。

        异常:
            KeyError: 节点不存在。
        """
        node = self._id_to_node.get(node_id)
        if node is None:
            raise KeyError(f"节点不存在: {node_id}")

        children = self._get_children(node)
        result: list[dict[str, Any]] = []
        for child in children:
            desc = _node_description(child)
            if len(desc) > 120:
                desc = desc[:120] + "..."
            result.append(
                {
                    "id": child.id,
                    "time_range": self._node_time_range_raw(child),
                    "summary": desc,
                }
            )
        return result

    def get_subtitle(self, node_id: str) -> str:
        """返回节点字幕文本。

        参数:
            node_id: 节点 ID。

        返回:
            字幕文本；无字幕或节点不存在时返回空字符串。
        """
        node = self._id_to_node.get(node_id)
        if node is None:
            return ""
        if isinstance(node, L3Node):
            return node.subtitle or ""
        return ""

    def resolve_frame_paths(self, node_ids: list[str]) -> list[Path]:
        """node_id → 帧文件路径。支持 L3（直接映射）和 L2（展开为 L3 children）。

        参数:
            node_ids: 节点 ID 列表。

        返回:
            帧文件 Path 列表。

        异常:
            KeyError: 节点不存在。
        """
        if not node_ids:
            return []

        paths: list[Path] = []
        for nid in node_ids:
            node = self._id_to_node.get(nid)
            if node is None:
                raise KeyError(f"节点不存在: {nid}")

            if isinstance(node, L3Node):
                paths.append(self._l3_frame_path(node))
            elif isinstance(node, L2Node):
                # 展开为所有 L3 子节点
                for l3 in node.children:
                    paths.append(self._l3_frame_path(l3))
            else:
                # L1 节点：展开为所有 L2 下的 L3
                assert isinstance(node, L1Node)
                for l2 in node.children:
                    for l3 in l2.children:
                        paths.append(self._l3_frame_path(l3))

        return paths

    # ------------------------------------------------------------------
    # 内部辅助方法
    # ------------------------------------------------------------------

    def _l3_frame_path(self, node: L3Node) -> Path:
        """将 L3 节点映射到帧文件路径。

        参数:
            node: L3 节点。

        返回:
            帧文件 Path。
        """
        if self._frames_dir is not None:
            # 从 node.id 中提取后缀（去掉 video_id 前缀）
            # ID 格式: {video_id}_{L1_xxx_L2_xxx_L3_xxx}
            # frame_path 格式: frames/{L1_xxx_L2_xxx_L3_xxx}.jpg
            if node.frame_path:
                return self._frames_dir / Path(node.frame_path).name
            # fallback: 从 ID 推断
            parts = node.id.split("_", 1)
            suffix = parts[1] if len(parts) > 1 else node.id
            return self._frames_dir / f"{suffix}.jpg"

        # 无 frames_dir 时使用节点自带路径
        if node.frame_path:
            return Path(node.frame_path)
        raise ValueError(f"L3 节点无 frame_path 且未提供 frames_dir: {node.id}")

    def _node_full_text(self, node: AnyNode) -> str:
        """获取节点完整文本（card 所有字段 + subtitle）。

        参数:
            node: 树节点。

        返回:
            拼接后的全文本。
        """
        card_strings = _collect_card_strings(node)
        text = "\n".join(card_strings)
        if isinstance(node, L3Node) and node.subtitle:
            text += f"\n字幕: {node.subtitle}"
        return text

    def _node_anchored_text(self, node: AnyNode) -> str:
        """获取带行号锚的节点文本。

        card 字符串逐行编 [c1]..[cN]，字幕逐行编 [s1]..[sM]。

        参数:
            node: 树节点。

        返回:
            带锚文本。
        """
        card_strings = _collect_card_strings(node)
        # 拆分内嵌换行，确保一锚一行
        card_lines: list[str] = []
        for s in card_strings:
            card_lines.extend(ln for ln in s.splitlines() if ln.strip())

        sub_lines: list[str] = []
        if isinstance(node, L3Node) and node.subtitle:
            sub_lines = [ln for ln in node.subtitle.splitlines() if ln.strip()]

        anchored: list[str] = []
        for i, line in enumerate(card_lines, 1):
            anchored.append(f"[c{i}] {line}")
        for i, line in enumerate(sub_lines, 1):
            anchored.append(f"[s{i}] {line}")

        return "\n".join(anchored)

    @staticmethod
    def _format_time_range(node: AnyNode) -> str:
        """格式化节点的时间范围。

        参数:
            node: 树节点。

        返回:
            "start-end s" 格式字符串，或 timestamp，或 "N/A"。
        """
        if isinstance(node, (L1Node, L2Node)) and node.time_range:
            return f"{node.time_range[0]:.1f}-{node.time_range[1]:.1f}s"
        if isinstance(node, L3Node) and node.timestamp is not None:
            return f"{node.timestamp:.1f}s"
        return "N/A"

    @staticmethod
    def _node_time_range_raw(node: AnyNode) -> tuple[float, float]:
        """提取节点时间范围的原始数值元组。

        L1/L2 返回 time_range 元组；L3 退化为 (timestamp, timestamp)；
        全部为 None 时兜底 (0.0, 0.0)。

        参数:
            node: 树节点。

        返回:
            (start, end) 秒级数值元组。
        """
        if isinstance(node, (L1Node, L2Node)) and node.time_range:
            return node.time_range
        if isinstance(node, L3Node) and node.timestamp is not None:
            return (node.timestamp, node.timestamp)
        return (0.0, 0.0)

    @staticmethod
    def _get_children(node: AnyNode) -> list[AnyNode]:
        """获取节点的直接子节点列表。

        参数:
            node: 树节点。

        返回:
            子节点列表（L3 节点返回空列表）。
        """
        if isinstance(node, L1Node):
            return list(node.children)
        if isinstance(node, L2Node):
            return list(node.children)
        return []