"""
YAML 读写：order_manifest + extrinsics 目录。
输入 yaml_dir，返回/接收 (transforms_dict, order_list)，不持有图。
"""
import yaml
import numpy as np
from pathlib import Path
from typing import Dict, List, Tuple

from .transform import Transform
from .tftree_errors import TFTreeNotFoundError


def load_extrinsics_and_order(yaml_dir: Path) -> Tuple[Dict[str, "Transform"], List[str]]:
    """
    从 yaml_dir 读取 order_manifest 与 extrinsics/*.yaml，返回 (transforms, order)。
    外参目录不存在时抛出 TFTreeNotFoundError。
    """
    extrinsics_dir = yaml_dir / "extrinsics"
    if not extrinsics_dir.exists():
        raise TFTreeNotFoundError(f"外参目录不存在: {extrinsics_dir}", {"path": str(extrinsics_dir)})

    order: List[str] = []
    order_file = yaml_dir / "order_manifest.yaml"
    if order_file.exists():
        with open(order_file, "r", encoding="utf-8") as f:
            order_data = yaml.safe_load(f)
            if order_data:
                if "extrinsics_order" in order_data:
                    order = list(order_data["extrinsics_order"])
                elif "extrinsics" in order_data:
                    for item in order_data["extrinsics"]:
                        rel_file = item.get("file", "")
                        path = yaml_dir / rel_file
                        if path.exists():
                            with open(path, "r", encoding="utf-8") as fp:
                                data = yaml.safe_load(fp) or {}
                            src = data.get("source_frame", "")
                            tgt = data.get("target_frame", "")
                            if src or tgt:
                                order.append(f"{src}->{tgt}")

    transforms: Dict[str, Transform] = {}
    for yaml_file in sorted(extrinsics_dir.glob("*.yaml")):
        with open(yaml_file, "r", encoding="utf-8") as f:
            data = yaml.safe_load(f)
            if data:
                source = data["source_frame"]
                target = data["target_frame"]
                translation = np.array([
                    data["translation"]["x"],
                    data["translation"]["y"],
                    data["translation"]["z"],
                ])
                rotation = np.array([
                    data["rotation"]["qx"],
                    data["rotation"]["qy"],
                    data["rotation"]["qz"],
                    data["rotation"]["qw"],
                ])
                t = Transform(source, target, translation, rotation)
                key = f"{source}->{target}"
                transforms[key] = t
                if key not in order:
                    order.append(key)

    return transforms, order


def save_extrinsics_and_order(
    yaml_dir: Path,
    transforms: Dict[str, "Transform"],
    order: List[str],
) -> None:
    """
    将 transforms 与 order 写回 yaml_dir：extrinsics/000.yaml 等 + order_manifest.yaml。
    """
    extrinsics_dir = yaml_dir / "extrinsics"
    extrinsics_dir.mkdir(exist_ok=True)

    for yaml_file in extrinsics_dir.glob("*.yaml"):
        yaml_file.unlink()

    for i, key in enumerate(order):
        if key not in transforms:
            continue
        transform = transforms[key]
        filepath = extrinsics_dir / f"{i:03d}.yaml"
        with open(filepath, "w", encoding="utf-8") as f:
            yaml.dump(transform.to_dict(), f, default_flow_style=False, allow_unicode=True)

    order_file = yaml_dir / "order_manifest.yaml"
    order_data = {}
    if order_file.exists():
        with open(order_file, "r", encoding="utf-8") as f:
            order_data = yaml.safe_load(f) or {}
    order_data["extrinsics"] = [{"index": i, "file": f"extrinsics/{i:03d}.yaml"} for i in range(len(order))]
    order_data["extrinsics_order"] = order
    with open(order_file, "w", encoding="utf-8") as f:
        yaml.dump(order_data, f, default_flow_style=False, allow_unicode=True, sort_keys=False)
