"""
TF Tree Manager - 仅调度：持有 transforms/order/graph，对外 API 委托给 transform、yaml_io、graph_ops、transform_math。
"""
from pathlib import Path
from typing import Dict, List, Optional, Set

import numpy as np

from .transform import Transform
from .yaml_io import load_extrinsics_and_order, save_extrinsics_and_order
from .graph_ops import build_bidirectional_graph, find_path_bfs, detect_directed_cycles
from .transform_math import get_transform_along_path, compute_all_from_root
from .tftree_errors import TFTreeValidationError, TFTreeKeyError


class TFTreeManager:
    """TF 树管理器：调度层，状态 + 委托。"""

    def __init__(self, yaml_dir: str):
        self.yaml_dir = Path(yaml_dir)
        self.transforms: Dict[str, Transform] = {}
        self.order: List[str] = []
        self.graph: Dict[str, Set[str]] = {}

    def load_from_yaml(self):
        """从 YAML 目录加载外参与顺序，构建图。"""
        self.transforms, self.order = load_extrinsics_and_order(self.yaml_dir)
        self.graph = build_bidirectional_graph(self.transforms)

    def _build_graph(self):
        """重建无向图（invert 后调用）。"""
        self.graph = build_bidirectional_graph(self.transforms)

    def detect_cycles(self) -> List[List[str]]:
        """有向图环检测，委托 graph_ops。"""
        return detect_directed_cycles(self.transforms)

    def get_transform(
        self, source: str, target: str, root: Optional[str] = None
    ) -> Optional[np.ndarray]:
        """沿路径累积变换矩阵。"""
        if source == target:
            return np.eye(4)
        path = find_path_bfs(self.graph, source, target)
        if not path:
            return None
        return get_transform_along_path(self.transforms, path)

    def compute_all_transforms_from_root(self, root: str) -> Dict[str, np.ndarray]:
        """从 root 到所有可达节点的变换。"""
        if root not in self.graph:
            raise TFTreeValidationError(f"根节点 {root} 不存在于图中", {"root": root})
        return compute_all_from_root(self.transforms, self.graph, root)

    def update_transform(
        self,
        source: str,
        target: str,
        translation: Optional[np.ndarray] = None,
        rotation: Optional[np.ndarray] = None,
    ):
        """更新指定边的平移/旋转。"""
        key = f"{source}->{target}"
        if key not in self.transforms:
            raise TFTreeKeyError(f"变换 {key} 不存在", {"key": key})
        t = self.transforms[key]
        if translation is not None:
            t.translation = translation.copy()
        if rotation is not None:
            t.rotation = rotation.copy()

    def invert_transform(self, source: str, target: str):
        """将边 source->target 求逆为 target->source。"""
        key = f"{source}->{target}"
        if key not in self.transforms:
            raise TFTreeKeyError(f"变换 {key} 不存在", {"key": key})
        t = self.transforms[key]
        inv = t.inverse()
        del self.transforms[key]
        new_key = f"{target}->{source}"
        self.transforms[new_key] = inv
        if key in self.order:
            idx = self.order.index(key)
            self.order[idx] = new_key
        self._build_graph()

    def reorder_transforms(self, new_order: List[str]):
        """重排顺序列表。"""
        if set(new_order) != set(self.order):
            raise TFTreeValidationError(
                "新顺序必须包含所有现有变换", {"order": list(self.order)}
            )
        self.order = new_order.copy()

    def save_to_yaml(self):
        """写回 extrinsics 与 order_manifest。"""
        save_extrinsics_and_order(self.yaml_dir, self.transforms, self.order)

    def get_all_frames(self) -> Set[str]:
        """所有节点（frame）集合。"""
        return set(self.graph.keys())

    def print_tree(self, root: Optional[str] = None):
        """打印树结构（调试）。"""
        if not self.graph:
            print("图为空")
            return
        if root is None:
            all_targets = set()
            all_sources = set()
            for key in self.transforms:
                s, t = key.split("->")
                all_sources.add(s)
                all_targets.add(t)
            roots = all_sources - all_targets
            root = list(roots)[0] if roots else list(self.graph.keys())[0]
        print(f"树结构（根节点: {root}）:")
        visited = set()

        def print_node(node: str, level: int):
            if node in visited:
                print("  " * level + f"└─ {node} (已访问)")
                return
            visited.add(node)
            print("  " * level + f"└─ {node}")
            for neighbor in sorted(self.graph.get(node, [])):
                if neighbor not in visited:
                    print_node(neighbor, level + 1)

        print_node(root, 0)


def main():
    """示例用法。"""
    import sys
    if len(sys.argv) < 2:
        print("用法: tftree-manager <yaml_dir>")
        return
    yaml_dir = sys.argv[1]
    manager = TFTreeManager(yaml_dir)
    manager.load_from_yaml()
    print(f"加载了 {len(manager.transforms)} 个变换")
    print(f"坐标系: {manager.get_all_frames()}")
    cycles = manager.detect_cycles()
    if cycles:
        print(f"检测到 {len(cycles)} 个环:")
        for i, c in enumerate(cycles):
            print(f"  环 {i+1}: {' -> '.join(c)}")
    else:
        print("未检测到环")
    manager.print_tree()
    frames = manager.get_all_frames()
    if frames:
        root = list(frames)[0]
        result = manager.compute_all_transforms_from_root(root)
        print(f"从根节点 {root} 到所有节点的变换: {len(result)} 个")
    print()


if __name__ == "__main__":
    main()
