#!/usr/bin/env python3
"""Dense contact sheets — LOOK at the footage before planning anything.

One PNG grid per clip, frames sampled evenly across the FULL duration.
This is how you catch what probe metadata can't: talking heads, burned-in
captions/lyrics (often intermittent — mid-clip only), screen recordings,
misnamed files, shaky/unusable sections, and where the actual action lives.

Frames are laid out row-major. No timestamps are burned in (Homebrew ffmpeg
lacks drawtext), so use the printed "frame N ≈ N x interval" mapping — the
script prints the sampling interval per clip.

Two samplers, chosen by clip length (--sampler forces one):

  decode  one ffmpeg pass with an fps= filter. Decodes the whole clip, but
          spawns a single process. Cheapest for short clips.
  seek    one keyframe seek per frame (-ss BEFORE -i). Skips decoding, but
          spawns a process per frame. Cheapest for long clips.

Which is cheaper depends on how expensive the codec is to decode, and that
varies enormously: a flat synthetic clip decodes ~30x faster than real camera
4K HEVC at the same resolution and length. A fixed rule tuned on one of those
is wrong for the other, so `auto` MEASURES the clip instead of guessing:

  - time one keyframe seek
  - time decoding a short slice, and extrapolate over the full duration
  - take whichever projected total is smaller

The probe costs about a second per clip. On footage where a full decode runs
into minutes that is a rounding error, and it means the decision adapts to
whatever camera and codec the shoot actually used rather than to whatever was
convenient to benchmark.

--sampler seek|decode forces one and skips the probe.

Clips are also processed concurrently — previously serial.

Usage:
  contact_sheet.py work/clips.json --outdir work/sheets
  contact_sheet.py work/clips.json --outdir work/sheets --interval 1.0  # denser
"""
import argparse, json, math, os, shutil, subprocess, sys, tempfile, time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

COLS = 6
MIN_FRAMES, MAX_FRAMES = 12, 60
THUMB_W = 320
PROBE_SLICE_S = 2.0       # how much footage the decode probe times


def default_workers():
    """Leave headroom — ffmpeg is already multi-threaded per invocation."""
    return max(2, min(8, (os.cpu_count() or 4) - 2))


def cheaper_sampler(seek_one_s, decode_slice_s, duration, n_frames,
                    slice_s=PROBE_SLICE_S):
    """Project both totals from measured costs and pick the smaller.

    decode cost scales with duration; seek cost scales with frame count.
    Returns (choice, projected_decode_s, projected_seek_s).
    """
    projected_decode = decode_slice_s * (duration / slice_s) if slice_s else 0.0
    projected_seek = seek_one_s * n_frames
    return (("seek" if projected_seek < projected_decode else "decode"),
            projected_decode, projected_seek)


def probe_costs(src, duration, keyframe_only=False):
    """Measure one seek and one short decode on THIS clip."""
    with tempfile.TemporaryDirectory() as td:
        td = Path(td)
        mid = duration / 2
        t0 = time.perf_counter()
        grab_frame(src, mid, td / "probe.png", keyframe_only=keyframe_only)
        seek_one = time.perf_counter() - t0

        t0 = time.perf_counter()
        subprocess.run(
            ["ffmpeg", "-y", "-v", "error", "-i", str(src),
             "-t", f"{PROBE_SLICE_S}", "-f", "null", "-"],
            capture_output=True, text=True)
        decode_slice = time.perf_counter() - t0
    return seek_one, decode_slice


def grab_frame(src, t, dst, width=THUMB_W, keyframe_only=False):
    """One frame at time t. -ss before -i seeks instead of decoding to t.

    keyframe_only adds -skip_frame nokey, which decodes ONLY keyframes instead
    of decoding forward from the nearest one — 7x faster on real 4K HEVC. The
    frame returned is then the nearest keyframe rather than the exact time,
    which is fine for a survey sheet but only while keyframes are closer
    together than the sampling step (see keyframe_interval).
    """
    pre = ["-skip_frame", "nokey", "-noaccurate_seek"] if keyframe_only else []
    r = subprocess.run(
        ["ffmpeg", "-y", "-v", "error"] + pre + ["-ss", f"{t:.3f}", "-i", str(src),
         "-frames:v", "1", "-vf", f"scale={width}:-2", "-fps_mode", "passthrough",
         str(dst)], capture_output=True, text=True)
    return dst if r.returncode == 0 and dst.exists() else None


