"""
批量视频建树脚本（视频间并行 + 视频内 asyncio 真并发）
==========================================================
扫描指定目录下所有 MP4，跳过已有 JSON 树的视频，
使用 ThreadPoolExecutor 进行视频间并行（--jobs，默认 1），
每个视频内部使用 asyncio + Semaphore(concurrency=16) 真并发 VLM 调用。

并发架构::

    外层: ThreadPoolExecutor(max_workers=jobs=1)  — 视频间并行（默认 1 路）
    内层: asyncio.run(_build_async())              — 每视频独立事件循环
          asyncio.Semaphore(concurrency=16)        — 限制同时在途 VLM 请求数
    总最大 VLM 并发: jobs × concurrency = 1 × 16 = 16

用法::

    # 默认 1 路 16 并发
    conda run -n Video-Tree-TRM python scripts/build_trees_batch.py

    # 指定路数
    conda run -n Video-Tree-TRM python scripts/build_trees_batch.py --jobs 2

    # 指定目录和配置
    conda run -n Video-Tree-TRM python scripts/build_trees_batch.py \\
        --video-dir data/videomme/videos \\
        --config config/videomme.yaml \\
        --jobs 3
"""

from __future__ import annotations

import argparse
import sys
import time
from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor
from concurrent.futures import wait as cfwait
from pathlib import Path
from typing import Dict, List, Optional, Tuple

# 项目根目录加入 sys.path
sys.path.insert(0, str(Path(__file__).parent.parent))

from video_tree_trm.config import Config
from video_tree_trm.llm_client import LLMClient
from video_tree_trm.video_tree_builder import VideoTreeBuilder
from utils.logger_system import log_msg


def _build_one(
    video_path: Path,
    cfg: Config,
) -> Tuple[str, bool, str]:
    """构建单个视频的 TreeIndex 并保存 JSON。

    参数:
        video_path: 视频文件绝对路径。
        cfg: 已加载的配置对象（由主进程共享，仅读取）。

    返回:
        (stem, success, message) 三元组。

    实现细节:
        每次调用独立初始化 LLMClient（避免多线程共享同一 httpx.Client 内部状态），
        使用 VideoTreeBuilder.build() 内部的异步事件循环（L2→L3→L1 链式并发）。
    """
    stem = video_path.stem
    try:
        progress_dir = Path(cfg.tree.cache_dir) / "progress"
        progress_path = progress_dir / f"{stem}.json"
        if progress_path.is_file():
            log_msg("INFO", "检测到中间进度，启用断点续跑", stem=stem, progress_path=str(progress_path))

        # 每线程独立 LLMClient（httpx.Client 线程安全，但独立更稳健）
        vlm = LLMClient(cfg.vlm)
        builder = VideoTreeBuilder(vlm, cfg.tree)
        tree = builder.build(str(video_path))

        # 保存 JSON
        cache_dir = Path(cfg.tree.cache_dir)
        cache_dir.mkdir(parents=True, exist_ok=True)
        out_path = cache_dir / f"{stem}_video.json"
        tree.save_json(str(out_path))

        l1 = len(tree.roots)
        l2 = sum(len(r.children) for r in tree.roots)
        l3 = sum(len(l2n.children) for r in tree.roots for l2n in r.children)
        msg = f"L1={l1} L2={l2} L3={l3} → {out_path}"
        log_msg("INFO", "视频建树完成", stem=stem, l1=l1, l2=l2, l3=l3)
        return stem, True, msg

    except Exception as e:  # noqa: BLE001
        log_msg("ERROR", "视频建树失败", stem=stem, error=str(e))
        return stem, False, str(e)


