#!/usr/bin/env python3
"""Render a reel from cut_plan.json: trim segments, concat, lay music, grade.

Grades:
  gritty  - teal/orange "Gritty Authentic Energy": cool shadows, warm mids/highs,
            lifted contrast, mild saturation, subtle vignette
  clean   - contrast + slight warmth only
  none    - passthrough
  (or --lut file.cube for a custom 3D LUT)

Usage:
  render_reel.py plan/cut_plan.json --preview --out preview/draft_v1.mp4
  render_reel.py plan/cut_plan.json --grade gritty --out work/master.mp4 [--srt caps.srt]
"""
import argparse, json, subprocess, sys, tempfile
from pathlib import Path

GRADES = {
    "gritty": (
        "eq=contrast=1.08:saturation=1.12,"
        "colorbalance=rs=-0.06:gs=-0.02:bs=0.08:rm=0.05:gm=0.01:bm=-0.05:"
        "rh=0.04:bh=-0.04,"
        "curves=all='0/0.02 0.5/0.5 1/0.98',"
        "vignette=PI/5:mode=backward"
    ),
    "clean": "eq=contrast=1.05:saturation=1.05,colorbalance=rm=0.03:bm=-0.03",
    "none": None,
}


def run(cmd, label):
    r = subprocess.run(cmd, capture_output=True, text=True)
    if r.returncode != 0:
        sys.exit(f"{label} failed:\n{r.stderr[-1200:]}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("plan")
    ap.add_argument("--out", required=True)
    ap.add_argument("--preview", action="store_true", help="fast 540x960 draft")
    ap.add_argument("--grade", default="none", choices=list(GRADES))
    ap.add_argument("--lut", help="Custom .cube LUT (overrides --grade)")
    ap.add_argument("--srt", help="Burn in captions from SRT file")
    ap.add_argument("--music-gain-db", type=float, default=-3.0)
    ap.add_argument("--natural-audio", type=float, metavar="DB", default=None,
                    help="Keep clips' natural audio mixed under music at this gain "
                         "(e.g. -18). 30ms fades at every cut prevent pops.")
    args = ap.parse_args()

    plan = json.loads(Path(args.plan).read_text())
    segs = plan["segments"]
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    # render to .tmp, rename after validation — never expose a partial file
    out_tmp = out.with_name(out.stem + ".tmp" + out.suffix)

    with tempfile.TemporaryDirectory() as td:
        td = Path(td)
        # 1. Cut each segment losslessly-ish (re-encode for frame accuracy)
        parts = []
        scale = "540:960" if args.preview else "1080:1920"
        crf = "26" if args.preview else "16"
        for n, s in enumerate(segs):
            p = td / f"seg{n:03d}.mp4"
            cut = ["ffmpeg", "-y", "-ss", str(s["in"]), "-i", s["clip"],
                   "-t", str(s["duration"]),
                   "-vf", f"scale={scale},fps=30,format=yuv420p",
                   "-c:v", "libx264", "-preset", "veryfast", "-crf", crf]
            if args.natural_audio is not None:
                fade_out = max(0.0, s["duration"] - 0.03)
                cut += ["-af", f"afade=t=in:d=0.03,afade=t=out:st={fade_out}:d=0.03",
                        "-c:a", "aac", "-b:a", "192k", "-ar", "48000"]
            else:
                cut += ["-an"]
            run(cut + [str(p)], f"segment {n}")
            parts.append(p)
            print(f"  cut {n+1}/{len(segs)}", flush=True)

        # 2. Concat
        lst = td / "list.txt"
        lst.write_text("".join(f"file '{p}'\n" for p in parts))
        video = td / "video.mp4"
        run(["ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", str(lst),
             "-c", "copy", str(video)], "concat")

        # 3. Grade + captions
        vf = []
        if args.lut:
            vf.append(f"lut3d='{args.lut}'")
        elif GRADES.get(args.grade):
            vf.append(GRADES[args.grade])
        if args.srt:
            vf.append(f"subtitles='{args.srt}':force_style="
                      "'FontName=Arial Black,FontSize=14,PrimaryColour=&HFFFFFF,"
                      "OutlineColour=&H000000,Outline=2,Alignment=2,MarginV=60'")

        # 4. Music
        cmd = ["ffmpeg", "-y", "-i", str(video),
               "-ss", str(plan.get("music_start", 0)), "-i", plan["music"]]
        fade = f"afade=t=out:st={max(0, plan['total_duration']-1.2)}:d=1.2"
        enc = ["-c:v", "libx264", "-preset",
               "veryfast" if args.preview else "slow", "-crf", crf]
        if args.natural_audio is not None:
            # single filter_complex graph for both video grade and audio mix
            vchain = f"[0:v]{','.join(vf)}[vout];" if vf else ""
            vmap = "[vout]" if vf else "0:v:0"
            cmd += ["-filter_complex",
                    vchain +
                    f"[0:a]volume={args.natural_audio}dB[nat];"
                    f"[1:a]volume={args.music_gain_db}dB[mus];"
                    f"[nat][mus]amix=inputs=2:duration=first:normalize=0,{fade}[aout]",
                    "-map", vmap, "-map", "[aout]"]
            cmd += enc if vf else ["-c:v", "copy"]
        else:
            cmd += (["-vf", ",".join(vf)] + enc) if vf else ["-c:v", "copy"]
            cmd += ["-filter:a", f"volume={args.music_gain_db}dB,{fade}",
                    "-map", "0:v:0", "-map", "1:a:0"]
        cmd += ["-c:a", "aac", "-b:a", "192k", "-shortest",
                "-movflags", "+faststart", str(out_tmp)]
        run(cmd, "final mux")

    # validate, then publish under the final name
    r = subprocess.run(["ffprobe", "-v", "quiet", "-show_entries", "format=duration",
                        "-of", "csv=p=0", str(out_tmp)], capture_output=True, text=True)
    if r.returncode != 0 or not r.stdout.strip():
        out_tmp.unlink(missing_ok=True)
        sys.exit("Rendered file failed ffprobe validation — output discarded")
    out_tmp.rename(out)
    print(f"Rendered {out} ({float(r.stdout.strip()):.1f}s, "
          f"{'preview' if args.preview else 'master'}, grade={args.lut or args.grade})")


if __name__ == "__main__":
    main()