def keyframe_interval(src, sample=12):
    """Median gap between the first few keyframes, or None if undetectable."""
    r = subprocess.run(
        ["ffprobe", "-v", "error", "-select_streams", "v:0", "-skip_frame",
         "nokey", "-show_entries", "frame=pts_time", "-of", "csv=p=0",
         "-read_intervals", "%+60", str(src)],
        capture_output=True, text=True)
    ts = []
    for line in r.stdout.splitlines():
        line = line.strip().rstrip(",")
        try:
            ts.append(float(line))
        except ValueError:
            continue
    ts = sorted(ts)[:sample]
    if len(ts) < 3:
        return None
    gaps = sorted(b - a for a, b in zip(ts, ts[1:]) if b > a)
    return gaps[len(gaps) // 2] if gaps else None


def sheet_for(clip, outdir, interval, sampler="auto"):
    src = Path(clip["file"])
    if clip.get("type") == "image":     # a still IS its own contact sheet
        dst = outdir / f"{src.stem}_sheet.png"
        tmp = outdir / f"{src.stem}_sheet.tmp.png"
        r = subprocess.run(["ffmpeg", "-y", "-v", "error", "-i", str(src),
                            "-vf", "scale=640:-2", "-frames:v", "1", str(tmp)],
                           capture_output=True, text=True)
        if r.returncode != 0 or not tmp.exists():
            tmp.unlink(missing_ok=True)
            print(f"  FAILED {src.name}: {r.stderr.strip()[-300:]}", file=sys.stderr)
            return None
        tmp.rename(dst)
        print(f"  {dst.name}: still image (becomes a Ken Burns clip)")
        return dst
    dur = clip["duration"]
    n = max(MIN_FRAMES, min(MAX_FRAMES, math.ceil(dur / interval)))
    step = dur / n
    rows = math.ceil(n / COLS)
    dst = outdir / f"{src.stem}_sheet.png"
    tmp = outdir / f"{src.stem}_sheet.tmp.png"

    # Keyframe-only decoding is the big win on long-GOP 4K, but only while
    # keyframes are denser than the sampling step — otherwise adjacent samples
    # snap to the same keyframe and the sheet shows duplicates.
    kf = keyframe_interval(src)
    kf_ok = kf is not None and kf < step * 0.9

    if sampler == "auto":
        seek_one, decode_slice = probe_costs(src, dur, kf_ok)
        sampler, est_d, est_s = cheaper_sampler(seek_one, decode_slice, dur, n)
        note = (f" [{sampler}{'+keyframe' if kf_ok and sampler == 'seek' else ''}; "
                f"decode~{est_d:.0f}s vs seek~{est_s:.0f}s]")
    else:
        note = f" [{sampler}{'+keyframe' if kf_ok and sampler == 'seek' else ''}]"

    if sampler == "decode":
        vf = f"fps=1/{step:.4f},scale={THUMB_W}:-2,tile={COLS}x{rows}:padding=2"
        r = subprocess.run(
            ["ffmpeg", "-y", "-v", "error", "-i", str(src),
             "-vf", vf, "-frames:v", "1", str(tmp)],
            capture_output=True, text=True)
        if r.returncode != 0 or not tmp.exists():
            tmp.unlink(missing_ok=True)
            print(f"  FAILED {src.name}: {r.stderr.strip()[-300:]}",
                  file=sys.stderr)
            return None
        tmp.rename(dst)
        print(f"  {dst.name}: {n} frames, 1 every {step:.1f}s{note} "
              f"(frame N is at ~N x {step:.1f}s, row-major)")
        return dst

    # Sample at the MIDPOINT of each slot: t=0 on many clips is a black or
    # fade-in frame, which reads as an empty shot on the sheet.
    times = [min(dur - 0.05, (i + 0.5) * step) for i in range(n)]
    work = Path(tempfile.mkdtemp(prefix="sheet_", dir=str(outdir)))
    try:
        frames = [grab_frame(src, t, work / f"f{i:03d}.png",
                             keyframe_only=kf_ok)
                  for i, t in enumerate(times)]
        got = [f for f in frames if f]
        if not got:
            print(f"  FAILED {src.name}: no frames could be read",
                  file=sys.stderr)
            return None
        r = subprocess.run(
            ["ffmpeg", "-y", "-v", "error", "-framerate", "1",
             "-pattern_type", "glob", "-i", str(work / "*.png"),
             "-vf", f"tile={COLS}x{rows}:padding=2", "-frames:v", "1",
             str(tmp)], capture_output=True, text=True)
        if r.returncode != 0 or not tmp.exists():
            tmp.unlink(missing_ok=True)
            print(f"  FAILED {src.name}: {r.stderr.strip()[-300:]}",
                  file=sys.stderr)
            return None
        n = len(got)
    finally:
        shutil.rmtree(work, ignore_errors=True)
    tmp.rename(dst)
    print(f"  {dst.name}: {n} frames, 1 every {step:.1f}s{note} "
          f"(frame N is at ~N x {step:.1f}s, row-major)")
    return dst


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("clips_json")
    ap.add_argument("--outdir", default="work/sheets")
    ap.add_argument("--sampler", choices=("auto", "seek", "decode"),
                    default="auto",
                    help="auto (default) times a seek and a short decode on "
                         "each clip and picks the cheaper; seek/decode force "
                         "one and skip the probe.")
    ap.add_argument("--workers", type=int, default=default_workers(),
                    help="Clips processed concurrently (default: cores-2, "
                         "capped at 8). Use 1 to serialize for debugging.")
    ap.add_argument("--interval", type=float, default=2.0,
                    help="Target seconds between sampled frames (default 2.0; "
                         "capped at %d frames for very long clips — outliers "
                         "get coarser sampling, extract a segment first)" % MAX_FRAMES)
    args = ap.parse_args()

    clips = [c for c in json.loads(Path(args.clips_json).read_text())["clips"]
             if "error" not in c]
    outdir = Path(args.outdir)
    outdir.mkdir(parents=True, exist_ok=True)

    print(f"Contact sheets -> {outdir}/  ({args.workers} workers)")
    if args.workers > 1:
        with ThreadPoolExecutor(max_workers=args.workers) as pool:
            made = list(pool.map(
                lambda c: sheet_for(c, outdir, args.interval, args.sampler),
                clips))
    else:
        made = [sheet_for(c, outdir, args.interval, args.sampler)
                for c in clips]
    ok = [m for m in made if m]
    print(f"\n{len(ok)}/{len(clips)} sheets written. VIEW EVERY SHEET before "
          f"building a cut plan: reject clips with burned-in text, talking "
          f"heads, or screen content; note where the best action moments are "
          f"(for in-points, hook, and closer).")
    if len(ok) < len(clips):
        sys.exit(1)


if __name__ == "__main__":
    main()