#!/usr/bin/env python3
"""Normalize mega_curated videos: 24 fps CFR, square pixels (SAR 1:1), copy captions.

Non-destructive: writes derived files to OUT_DIR, sources untouched.
Skips clips whose 24fps-converted frame count would fall below 22 (17n+5 grid minimum for video).
Usage: normalize_videos.py SRC_DIR OUT_DIR [--jobs N]
"""
import argparse
import fractions
import json
import shutil
import subprocess
import sys
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path

MIN_FRAMES = 22  # smallest useful 17n+5 video length


def probe(path: Path):
    out = subprocess.run(
        ["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries",
         "stream=width,height,sample_aspect_ratio,avg_frame_rate,nb_frames,duration",
         "-of", "json", str(path)],
        capture_output=True, text=True, timeout=120)
    if out.returncode != 0:
        return None
    st = json.loads(out.stdout)["streams"][0]
    w, h = int(st["width"]), int(st["height"])
    sar = st.get("sample_aspect_ratio") or "1:1"
    if sar in ("0:1", "N/A"):
        sar = "1:1"
    fps = float(fractions.Fraction(st["avg_frame_rate"]))
    nb = st.get("nb_frames")
    if nb is None or nb == "N/A":
        dur = float(st.get("duration") or 0)
        nb = int(dur * fps)
    return w, h, sar, fps, int(nb)


def process(src: Path, dst_dir: Path):
    info = probe(src)
    if info is None:
        return ("probe_fail", src.name, None)
    w, h, sar, fps, nb = info
    out_frames = int(nb * 24.0 / fps) if fps > 0 else 0
    if out_frames < MIN_FRAMES:
        return ("too_short", src.name, out_frames)
    sn, sd = (int(x) for x in sar.split(":"))
    vf = ["fps=24"]
    if sn != sd:  # anamorphic: resample to square pixels, keep height
        out_w = round(w * sn / sd / 2) * 2
        vf.insert(0, f"scale={out_w}:{h}:flags=lanczos")
        final_w = out_w
    else:
        final_w = w
    vf.append("setsar=1")
    dst = dst_dir / src.name
    if dst.exists():
        return ("exists", src.name, (final_w, h))
    r = subprocess.run(
        ["ffmpeg", "-y", "-v", "error", "-i", str(src),
         "-vf", ",".join(vf), "-c:v", "libx264", "-crf", "15",
         "-preset", "fast", "-pix_fmt", "yuv420p", "-an", "-map_metadata", "-1",
         str(dst) + ".tmp.mp4"],
        capture_output=True, text=True, timeout=600)
    if r.returncode != 0:
        Path(str(dst) + ".tmp.mp4").unlink(missing_ok=True)
        return ("encode_fail", src.name, r.stderr[-200:])
    Path(str(dst) + ".tmp.mp4").rename(dst)
    txt = src.with_suffix(".txt")
    if txt.exists():
        shutil.copy2(txt, dst_dir / txt.name)
    return ("ok", src.name, (final_w, h))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("src")
    ap.add_argument("out")
    ap.add_argument("--jobs", type=int, default=30)
    args = ap.parse_args()
    src_dir, out_dir = Path(args.src), Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)
    videos = sorted(src_dir.glob("*.mp4"))
    print(f"{len(videos)} source videos", flush=True)
    results = {"ok": 0, "exists": 0, "too_short": 0, "probe_fail": 0, "encode_fail": 0}
    failures = []
    with ProcessPoolExecutor(max_workers=args.jobs) as ex:
        futs = {ex.submit(process, v, out_dir): v for v in videos}
        for i, fut in enumerate(as_completed(futs)):
            status, name, extra = fut.result()
            results[status] += 1
            if status in ("too_short", "probe_fail", "encode_fail"):
                failures.append((status, name, extra))
            if (i + 1) % 500 == 0:
                print(f"{i+1}/{len(videos)} {results}", flush=True)
    print("FINAL:", results, flush=True)
    with open(out_dir.parent / "normalize_failures.log", "w") as f:
        for status, name, extra in failures:
            f.write(f"{status}\t{name}\t{extra}\n")


if __name__ == "__main__":
    main()