#!/usr/bin/env python3
"""一次性格式转换：TRM4 flat tree.json -> TRM5 TreeIndex JSON。

用法: python tools/convert_flat_to_treeindex.py <src_dir> <dst_dir>

遍历 src_dir 下每个 video_id 子目录中的 tree.json（TRM4 flat 格式），
转换为 TRM5 TreeIndex 嵌套格式并写入 dst_dir 对应子目录。

app/core/adapters 不 import 此脚本。迁移完成后归档至 tools/archived/。
"""

from __future__ import annotations

import json
import sys
from datetime import datetime
from pathlib import Path
from typing import Any

# ---------------------------------------------------------------------------
# Card 字段默认值（处理 TRM4 可能缺失的字段）
# ---------------------------------------------------------------------------

_L3_CARD_DEFAULTS: dict[str, Any] = {
    "frame_summary": "",
    "visible_entities": [],
    "ongoing_actions": [],
    "visible_text": [],
    "spatial_layout": "",
    "visual_attributes": {},
}

_L2_CARD_DEFAULTS: dict[str, Any] = {
    "event_description": "",
    "entities": [],
    "actions": [],
    "action_subjects": [],
    "visible_text": [],
    "spatial_relations": "",
    "state_changes": None,
}

_L1_CARD_DEFAULTS: dict[str, Any] = {
    "scene_summary": "",
    "main_setting": "",
    "key_entities": [],
    "main_actions": [],
    "topic_keywords": [],
    "visible_text": [],
    "temporal_flow": "",
}


# ---------------------------------------------------------------------------
# Card 构建辅助
# ---------------------------------------------------------------------------


def _build_card(raw_card: dict[str, Any], defaults: dict[str, Any]) -> dict[str, Any]:
    """从 TRM4 原始 card 字典构建 TRM5 card，缺失字段用默认值填充。

    参数:
        raw_card: TRM4 tree.json 中节点的 card 字典。
        defaults: 该层级的默认值字典。

    返回:
        仅包含目标字段的 card 字典（字段集合与 defaults 一致）。
    """
    return {key: raw_card.get(key, default) for key, default in defaults.items()}


# ---------------------------------------------------------------------------
# 节点转换
# ---------------------------------------------------------------------------


def _convert_l3(
    node: dict[str, Any],
    video_id: str,
) -> dict[str, Any]:
    """将 TRM4 flat L3 节点转换为 TRM5 嵌套 L3 节点。

    参数:
        node: TRM4 flat 格式的 L3 节点字典。
        video_id: 视频 ID，用于计算 frame_path 的相对路径后缀。

    返回:
        TRM5 格式的 L3 节点字典。
    """
    node_id: str = node["node_id"]
    raw_card = node.get("card") or {}
    card = _build_card(raw_card, _L3_CARD_DEFAULTS)

    # frame_path: frames/{suffix}.jpg，suffix = node_id 去掉 video_id 前缀 + 下划线
    prefix = f"{video_id}_"
    suffix = node_id[len(prefix) :] if node_id.startswith(prefix) else node_id
    frame_path = f"frames/{suffix}.jpg"

    return {
        "id": node_id,
        "card": card,
        "timestamp": node.get("frame_timestamp"),
        "frame_path": frame_path,
        "subtitle": node.get("subtitle"),
    }


def _convert_l2(
    node: dict[str, Any],
    l3_children: list[dict[str, Any]],
) -> dict[str, Any]:
    """将 TRM4 flat L2 节点转换为 TRM5 嵌套 L2 节点。

    参数:
        node: TRM4 flat 格式的 L2 节点字典。
        l3_children: 已转换的 L3 子节点列表（按 time_range 排序）。

    返回:
        TRM5 格式的 L2 节点字典。
    """
    raw_card = node.get("card") or {}
    card = _build_card(raw_card, _L2_CARD_DEFAULTS)
    time_range = node.get("time_range")

    return {
        "id": node["node_id"],
        "card": card,
        "time_range": time_range,
        "children": l3_children,
    }


