#!/usr/bin/env python3
"""Same frames, every grade, side by side — pick a look without guessing.

Choosing a grade by rendering the reel once per look is slow and, worse, you
end up comparing from memory. This applies each grade to the SAME still frames
and tiles them: one row per frame, one column per grade, labelled.

Frames come from a cut plan (one per segment, spread across the reel) or from
any video at even intervals. Stills only — no encode, so it takes seconds.

Requires ImageMagick for the labels; without it the sheet is unlabelled and the
column order is printed instead.

Usage:
  compare_grades.py plan/cut_plan.json --out work/grades.png
  compare_grades.py --video work/master.mp4 --grades rapha,cine,punch --frames 4
"""
import argparse, json, math, os, shutil, subprocess, sys, tempfile
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from render_reel import GRADES              # noqa: E402

TILE_W = 300
LABEL_H = 30
DEFAULT_FRAMES = 3


def default_workers():
    return max(2, min(8, (os.cpu_count() or 4) - 2))


def graded_still(src, t, grade, dst, width=TILE_W):
    """One frame with one grade applied. Stills only — no encode."""
    vf = [f"scale={width}:-2"]
    chain = GRADES.get(grade)
    if chain:
        vf.append(chain)
    # Pin the pixel format. Without a grade chain ffmpeg writes rgb48be while
    # graded stills come out rgb24, and the tile filter silently drops the odd
    # one out — the ungraded column vanishes and every label shifts by one.
    vf.append("format=rgb24")
    r = subprocess.run(
        ["ffmpeg", "-y", "-v", "error", "-skip_frame", "nokey",
         "-noaccurate_seek", "-ss", f"{t:.3f}", "-i", str(src),
         "-frames:v", "1", "-vf", ",".join(vf), str(dst)],
        capture_output=True, text=True)
    return dst if r.returncode == 0 and dst.exists() else None


def png_pix_fmt(path):
    r = subprocess.run(["ffprobe", "-v", "error", "-show_entries",
                        "stream=pix_fmt", "-of", "csv=p=0", str(path)],
                       capture_output=True, text=True)
    return r.stdout.strip()


def label(magick, img, text=""):
    """Splice a strip under every tile; only some carry text.

    The strip goes on EVERY tile even when blank: ffmpeg's tile filter needs
    uniform input sizes, so labelling just the top row makes those tiles taller
    and the remaining rows get silently dropped.

    -splice follows -gravity, so south or the strip lands mid-frame.
    """
    cmd = [magick, str(img), "-background", "#111318", "-gravity", "south",
           "-splice", f"0x{LABEL_H}"]
    if text:
        cmd += ["-fill", "white", "-pointsize", "17", "-annotate", "+0+6", text]
    # Flatten and drop alpha. Drawing text makes ImageMagick write rgba while
    # a blank splice stays rgb24, and ffmpeg's glob demuxer drops whichever
    # format is in the minority — silently, with exit code 0. The labelled row
    # simply vanishes from the sheet.
    cmd += ["-background", "#111318", "-alpha", "remove", "-alpha", "off",
            "-define", "png:color-type=2"]
    subprocess.run(cmd + [str(img)], capture_output=True, text=True)


