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

The plan's "music" key is optional. With it, the track is the spine (music-video
mode) and --natural-audio mixes clip sound underneath. Without it, the clips'
own audio is the spine (talking-head mode) — the VO carries the reel.

--vo-overlay lays a CONTINUOUS voiceover over the finished picture instead of
taking audio from the segments. Segments are concatenated, so a b-roll segment
otherwise replaces its slice of VO and punches a hole in the narration; with
--vo-overlay the picture cuts freely while the voice runs unbroken.

Grades:
  gritty  - teal/orange "Gritty Authentic Energy": cool shadows, warm mids/highs,
            lifted contrast, mild saturation, subtle vignette
  clean   - contrast + slight warmth only
  rapha   - muted painterly editorial: desaturated, faded blacks, warm mids,
            subtle grain, no vignette
  cine    - contrasty and grainy, colour-faithful: S-curve + grain with
            saturation held near original so hues do not shift
  punch   - harder contrast and grain than cine, slightly desaturated and
            warmer; the aggressive end of the set
  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]
  render_reel.py plan/cut_plan.json --resolution 4k --out work/master_4k.mp4
"""
import argparse, json, subprocess, sys, tempfile
from pathlib import Path

PAN_X = {"left": "0", "center": "(iw-ow)/2", "right": "iw-ow"}


def pan_expr(pan):
    """ffmpeg crop x-offset for a segment's pan.

    Accepts left/center/right or a 0..1 fraction, so a clip reused at several
    in-points can be reframed per cut. Requires a mezzanine built with
    --no-crop; a pre-cropped one has no pixels left to move to.
    """
    if pan is None:
        return PAN_X["center"]
    if isinstance(pan, (int, float)):
        f = min(1.0, max(0.0, float(pan)))
        return f"(iw-ow)*{f:.4f}"
    return PAN_X.get(str(pan).lower(), PAN_X["center"])


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",
    "rapha": (
        # muted painterly editorial: desaturated, warm mids/highs, cool-neutral
        # shadows, lifted blacks + softened highs (faded film), gentle contrast,
        # subtle grain, no vignette
        "eq=contrast=1.04:saturation=0.82:gamma=1.02,"
        "colorbalance=rs=0.02:gs=-0.01:bs=-0.03:rm=0.05:gm=0.00:bm=-0.05:"
        "rh=0.03:gh=0.00:bh=-0.04,"
        "curves=all='0/0.03 0.5/0.5 1/0.97',"
        "noise=alls=5:allf=t+u"
    ),
    "cine": (
        # rapha's tonal shape with more contrast and grain, but saturation
        # held near original (0.97) so the look does not shift colour
        "eq=contrast=1.14:saturation=0.97:gamma=1.01,"
        "colorbalance=rs=0.01:bs=-0.02:rm=0.03:bm=-0.03:rh=0.02:bh=-0.03,"
        "curves=all='0/0.02 0.25/0.21 0.75/0.79 1/0.98',"
        "noise=alls=8:allf=t+u"
    ),
    "punch": (
        # harder contrast and grain than cine, slightly desaturated and warmer
        "eq=contrast=1.20:saturation=0.94:gamma=0.99,"
        "colorbalance=rs=0.02:bs=-0.03:rm=0.06:bm=-0.05:rh=0.03:bh=-0.04,"
        "curves=all='0/0.0 0.25/0.19 0.75/0.82 1/1.0',"
        "noise=alls=10:allf=t+u"
    ),
    "none": None,
}

# Frame sizes by aspect then resolution. Must match the mezzanine built by
# normalize_clips.py, which takes the same --aspect / --resolution flags.
# 9:16 is the reel/short default; 16:9 is for landscape YouTube cuts.
ASPECTS = {
    "9:16": {"1080": (1080, 1920), "4k": (2160, 3840)},
    "16:9": {"1080": (1920, 1080), "4k": (3840, 2160)},
    "1:1":  {"1080": (1080, 1080), "4k": (2160, 2160)},
}
# Kept for importers that only ever needed the vertical sizes.
RESOLUTIONS = ASPECTS["9:16"]

PREVIEW = {"9:16": (540, 960), "16:9": (960, 540), "1:1": (540, 540)}


def dims_for(aspect, resolution, preview=False):
    return PREVIEW[aspect] if preview else ASPECTS[aspect][resolution]


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("--resolution", choices=("1080", "4k"), default="1080",
                    help="Output size (default 1080). Use the same value "
                         "normalize_clips.py built the mezzanine at.")
    ap.add_argument("--vo-overlay", metavar="FILE",
                    help="Lay continuous voiceover audio from FILE over the "
                         "cut picture, so b-roll does not interrupt the VO.")
    ap.add_argument("--vo-start", type=float, default=0.0, metavar="SEC",
                    help="Where in the VO file the reel's t=0 sits "
                         "(default 0).")
    ap.add_argument("--aspect", choices=sorted(ASPECTS), default="9:16",
                    help="Frame aspect (default 9:16 for reels/shorts). Use "
                         "16:9 for a landscape YouTube cut. Must match the "
                         "mezzanine.")
    ap.add_argument("--crf", type=int, default=None, metavar="N",
                    help="x264 quality, lower is better (default 16 master / "
                         "26 preview). 12 is near-lossless for an archival "
                         "master; below 10 mostly grows the file.")
    ap.add_argument("--audio-bitrate", default=None, metavar="RATE",
                    help="AAC bitrate (default 192k). 320k for an archival "
                         "master.")
    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. "
                         "Music-less plans use the clip audio at 0dB by default.")
    args = ap.parse_args()

    plan = json.loads(Path(args.plan).read_text())
    segs = plan["segments"]

    # Audio spine. Music-video plans carry a "music" track; talking-head plans
    # omit it and let the clips' own audio (the VO) carry the reel, so keep the
    # segment audio at unity unless the caller asked for a different gain.
    music = plan.get("music")
    nat_db = args.natural_audio
    if args.vo_overlay:
        # The VO file is the spine; segment audio would only fight it.
        if not Path(args.vo_overlay).exists():
            sys.exit(f"--vo-overlay file not found: {args.vo_overlay}")
        music = None
        nat_db = None
    elif not music and nat_db is None:
        nat_db = 0.0
    # Hand-authored plans omit total_duration; derive it from the segments.
    total = plan.get("total_duration") or sum(s["duration"] for s in segs)

    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 = []
        rw, rh = dims_for(args.aspect, args.resolution, args.preview)
        crf = str(args.crf) if args.crf else ("26" if args.preview else "16")
        abr = args.audio_bitrate or "192k"
        for n, s in enumerate(segs):
            p = td / f"seg{n:03d}.mp4"
            # scale-to-cover then crop, so a full-frame mezzanine is reframed
            # rather than squashed; on an already-cropped one the crop is a
            # no-op and the pan is simply ignored.
            geom = (f"scale={rw}:{rh}:force_original_aspect_ratio=increase,"
                    f"crop={rw}:{rh}:{pan_expr(s.get('pan'))}:(ih-oh)/2")
            cut = ["ffmpeg", "-y", "-ss", str(s["in"]), "-i", s["clip"],
                   "-t", str(s["duration"]),
                   "-vf", f"{geom},fps=30,format=yuv420p",
                   "-c:v", "libx264", "-preset", "veryfast", "-crf", crf]
            if nat_db 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", abr, "-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. Lay the audio spine
        cmd = ["ffmpeg", "-y", "-i", str(video)]
        if args.vo_overlay:
            cmd += ["-ss", str(args.vo_start), "-i", str(args.vo_overlay)]
        if music:
            cmd += ["-ss", str(plan.get("music_start", 0)), "-i", music]
        fade = f"afade=t=out:st={max(0, total-1.2)}:d=1.2"
        enc = ["-c:v", "libx264", "-preset",
               "veryfast" if args.preview else "slow", "-crf", crf]
        if music and nat_db 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={nat_db}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"]
        elif music:
            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"]
        elif args.vo_overlay:
            # picture from the segments, voice from the VO file, unbroken
            cmd += (["-vf", ",".join(vf)] + enc) if vf else ["-c:v", "copy"]
            cmd += ["-filter:a", fade, "-map", "0:v:0", "-map", "1:a:0"]
        else:
            # talking-head: the segments' own audio is the only spine
            cmd += (["-vf", ",".join(vf)] + enc) if vf else ["-c:v", "copy"]
            cmd += ["-filter:a", f"volume={nat_db}dB,{fade}",
                    "-map", "0:v:0", "-map", "0:a:0"]
        cmd += ["-c:a", "aac", "-b:a", abr, "-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)
    rw, rh = dims_for(args.aspect, args.resolution, args.preview)
    print(f"Rendered {out} ({float(r.stdout.strip()):.1f}s, {rw}x{rh} "
          f"{args.aspect}, {'preview' if args.preview else 'master'}, "
          f"grade={args.lut or args.grade})")


if __name__ == "__main__":
    main()