def _convert_l1(
    node: dict[str, Any],
    l2_children: list[dict[str, Any]],
) -> dict[str, Any]:
    """将 TRM4 flat L1 节点转换为 TRM5 嵌套 L1 节点。

    参数:
        node: TRM4 flat 格式的 L1 节点字典。
        l2_children: 已转换的 L2 子节点列表（按 time_range 排序）。

    返回:
        TRM5 格式的 L1 节点字典。
    """
    raw_card = node.get("card") or {}
    card = _build_card(raw_card, _L1_CARD_DEFAULTS)
    time_range = node.get("time_range")

    return {
        "id": node["node_id"],
        "card": card,
        "time_range": time_range,
        "children": l2_children,
    }


# ---------------------------------------------------------------------------
# 排序辅助
# ---------------------------------------------------------------------------


def _sort_key_time_range(node: dict[str, Any]) -> float:
    """按 time_range 的起始时间排序。

    参数:
        node: TRM4 节点字典。

    返回:
        起始时间（float），无 time_range 时返回 0.0。
    """
    tr = node.get("time_range")
    if tr and len(tr) >= 1:
        return float(tr[0])
    return 0.0


def _sort_key_timestamp(converted: dict[str, Any]) -> float:
    """按 timestamp 排序（L3 转换后的字典）。

    参数:
        converted: 已转换的 TRM5 L3 节点字典。

    返回:
        timestamp（float），无值时返回 0.0。
    """
    ts = converted.get("timestamp")
    return float(ts) if ts is not None else 0.0


# ---------------------------------------------------------------------------
# 单棵树转换
# ---------------------------------------------------------------------------


def convert_single_tree(flat_data: dict[str, Any], source_path: str) -> dict[str, Any]:
    """将单个 TRM4 flat tree.json 转换为 TRM5 TreeIndex 字典。

    参数:
        flat_data: TRM4 flat tree.json 解析后的字典。
        source_path: 原始数据路径（写入 metadata.source_path）。

    返回:
        TRM5 TreeIndex 格式的字典（可直接 json.dump 或传入 TreeIndex.from_dict）。

    异常:
        ValueError: 无法从 flat_data 中提取 video_id。
    """
    video_id = flat_data.get("video_id") or flat_data.get("videoID")
    if not video_id:
        raise ValueError("flat tree.json 中缺少 video_id / videoID 字段")

    nodes: dict[str, dict[str, Any]] = flat_data.get("nodes", {})

    # Phase 1: 按层级分组
    l1_nodes: list[dict[str, Any]] = []
    l2_nodes: list[dict[str, Any]] = []
    l3_nodes: list[dict[str, Any]] = []

    for node in nodes.values():
        level = node.get("level")
        if level == 1:
            l1_nodes.append(node)
        elif level == 2:
            l2_nodes.append(node)
        elif level == 3:
            l3_nodes.append(node)

    # Phase 2: 构建 parent -> children 映射
    # L3 按 parent_id 分组
    l3_by_parent: dict[str, list[dict[str, Any]]] = {}
    for n in l3_nodes:
        pid = n.get("parent_id", "")
        l3_by_parent.setdefault(pid, []).append(n)

    # L2 按 parent_id 分组
    l2_by_parent: dict[str, list[dict[str, Any]]] = {}
    for n in l2_nodes:
        pid = n.get("parent_id", "")
        l2_by_parent.setdefault(pid, []).append(n)

    # Phase 3: 自底向上构建嵌套结构
    # 转换 L2 -> 附带转换后的 L3 children
    converted_l2_by_id: dict[str, dict[str, Any]] = {}
    for l2 in l2_nodes:
        l2_id = l2["node_id"]
        raw_l3_children = l3_by_parent.get(l2_id, [])
        # 先转换 L3，再按 timestamp 排序
        converted_l3 = [_convert_l3(n, video_id) for n in raw_l3_children]
        converted_l3.sort(key=_sort_key_timestamp)
        converted_l2_by_id[l2_id] = _convert_l2(l2, converted_l3)

    # 转换 L1 -> 附带转换后的 L2 children
    roots: list[dict[str, Any]] = []
    l1_nodes.sort(key=_sort_key_time_range)

    for l1 in l1_nodes:
        l1_id = l1["node_id"]
        raw_l2_children = l2_by_parent.get(l1_id, [])
        raw_l2_children.sort(key=_sort_key_time_range)
        l2_children = [converted_l2_by_id[n["node_id"]] for n in raw_l2_children]
        roots.append(_convert_l1(l1, l2_children))

    # Phase 4: 构建 TreeIndex 字典
    return {
        "metadata": {
            "source_path": source_path,
            "modality": "video",
            "created_at": datetime.now().isoformat(),
        },
        "roots": roots,
    }


