"""
Walk extracted videos, verify each is decodable with decord and has ≥ FRAMES_MIN
frames. Outputs a per-source manifest jsonl with one valid entry per video:

    {"src": "tb", "rel_path": "long_video/ActivityNet/abc.mp4",
     "frames": 240, "fps": 30.0, "duration_s": 8.0}

Multiprocessing; safe to interrupt and re-run (writes incrementally).
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path

# Decord on CPU only — we don't need GPU for frame-count check.
os.environ.setdefault("DECORD_EOF_RETRY_MAX", "1024")

ROOT = Path("/mnt/local-fast/opd_zt")
VIDEOS = ROOT / "data" / "videos"
FILTERED = ROOT / "data" / "filtered"
FRAMES_MIN = 32

VIDEO_EXTS = {".mp4", ".mkv", ".webm", ".avi", ".mov", ".m4v"}


def probe(path: Path) -> dict | None:
    """Return manifest entry or None if invalid / too short / unreadable."""
    try:
        # Lazy import inside worker to avoid heavy load on parent.
        from decord import VideoReader, cpu
        vr = VideoReader(str(path), ctx=cpu(0), num_threads=1)
        n = len(vr)
        if n < FRAMES_MIN:
            return None
        fps = float(vr.get_avg_fps())
        # Sanity: try decoding the last frame to catch corrupt tails.
        _ = vr[n - 1].asnumpy()
        duration = n / fps if fps > 0 else 0.0
        return {
            "frames": int(n),
            "fps": fps,
            "duration_s": duration,
        }
    except Exception:
        return None


def worker(src_root: str, rel_path: str) -> tuple[str, dict | None]:
    full = Path(src_root) / rel_path
    info = probe(full)
    if info is None:
        return rel_path, None
    info["rel_path"] = rel_path
    return rel_path, info


def gather_videos(src_root: Path) -> list[str]:
    files: list[str] = []
    for p in src_root.rglob("*"):
        if p.is_file() and p.suffix.lower() in VIDEO_EXTS:
            files.append(str(p.relative_to(src_root)))
    return files


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--src", choices=["tb", "lv178k"], required=True)
    parser.add_argument("--workers", type=int, default=os.cpu_count() or 8)
    parser.add_argument("--limit", type=int, default=0, help="0 = no limit")
    args = parser.parse_args()

    src_root = VIDEOS / args.src
    FILTERED.mkdir(parents=True, exist_ok=True)
    out_path = FILTERED / f"{args.src}_manifest.jsonl"
    bad_path = FILTERED / f"{args.src}_invalid.txt"

    if not src_root.exists():
        print(f"FATAL: source dir does not exist: {src_root}", file=sys.stderr)
        sys.exit(2)

    # Resume support: if manifest exists, skip files already there.
    seen: set[str] = set()
    if out_path.exists():
        with out_path.open() as f:
            for line in f:
                try:
                    seen.add(json.loads(line)["rel_path"])
                except Exception:
                    pass
    if bad_path.exists():
        with bad_path.open() as f:
            for line in f:
                seen.add(line.strip())
    print(f"[scan] gathering videos under {src_root}", flush=True)
    all_files = gather_videos(src_root)
    print(f"[scan] found {len(all_files)} video files; {len(seen)} already processed", flush=True)
    todo = [p for p in all_files if p not in seen]
    if args.limit:
        todo = todo[: args.limit]
    print(f"[scan] {len(todo)} to probe", flush=True)

    # Append-mode writers.
    n_ok = 0
    n_bad = 0
    t0 = time.time()
    with out_path.open("a") as out_f, bad_path.open("a") as bad_f, \
         ProcessPoolExecutor(max_workers=args.workers) as ex:
        futs = {ex.submit(worker, str(src_root), p): p for p in todo}
        for i, fut in enumerate(as_completed(futs), 1):
            rel, info = fut.result()
            if info is None:
                bad_f.write(rel + "\n")
                bad_f.flush()
                n_bad += 1
            else:
                info["src"] = args.src
                out_f.write(json.dumps(info) + "\n")
                out_f.flush()
                n_ok += 1
            if i % 200 == 0:
                rate = i / max(1.0, time.time() - t0)
                print(f"[probe] {i}/{len(todo)}  ok={n_ok}  bad={n_bad}  "
                      f"({rate:.1f}/s)", flush=True)
    print(f"DONE  ok={n_ok}  bad={n_bad}  out={out_path}", flush=True)


if __name__ == "__main__":
    main()