def main() -> None:
    """批量建树主函数：视频间 ThreadPoolExecutor 并行 + 视频内异步事件循环。"""
    parser = argparse.ArgumentParser(description="批量视频建树（视频间并行）")
    parser.add_argument(
        "--video-dir",
        default="data/videomme/videos",
        help="MP4 视频目录（默认: data/videomme/videos）",
    )
    parser.add_argument(
        "--config",
        default="config/videomme.yaml",
        help="配置文件路径（默认: config/videomme.yaml）",
    )
    parser.add_argument(
        "--jobs",
        type=int,
        default=1,
        help="视频间并行数（默认: 1，每路视频内 asyncio Semaphore(concurrency) 并发 VLM）",
    )
    args = parser.parse_args()

    # Phase 1: 加载配置（所有线程共享，只读）
    cfg = Config.load(args.config)
    log_msg(
        "INFO",
        "批量建树配置",
        video_dir=args.video_dir,
        cache_dir=cfg.tree.cache_dir,
        jobs=args.jobs,
        intra_concurrency=cfg.tree.concurrency,
    )

    # Phase 2: 扫描视频 + 过滤已有树
    video_dir = Path(args.video_dir)
    assert video_dir.is_dir(), f"视频目录不存在: {video_dir}"
    all_videos: List[Path] = sorted(video_dir.glob("*.mp4"))

    cache_dir = Path(cfg.tree.cache_dir)
    progress_dir = cache_dir / "progress"
    with_progress: List[Path] = []
    without_progress: List[Path] = []
    for v in all_videos:
        if (cache_dir / f"{v.stem}_video.json").exists():
            continue
        progress_path = progress_dir / f"{v.stem}.json"
        if progress_path.is_file():
            with_progress.append(v)
        else:
            without_progress.append(v)
    pending: List[Path] = with_progress + without_progress
    skipped = len(all_videos) - len(pending)

    print(f"\n===== 批量建树 =====")
    print(f"  总视频数:    {len(all_videos)}")
    print(f"  已跳过(已建): {skipped}")
    print(f"  待处理:      {len(pending)}")
    print(f"  视频间并行:  {args.jobs}")
    print(f"  视频内并发:  {cfg.tree.concurrency}")
    print(f"  输出目录:    {cache_dir}\n")

    if not pending:
        print("所有视频均已建树，无需处理。")
        return

    # Phase 3: 异步视频间并行（ThreadPoolExecutor + FIRST_COMPLETED 事件循环）
    built: List[str] = []
    failed: List[Tuple[str, str]] = []
    start_total = time.time()

    pending_futures: Dict[Future, Path] = {}
    pending_queue = list(pending)  # 待提交队列
    pool = ThreadPoolExecutor(max_workers=args.jobs)

    # 初始填满 jobs 个任务
    while pending_queue and len(pending_futures) < args.jobs:
        video = pending_queue.pop(0)
        fut = pool.submit(_build_one, video, cfg)
        pending_futures[fut] = video
        print(f"[提交] {video.stem} ({len(built)+len(failed)+len(pending_futures)}/{len(pending)})")

    # 事件循环：完成一个，补一个
    done_count = 0
    while pending_futures:
        done, _ = cfwait(list(pending_futures), return_when=FIRST_COMPLETED)
        for fut in done:
            video = pending_futures.pop(fut)
            stem, success, msg = fut.result()
            done_count += 1
            elapsed = time.time() - start_total
            if success:
                built.append(stem)
                print(f"[OK {done_count}/{len(pending)}] {stem}  {msg}  (累计 {elapsed:.0f}s)")
            else:
                failed.append((stem, msg))
                print(f"[FAIL {done_count}/{len(pending)}] {stem}  {msg}")

            # 补充下一个任务（非阻塞提交）
            if pending_queue:
                next_video = pending_queue.pop(0)
                next_fut = pool.submit(_build_one, next_video, cfg)
                pending_futures[next_fut] = next_video
                print(f"[提交] {next_video.stem} ({done_count+len(pending_futures)}/{len(pending)})")

    pool.shutdown(wait=False)

    # Phase 4: 汇总
    total_elapsed = time.time() - start_total
    print(f"\n===== 汇总 =====")
    print(f"  成功: {len(built)}")
    print(f"  失败: {len(failed)}")
    print(f"  跳过: {skipped}")
    print(f"  总耗时: {total_elapsed:.1f}s")
    if failed:
        print("\n  失败列表:")
        for stem, err in failed:
            print(f"    {stem}: {err}")

    log_msg(
        "INFO",
        "批量建树完成",
        built=len(built),
        failed=len(failed),
        skipped=skipped,
        elapsed_s=round(total_elapsed, 1),
    )


if __name__ == "__main__":
    main()