# ---------------------------------------------------------------------------
# 批量转换入口
# ---------------------------------------------------------------------------


def convert_directory(src_dir: str, dst_dir: str) -> tuple[int, int]:
    """批量转换目录下所有 TRM4 tree.json 到 TRM5 TreeIndex 格式。

    参数:
        src_dir: 源目录（TRM4 store/videos/），其下每个子目录含 tree.json。
        dst_dir: 目标目录（TRM5 store/videos/），保持同名子目录结构。

    返回:
        (成功数, 失败数) 元组。
    """
    src_path = Path(src_dir)
    dst_path = Path(dst_dir)

    if not src_path.is_dir():
        print(f"错误: 源目录不存在: {src_dir}", file=sys.stderr)
        sys.exit(1)

    success_count = 0
    fail_count = 0

    tree_files = sorted(src_path.glob("*/tree.json"))
    total = len(tree_files)
    print(f"发现 {total} 个 tree.json 待转换")

    for idx, tree_file in enumerate(tree_files, 1):
        video_id = tree_file.parent.name
        out_dir = dst_path / video_id
        out_file = out_dir / "tree.json"

        try:
            with open(tree_file, encoding="utf-8") as f:
                flat_data = json.load(f)

            result = convert_single_tree(flat_data, source_path=str(tree_file))

            out_dir.mkdir(parents=True, exist_ok=True)
            with open(out_file, "w", encoding="utf-8") as f:
                json.dump(result, f, ensure_ascii=False, indent=2)

            # 统计节点数
            n_l1 = len(result["roots"])
            n_l2 = sum(len(r["children"]) for r in result["roots"])
            n_l3 = sum(len(l2["children"]) for r in result["roots"] for l2 in r["children"])
            print(f"[{idx}/{total}] {video_id}: L1={n_l1}, L2={n_l2}, L3={n_l3}")
            success_count += 1

        except Exception as e:
            print(f"[{idx}/{total}] {video_id}: 失败 - {e}", file=sys.stderr)
            fail_count += 1

    return success_count, fail_count


# ---------------------------------------------------------------------------
# CLI 入口
# ---------------------------------------------------------------------------


def main() -> None:
    """CLI 入口：解析参数并执行批量转换。"""
    if len(sys.argv) != 3:
        print(
            "用法: python tools/convert_flat_to_treeindex.py <src_dir> <dst_dir>",
            file=sys.stderr,
        )
        print("  src_dir: TRM4 store/videos/ 目录（含 video_id/tree.json）", file=sys.stderr)
        print("  dst_dir: TRM5 store/videos/ 目标目录", file=sys.stderr)
        sys.exit(1)

    src_dir = sys.argv[1]
    dst_dir = sys.argv[2]

    print(f"源目录: {src_dir}")
    print(f"目标目录: {dst_dir}")
    print()

    success, fail = convert_directory(src_dir, dst_dir)

    print()
    print(f"转换完成: 成功 {success}, 失败 {fail}")
    if fail > 0:
        sys.exit(1)


if __name__ == "__main__":
    main()