#!/usr/bin/env python3
"""Normalize mixed-source clips to a uniform 9:16 mezzanine for editing.

Target: 1080x1920, 30fps CFR, H.264 (high bitrate intra-friendly), yuv420p,
BT.709 SDR, 48kHz stereo AAC (silent track added if clip has no audio).

Handles per-clip (driven by clips.json from probe_clips.py):
  - landscape -> vertical center-cut (pan left/center/right honored)
  - HLG/PQ HDR -> SDR tonemap
  - D-Log -> Rec709 contrast/saturation lift (approximation; use official LUT
    via --dlog-lut for accurate conversion)
  - VFR -> CFR

Usage: normalize_clips.py work/clips.json --outdir work/mezz [--dlog-lut file.cube]
"""
import argparse, json, subprocess, sys
from pathlib import Path

W, H, FPS = 1080, 1920, 30

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


def build_vf(clip, dlog_lut=None):
    vf = []
    # 1. Color: get everything to BT.709 SDR first
    if clip["profile"] in ("hlg", "pq"):
        vf.append("zscale=t=linear:npl=100,tonemap=hable:desat=0,"
                  "zscale=p=bt709:t=bt709:m=bt709:r=tv")
    elif clip["profile"] == "dlog":
        if dlog_lut:
            vf.append(f"lut3d='{dlog_lut}'")
        else:
            # last-resort approximation: restore contrast + saturation
            vf.append("eq=contrast=1.35:saturation=1.4:gamma=0.92")
    # 2. Geometry: scale to cover 1080x1920, then crop with pan
    pan_x = PAN_X.get(clip.get("pan", "center"), PAN_X["center"])
    vf.append(f"scale={W}:{H}:force_original_aspect_ratio=increase")
    vf.append(f"crop={W}:{H}:{pan_x}:(ih-oh)/2")
    # 3. Uniform fps + pixel format
    vf.append(f"fps={FPS}")
    vf.append("format=yuv420p")
    return ",".join(vf)


def normalize(clip, outdir, dlog_lut=None):
    src = Path(clip["file"])
    dst = outdir / (src.stem + ".mp4")
    # write to .tmp.mp4, rename on success: a half-written mezzanine must never
    # be visible under its final name (downstream steps gate on these files)
    tmp = outdir / (src.stem + ".tmp.mp4")
    vf = build_vf(clip, dlog_lut)
    cmd = ["ffmpeg", "-y", "-i", str(src)]
    if not clip["has_audio"]:
        cmd += ["-f", "lavfi", "-i", "anullsrc=r=48000:cl=stereo", "-shortest"]
    cmd += ["-vf", vf,
            "-c:v", "libx264", "-preset", "medium", "-crf", "16",
            "-colorspace", "bt709", "-color_primaries", "bt709", "-color_trc", "bt709",
            "-c:a", "aac", "-b:a", "192k", "-ar", "48000", "-ac", "2",
            "-movflags", "+faststart", str(tmp)]
    r = subprocess.run(cmd, capture_output=True, text=True)
    if r.returncode != 0:
        tmp.unlink(missing_ok=True)
        print(f"FAILED {src.name}:\n{r.stderr[-800:]}", file=sys.stderr)
        return None
    tmp.rename(dst)
    return dst


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("clips_json")
    ap.add_argument("--outdir", default="work/mezz")
    ap.add_argument("--dlog-lut",
                    help="D-Log(M) to Rec709 .cube LUT — pass DJI's official "
                         "per-model LUT for exact conversion. Default: bundled "
                         "assets/luts/conversion/dji-dlogm-to-rec709.cube "
                         "(whitepaper approximation). Applied to dlog-profile "
                         "clips only — never to SDR footage")
    ap.add_argument("--profile", default="auto", help="(reserved) mezzanine profile")
    args = ap.parse_args()

    if not args.dlog_lut:
        bundled = (Path(__file__).resolve().parent.parent /
                   "assets/luts/conversion/dji-dlogm-to-rec709.cube")
        if bundled.exists():
            args.dlog_lut = str(bundled)

    clips = json.loads(Path(args.clips_json).read_text())["clips"]
    outdir = Path(args.outdir)
    outdir.mkdir(parents=True, exist_ok=True)

    done, failed = [], []
    for c in clips:
        if "error" in c:
            continue
        print(f"Normalizing {Path(c['file']).name} "
              f"({c['profile']}, {c['orientation']}, pan={c.get('pan','center')})...")
        out = normalize(c, outdir, args.dlog_lut)
        (done if out else failed).append(c["file"])

    print(f"\nDone: {len(done)} mezzanine files in {outdir}/")
    if failed:
        print(f"Failed: {failed}")
        sys.exit(1)


if __name__ == "__main__":
    main()