#!/usr/bin/env python3
"""批量并行建树入口：多视频并发构建三层 TreeIndex。

并发模型（Spec-2）:
    视频级 Semaphore(video_concurrency) + gather —— 复刻 repair_trees.py 惯例；
    全局共享一个 API Semaphore(api_concurrency) 注入所有 VideoTreeBuilder，
    端点压力与单视频建树完全一致，吞吐提升来自非 API 阶段跨视频重叠。

用法:
    conda activate Video-Tree-TRM
    python tools/build_trees.py --videos-dir <dir> [--out-dir store/videos]
        [--srt-dir <dir>] [--video-concurrency 16] [--limit 0]

api_concurrency 为工程配置，从 .env 读取 TREE_BUILD_API_CONCURRENCY（默认 16）。
app/core/adapters 不 import 此脚本。
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import sys
import time
from pathlib import Path
from typing import TYPE_CHECKING

# 确保项目根目录在 sys.path 中
PROJECT_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

import yaml  # noqa: E402
from dotenv import load_dotenv  # noqa: E402
from loguru import logger  # noqa: E402

load_dotenv(PROJECT_ROOT / ".env")

from app.tree.config import TreeConfig  # noqa: E402
from app.tree.index import TreeIndex  # noqa: E402
from app.tree.subtitle import parse_srt  # noqa: E402
from app.tree.video_builder import VideoTreeBuilder  # noqa: E402

if TYPE_CHECKING:
    from adapters.llm import GovernedLLMClient
    from adapters.vlm import GovernedVLMClient
    from app.tree.subtitle import SRTEntry

# ---------------------------------------------------------------------------
# 日志配置：不缓存，立即输出
# ---------------------------------------------------------------------------

logger.remove()
logger.add(
    sys.stderr,
    format="{time:HH:mm:ss} | {level:<7} | {message}",
    level="DEBUG",
    colorize=True,
)
logger.add(
    PROJECT_ROOT / "logs" / "build_trees.log",
    format="{time:YYYY-MM-DD HH:mm:ss} | {level:<7} | {message}",
    level="DEBUG",
    rotation="50 MB",
)

# ---------------------------------------------------------------------------
# 断点续跑 — progress 文件管理（复刻 repair_trees.py 惯例）
# ---------------------------------------------------------------------------

PROGRESS_FILE = "build_progress.json"

_VIDEO_SUFFIXES = frozenset({".mp4", ".mkv", ".avi", ".webm"})


def load_progress(path: Path) -> set[str]:
    """读取 progress 文件，返回已完成视频 ID 集合。

    参数:
        path: progress JSON 文件路径。

    返回:
        已完成视频 ID 集合。文件不存在或损坏时返回空集。
    """
    if not path.exists():
        return set()
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
        return set(data.get("finished_video_ids", []))
    except (json.JSONDecodeError, KeyError, TypeError, AttributeError):
        logger.warning("progress 文件损坏，忽略: {}", path)
        return set()


async def save_progress(path: Path, lock: asyncio.Lock, vid: str) -> None:
    """原子追加一个视频 ID 到 progress 文件。

    参数:
        path: progress JSON 文件路径。
        lock: asyncio.Lock，防并发读改写丢更新。
        vid: 要追加的视频 ID。
    """
    async with lock:
        finished = load_progress(path)
        finished.add(vid)
        tmp = path.with_suffix(".tmp")
        tmp.write_text(
            json.dumps({"finished_video_ids": sorted(finished)}, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        os.replace(str(tmp), str(path))


# ---------------------------------------------------------------------------
# 待建发现与完整性校验
# ---------------------------------------------------------------------------


def _tree_is_complete(tree_path: Path) -> bool:
    """判断 tree.json 是否存在且可加载为非空树。

    参数:
        tree_path: tree.json 路径。

    返回:
        True 表示完整（跳过重建）；文件缺失/损坏/空树返回 False。
    """
    if not tree_path.exists():
        return False
    try:
        index = TreeIndex.load_json(str(tree_path))
    except (json.JSONDecodeError, KeyError, TypeError, ValueError, AssertionError) as exc:
        logger.warning("tree.json 无法加载，视为不完整: {} ({})", tree_path, exc)
        return False
    return len(index.roots) > 0


def _discover_pending(
    videos_dir: Path,
    out_dir: Path,
    finished: set[str],
) -> list[Path]:
    """扫描视频目录，返回待建视频文件列表（按文件名排序）。

    跳过条件：video_id 在 progress 中，或 out_dir/<video_id>/tree.json 完整。

    参数:
        videos_dir: 视频文件目录。
        out_dir: 树输出根目录。
        finished: progress 中已完成的视频 ID 集合。

    返回:
        待建视频文件路径列表。
    """
    pending: list[Path] = []
    for f in sorted(videos_dir.iterdir()):
        if not f.is_file() or f.suffix.lower() not in _VIDEO_SUFFIXES:
            continue
        vid = f.stem
        if vid in finished:
            continue
        if _tree_is_complete(out_dir / vid / "tree.json"):
            continue
        pending.append(f)
    return pending


def _find_srt_entries(video_path: Path, srt_dir: Path) -> list[SRTEntry] | None:
    """按视频同名规则查找并解析 SRT 字幕。

    参数:
        video_path: 视频文件路径。
        srt_dir: SRT 目录。

    返回:
        SRTEntry 列表；无同名 .srt 时返回 None。
    """
    srt_path = srt_dir / f"{video_path.stem}.srt"
    if not srt_path.exists():
        return None
    return parse_srt(str(srt_path))


# ---------------------------------------------------------------------------
# LLM/VLM 客户端构建（复刻 repair_trees.py 惯例）
# ---------------------------------------------------------------------------


def _build_clients(api_concurrency: int) -> tuple[GovernedLLMClient, GovernedVLMClient]:
    """构建 GovernedLLMClient（LLM + VLM），熔断阈值随 API 并发缩放。

    参数:
        api_concurrency: 全局 API 并发上限（熔断阈值取 max(cfg, api_concurrency*2)）。

    返回:
        (llm_client, vlm_client) 元组。
    """
    from adapters.breaker import CircuitBreaker
    from adapters.llm import GovernedLLMClient
    from adapters.telemetry import SQLiteTelemetryRecorder
    from adapters.vlm import GovernedVLMClient

    (PROJECT_ROOT / "logs").mkdir(exist_ok=True)
    telemetry = SQLiteTelemetryRecorder(str(PROJECT_ROOT / "logs" / "build_trees_telemetry.db"))

    breaker_threshold = int(os.getenv("LLM_CIRCUIT_BREAKER_THRESHOLD", "5"))
    breaker_threshold = max(breaker_threshold, api_concurrency * 2)
    breaker_cooldown = int(os.getenv("LLM_CIRCUIT_BREAKER_COOLDOWN", "60"))
    timeout_s = float(os.getenv("LLM_TIMEOUT", "120"))
    max_retries = int(os.getenv("LLM_MAX_RETRIES", "3"))
    base_delay = float(os.getenv("LLM_RETRY_BASE_DELAY", "2.0"))
    max_delay = float(os.getenv("LLM_RETRY_MAX_DELAY", "30.0"))
    ttft = float(os.getenv("LLM_TTFT_TIMEOUT", "30"))
    inter_token = float(os.getenv("LLM_INTER_TOKEN_TIMEOUT", "15"))

    llm = GovernedLLMClient(
        model=os.environ["SEARCH_LLM_MODEL"],
        base_url=os.environ["SEARCH_LLM_BASE_URL"],
        api_key=os.environ["SEARCH_LLM_API_KEY"],
        provider="deepseek",
        thinking=False,
        breaker=CircuitBreaker(fail_threshold=breaker_threshold, cooldown_s=breaker_cooldown),
        cache=None,
        telemetry=telemetry,
        timeout_s=timeout_s,
        ttft_timeout_s=ttft,
        inter_token_timeout_s=inter_token,
        max_retries=max_retries,
        retry_base_delay_s=base_delay,
        retry_max_delay_s=max_delay,
    )
    vlm_base = GovernedLLMClient(
        model=os.environ["VL_LLM_MODEL"],
        base_url=os.environ["VL_LLM_BASE_URL"],
        api_key=os.environ["VL_LLM_API_KEY"],
        provider="qwen",
        thinking=False,
        breaker=CircuitBreaker(fail_threshold=breaker_threshold, cooldown_s=breaker_cooldown),
        cache=None,
        telemetry=telemetry,
        timeout_s=timeout_s,
        ttft_timeout_s=ttft,
        inter_token_timeout_s=inter_token,
        max_retries=max_retries,
        retry_base_delay_s=base_delay,
        retry_max_delay_s=max_delay,
    )
    return llm, GovernedVLMClient(vlm_base)


# ---------------------------------------------------------------------------
# 主编排
# ---------------------------------------------------------------------------


async def main_async(args: argparse.Namespace) -> None:
    """异步主流程：视频级并发建树 + 全局共享 API 信号量。

    参数:
        args: CLI 参数（videos_dir/out_dir/srt_dir/video_concurrency/limit/progress_path）。
    """
    videos_dir = Path(args.videos_dir)
    out_dir = Path(args.out_dir)
    srt_dir = Path(args.srt_dir)
    assert videos_dir.is_dir(), f"视频目录不存在: {videos_dir}"

    api_concurrency = int(os.getenv("TREE_BUILD_API_CONCURRENCY", "16"))
    if api_concurrency < 1:
        raise ValueError(f"TREE_BUILD_API_CONCURRENCY 必须 >= 1，当前值: {api_concurrency}")
    if args.video_concurrency < 1:
        raise ValueError(f"--video-concurrency 必须 >= 1，当前值: {args.video_concurrency}")

    # Phase 1: 待建发现（progress + 完整性双重跳过）
    progress_path = Path(args.progress_path)
    finished = load_progress(progress_path)
    if finished:
        logger.info("progress 已记录 {} 个完成视频", len(finished))
    pending = _discover_pending(videos_dir, out_dir, finished)
    if args.limit > 0:
        pending = pending[: args.limit]
    logger.info(
        "待建 {} 个视频, video_concurrency={}, api_concurrency={}",
        len(pending),
        args.video_concurrency,
        api_concurrency,
    )
    if not pending:
        return

    # Phase 2: 客户端与共享信号量
    llm, vlm = _build_clients(api_concurrency)
    with open(PROJECT_ROOT / "config" / "default.yaml", encoding="utf-8") as f:
        tree_cfg = TreeConfig.from_dict(yaml.safe_load(f)["tree"])
    api_sem = asyncio.Semaphore(api_concurrency)
    video_sem = asyncio.Semaphore(args.video_concurrency)
    progress_lock = asyncio.Lock()
    start_time = time.time()
    completed = 0
    failed: list[str] = []

    # Phase 3: 视频级并发编排（复刻 repair_trees 模式）
    async def _build_one(video_path: Path) -> None:
        nonlocal completed
        async with video_sem:
            vid = video_path.stem
            logger.info("开始建树 {}", vid)
            try:
                builder = VideoTreeBuilder(vlm=vlm, llm=llm, config=tree_cfg, api_semaphore=api_sem)
                srt_entries = _find_srt_entries(video_path, srt_dir)
                index = await builder.build_async(str(video_path), srt_entries)
                tree_path = out_dir / vid / "tree.json"
                tree_path.parent.mkdir(parents=True, exist_ok=True)
                index.save_json(str(tree_path))
                await save_progress(progress_path, progress_lock, vid)
                completed += 1
            except Exception as exc:
                # WHY: 单视频错误隔离——一个视频失败不拖垮整批；
                # 失败清单汇总上报，vid 不进 progress，下次重跑自动重建。
                logger.error("建树失败 {} ({}): {}", vid, type(exc).__name__, exc)
                failed.append(vid)
                return
            if completed % 5 == 0:
                elapsed = time.time() - start_time
                rate = completed / elapsed * 60 if elapsed > 0 else 0
                logger.info(
                    "进度: {}/{}, 已用 {:.0f}s, 速率 {:.2f} 视频/分钟",
                    completed,
                    len(pending),
                    elapsed,
                    rate,
                )

    await asyncio.gather(*[asyncio.create_task(_build_one(p)) for p in pending])

    # Phase 4: 汇总
    elapsed = time.time() - start_time
    logger.info(
        "批量建树完成: 成功 {}, 失败 {}, 总耗时 {:.0f}s{}",
        completed,
        len(failed),
        elapsed,
        f", 失败清单: {failed}" if failed else "",
    )


def main() -> None:
    """同步入口。"""
    parser = argparse.ArgumentParser(description="批量并行建树")
    parser.add_argument("--videos-dir", type=str, required=True, help="视频文件目录")
    parser.add_argument("--out-dir", type=str, default="store/videos", help="树输出根目录")
    parser.add_argument("--srt-dir", type=str, default="", help="SRT 目录（默认同 videos-dir）")
    parser.add_argument("--video-concurrency", type=int, default=16, help="同时在建视频数")
    parser.add_argument("--limit", type=int, default=0, help="只建前 N 个（0=全部，烟测用）")
    parser.add_argument(
        "--progress-path",
        type=str,
        default=str(PROJECT_ROOT / "logs" / PROGRESS_FILE),
        dest="progress_path",
        help="progress 文件路径",
    )
    args = parser.parse_args()
    if not args.srt_dir:
        args.srt_dir = args.videos_dir
    asyncio.run(main_async(args))


if __name__ == "__main__":
    main()