#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import shutil
import subprocess
from pathlib import Path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Summarize a prepared or completed poetry compare run.")
    parser.add_argument("run_root", nargs="?", type=Path, help="Compare run directory. Defaults to the latest run under outputs/compare_poetry.")
    parser.add_argument("--compare-root", type=Path, default=Path(__file__).resolve().parents[2] / "outputs" / "compare_poetry")
    parser.add_argument("--json", action="store_true", help="Print JSON instead of a text table.")
    parser.add_argument("--ffprobe-timeout", type=float, default=8.0)
    return parser.parse_args()


def find_latest_run(compare_root: Path) -> Path:
    candidates = [path for path in compare_root.iterdir() if path.is_dir()]
    if not candidates:
        raise FileNotFoundError(f"No compare runs found under {compare_root}")
    return max(candidates, key=lambda path: path.stat().st_mtime)


def load_json(path: Path) -> dict:
    return json.loads(path.read_text(encoding="utf-8"))


def human_size(num_bytes: int | None) -> str:
    if num_bytes is None:
      return "-"
    value = float(num_bytes)
    for unit in ["B", "KB", "MB", "GB", "TB"]:
        if value < 1024.0 or unit == "TB":
            return f"{value:.1f}{unit}"
        value /= 1024.0
    return f"{num_bytes}B"


def probe_duration(path: Path, timeout: float) -> float | None:
    if not path.is_file() or shutil.which("ffprobe") is None:
        return None
    command = [
        "ffprobe",
        "-v",
        "error",
        "-show_entries",
        "format=duration",
        "-of",
        "default=noprint_wrappers=1:nokey=1",
        str(path),
    ]
    try:
        result = subprocess.run(command, capture_output=True, text=True, check=True, timeout=timeout)
    except (subprocess.CalledProcessError, subprocess.TimeoutExpired, FileNotFoundError):
        return None
    text = result.stdout.strip()
    if not text:
        return None
    try:
        return float(text)
    except ValueError:
        return None


def collect_row(run_dir: Path, ffprobe_timeout: float) -> dict:
    spec_path = run_dir / "spec.json"
    spec = load_json(spec_path) if spec_path.is_file() else {}
    planner = str(spec.get("planner") or run_dir.name.split("__", 1)[0])
    backend = str(spec.get("backend") or run_dir.name.split("__", 1)[-1])
    shared_plan_json = Path(spec.get("shared_plan_json")) if spec.get("shared_plan_json") else None
    plan_path = run_dir / "plan.json"
    final_video = run_dir / "final_poetry_teaching.mp4"
    subtitled_video = run_dir / "final_poetry_teaching_subtitled.mp4"
    launch_path = run_dir / "launch.sh"
    segment_videos = sorted((run_dir / "videos").glob("segment_*.mp4"))
    plan_segments = None
    if plan_path.is_file():
        try:
            plan_segments = len(load_json(plan_path).get("segments", []))
        except Exception:
            plan_segments = None

    if final_video.is_file():
        status = "done"
    elif plan_path.is_file() or segment_videos:
        status = "partial"
    elif launch_path.is_file():
        status = "prepared"
    else:
        status = "missing"

    return {
        "planner": planner,
        "image_backend": str(spec.get("image_backend") or "-"),
        "backend": backend,
        "status": status,
        "shared_plan_ready": bool(shared_plan_json and shared_plan_json.is_file()),
        "plan_ready": plan_path.is_file(),
        "plan_segments": plan_segments,
        "segment_videos": len(segment_videos),
        "final_video": str(final_video) if final_video.is_file() else None,
        "final_video_size_bytes": final_video.stat().st_size if final_video.is_file() else None,
        "final_video_duration_seconds": probe_duration(final_video, ffprobe_timeout),
        "subtitled_video": str(subtitled_video) if subtitled_video.is_file() else None,
        "run_dir": str(run_dir),
    }


def print_table(rows: list[dict], run_root: Path, manifest: dict | None) -> None:
    print(f"Compare run: {run_root}")
    if manifest:
        print(f"Experiment JSON: {manifest.get('experiment_json', '-')}")
        print(f"Requested backends: {', '.join(manifest.get('backends_requested', []))}")
        print(f"Planner variants: {', '.join(manifest.get('planner_variants', []))}")
        print(f"Image backend: {manifest.get('image_backend', '-')}")
        print(f"Persistent video backends: {manifest.get('persistent_video_backends')}")
        print(f"Dry run: {manifest.get('dry_run')}")
    headers = ["planner", "backend", "image", "status", "shared_plan", "plan", "segments", "duration", "size"]
    table = []
    for row in rows:
        duration = row["final_video_duration_seconds"]
        duration_text = f"{duration:.1f}s" if duration is not None else "-"
        table.append([
            row["planner"],
            row["backend"],
            row["image_backend"],
            row["status"],
            "yes" if row["shared_plan_ready"] else "no",
            "yes" if row["plan_ready"] else "no",
            str(row["segment_videos"]),
            duration_text,
            human_size(row["final_video_size_bytes"]),
        ])
    widths = [len(header) for header in headers]
    for line in table:
        for idx, cell in enumerate(line):
            widths[idx] = max(widths[idx], len(cell))

    header_line = "  ".join(header.ljust(widths[idx]) for idx, header in enumerate(headers))
    print(header_line)
    print("  ".join("-" * width for width in widths))
    for line, row in zip(table, rows):
        print("  ".join(cell.ljust(widths[idx]) for idx, cell in enumerate(line)))
        if row["final_video"]:
            print(f"  final: {row['final_video']}")
        else:
            print(f"  run:   {row['run_dir']}")


def main() -> None:
    args = parse_args()
    run_root = args.run_root.resolve() if args.run_root else find_latest_run(args.compare_root.resolve())
    manifest_path = run_root / "manifest.json"
    manifest = load_json(manifest_path) if manifest_path.is_file() else None

    rows = []
    for run_dir in sorted(path for path in run_root.iterdir() if path.is_dir() and path.name != "shared_plan"):
        rows.append(collect_row(run_dir, args.ffprobe_timeout))

    if args.json:
        print(json.dumps({"run_root": str(run_root), "manifest": manifest, "rows": rows}, ensure_ascii=False, indent=2))
    else:
        print_table(rows, run_root, manifest)


if __name__ == "__main__":
    main()