def sample_points(plan_path, video, n_frames):
    """[(source, time)] to compare — spread across the reel, not clustered."""
    if video:
        r = subprocess.run(["ffprobe", "-v", "error", "-show_entries",
                            "format=duration", "-of", "csv=p=0", video],
                           capture_output=True, text=True)
        dur = float(r.stdout.strip() or 0)
        step = dur / (n_frames + 1)
        return [(video, round(step * (i + 1), 2)) for i in range(n_frames)]

    segs = json.loads(Path(plan_path).read_text()).get("segments", [])
    if not segs:
        sys.exit(f"{plan_path} has no segments.")
    # spread the picks across the plan rather than taking the first N
    step = max(1, len(segs) // n_frames)
    picked = segs[::step][:n_frames]
    return [(s["clip"], s["in"] + s["duration"] / 2) for s in picked]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("plan", nargs="?")
    ap.add_argument("--video", help="Compare on a rendered video instead")
    ap.add_argument("--out", default="work/grade_compare.png")
    ap.add_argument("--grades", default=None,
                    help="Comma-separated (default: every grade except none)")
    ap.add_argument("--frames", type=int, default=DEFAULT_FRAMES)
    ap.add_argument("--workers", type=int, default=default_workers())
    args = ap.parse_args()

    if not args.plan and not args.video:
        sys.exit("Pass a cut plan, or --video.")

    if args.grades:
        grades = [g.strip() for g in args.grades.split(",") if g.strip()]
        unknown = [g for g in grades if g not in GRADES]
        if unknown:
            sys.exit(f"Unknown grade(s): {', '.join(unknown)}. "
                     f"Available: {', '.join(sorted(GRADES))}")
    else:
        grades = ["none"] + [g for g in GRADES if g != "none"]

    points = sample_points(args.plan, args.video, args.frames)
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    magick = shutil.which("magick")
    work = Path(tempfile.mkdtemp(prefix="grades_", dir=str(out.parent)))

    try:
        jobs = [(r, c, src, t, g)
                for r, (src, t) in enumerate(points)
                for c, g in enumerate(grades)]

        def one(job):
            r, c, src, t, g = job
            dst = work / f"r{r:02d}c{c:02d}.png"
            got = graded_still(src, t, g, dst)
            if got and magick:
                # text on the top row, blank strip elsewhere — every tile must
                # end up the same height or tile() drops the odd rows
                label(magick, dst, g if r == 0 else "")
            return got

        with ThreadPoolExecutor(max_workers=args.workers) as pool:
            made = list(pool.map(one, jobs))
        missing = [j for j, got in zip(jobs, made) if not got]
        if not any(made):
            sys.exit("No frames could be read — check the clip paths.")
        if missing:
            # A dropped tile shifts every later column under the wrong label,
            # so refuse rather than hand back a mislabelled sheet.
            sys.exit(f"{len(missing)} tile(s) failed to render "
                     f"({', '.join(sorted({j[4] for j in missing}))}). "
                     f"The sheet would be mislabelled, so nothing was written.")

        # Guard the invariant rather than trusting it: a format split here
        # costs a whole row with no error anywhere.
        fmts = {png_pix_fmt(work / f"r{r:02d}c{c:02d}.png")
                for r, c, *_ in jobs}
        fmts.discard("")
        if len(fmts) > 1:
            sys.exit(f"Tiles came out in mixed pixel formats ({', '.join(sorted(fmts))}). "
                     f"ffmpeg's tile filter would drop a row silently, so "
                     f"nothing was written.")

        tmp = out.with_name(out.stem + ".tmp" + out.suffix)
        r = subprocess.run(
            ["ffmpeg", "-y", "-v", "error", "-framerate", "1",
             "-pattern_type", "glob", "-i", str(work / "r*.png"),
             "-vf", f"tile={len(grades)}x{len(points)}:padding=4:color=#111318",
             "-frames:v", "1", str(tmp)], capture_output=True, text=True)
        if r.returncode != 0:
            tmp.unlink(missing_ok=True)
            sys.exit(f"tile failed:\n{r.stderr[-600:]}")
        tmp.rename(out)
    finally:
        shutil.rmtree(work, ignore_errors=True)

    print(f"Grade comparison -> {out}")
    print(f"  {len(points)} frame(s) x {len(grades)} grade(s), "
          f"columns left to right: {', '.join(grades)}")
    if not magick:
        print("  (unlabelled — ImageMagick not installed)")
    print("\nPick from the sheet, then render with --grade <name>. Judge on a "
          "frame with skin and sky in it if you have one; grades diverge most "
          "there.")


if __name__ == "__main__":
    main()