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

Target: 1080x1920 by default (--resolution 4k for a 2160x3840 archival
master), 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

Still images (type=image in clips.json) become Ken Burns motion clips:
a slow push/pull/pan rendered at mezzanine spec, so downstream steps treat
them like any other footage. Direction via the clip's "kb" field:
  auto (default)  landscape stills pan right, everything else pushes in
  in | out        slow zoom (10% over the clip — tasteful, not a crash zoom)
  left | right    lateral drift across a landscape still

Usage: normalize_clips.py work/clips.json --outdir work/mezz [--dlog-lut file.cube]
       normalize_clips.py work/clips.json --resolution 4k
"""
import argparse, json, os, subprocess, sys
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path


def default_workers():
    """Each ffmpeg is already multi-threaded, so oversubscribing thrashes.

    A couple in flight still helps: one encode rarely saturates a modern
    many-core machine, and the reads overlap.
    """
    return max(1, min(4, (os.cpu_count() or 4) // 4))

W, H, FPS = 1080, 1920, 30
KB_ZOOM = 0.10   # total Ken Burns push over the clip

# Mezzanine sizes by aspect then resolution. 9:16 is the reel/short default;
# 16:9 produces a landscape master for YouTube proper. 4k is an archival
# master — platforms downscale on delivery either way.
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)},
}
RESOLUTIONS = ASPECTS["9:16"]

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


def build_vf(clip, dlog_lut=None, dims=(W, H), crop=True):
    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 the target, then crop with pan
    w, h = dims
    if crop:
        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")
    else:
        # Full frame at target height: leaves horizontal pixels for a
        # per-segment pan at render time. Nothing else can reframe later —
        # once a mezzanine is cropped, the pixels are gone.
        vf.append(f"scale=-2:{h}")
    # 3. Uniform fps + pixel format
    vf.append(f"fps={FPS}")
    vf.append("format=yuv420p")
    return ",".join(vf)


def build_kb_vf(clip, dims=(W, H)):
    """Ken Burns filter chain for a still image."""
    w, h = dims
    dur = clip["duration"]
    frames = round(dur * FPS)
    kb = clip.get("kb", "auto")
    if kb == "auto":
        kb = "right" if clip["orientation"] == "landscape" else "in"
    if kb in ("in", "out"):
        # pre-scale to 2x target so zoompan steps land on half-pixels (smooth)
        pan_x = PAN_X.get(clip.get("pan", "center"), PAN_X["center"])
        z = (f"1+{KB_ZOOM}*on/{frames}" if kb == "in"
             else f"{1 + KB_ZOOM}-{KB_ZOOM}*on/{frames}")
        vf = (f"scale={2*w}:{2*h}:force_original_aspect_ratio=increase,"
              f"crop={2*w}:{2*h}:{pan_x}:(ih-oh)/2,"
              f"zoompan=z='{z}':x='iw/2-iw/zoom/2':y='ih/2-ih/zoom/2'"
              f":d=1:s={w}x{h}:fps={FPS}")
    else:  # left / right lateral drift
        x = (f"(iw-ow)*t/{dur}" if kb == "right"
             else f"(iw-ow)*(1-t/{dur})")
        vf = f"scale=-2:{h},crop={w}:{h}:'{x}':0"
    return vf + f",fps={FPS},format=yuv420p", kb


def normalize_image(clip, outdir, dims=(W, H)):
    src = Path(clip["file"])
    dst = outdir / (src.stem + ".mp4")
    tmp = outdir / (src.stem + ".tmp.mp4")
    vf, kb = build_kb_vf(clip, dims)
    cmd = ["ffmpeg", "-y", "-loop", "1", "-framerate", str(FPS),
           "-t", str(clip["duration"]), "-i", str(src),
           "-f", "lavfi", "-i", "anullsrc=r=48000:cl=stereo",
           "-vf", vf, "-t", str(clip["duration"]), "-shortest",
           "-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 normalize(clip, outdir, dlog_lut=None, dims=(W, H), crop=True):
    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, dims, crop)
    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 thin_sources(clips, dims):
    """Sources whose crop to the target aspect lands below the target width.

    Cropping across aspect throws away pixels on one axis: a landscape frame
    cut to 9:16 yields only height*9/16 of width, a portrait frame cut to 16:9
    keeps only its own width. Either way the result upscales to fill the
    target: bigger, not sharper. Returns [(name, native_width), ...].
    """
    w, h = dims
    target_ar = w / h
    thin = []
    for c in clips:
        if "error" in c:
            continue
        sw, sh = c.get("width", 0), c.get("height", 0)
        if not sw or not sh:
            continue
        src_ar = sw / sh
        if src_ar > target_ar:      # source wider — crop sides, width is the limit
            native_w = round(sh * target_ar)
        else:                        # source taller — crop top/bottom
            native_w = sw
        if native_w < w:
            thin.append((Path(c["file"]).name, native_w))
    return thin


def warn_upscale(clips, dims, aspect):
    w, h = dims
    thin = thin_sources(clips, dims)
    if not thin:
        return thin
    print(f"NOTE: {len(thin)} clip(s) crop to less than {w}px wide "
          f"and will be upscaled to fill {w}x{h} — bigger, not sharper:",
          file=sys.stderr)
    for name, native in thin:
        print(f"  {name}: ~{native}px native after {aspect} crop", file=sys.stderr)
    return thin


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")
    ap.add_argument("--resolution", choices=("1080", "4k"), default="1080",
                    help="Mezzanine size (default 1080). 4k = archival master.")
    ap.add_argument("--aspect", choices=sorted(ASPECTS), default="9:16",
                    help="Frame aspect (default 9:16 for reels/shorts). 16:9 "
                         "for a landscape YouTube cut.")
    ap.add_argument("--no-crop", action="store_true",
                    help="Keep the full frame at target height instead of "
                         "cropping to aspect. Needed if you want per-segment "
                         "pan at render time — a cropped mezzanine has no "
                         "spare pixels to reframe with.")
    ap.add_argument("--workers", type=int, default=default_workers(),
                    help="Clips encoded concurrently (default: cores/4, capped "
                         "at 4). Use 1 to serialize.")
    ap.add_argument("--width", type=int, help="Explicit width (overrides --resolution)")
    ap.add_argument("--height", type=int, help="Explicit height (overrides --resolution)")
    args = ap.parse_args()

    dims = ASPECTS[args.aspect][args.resolution]
    if args.width or args.height:
        if not (args.width and args.height):
            sys.exit("--width and --height must be given together")
        dims = (args.width, args.height)

    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)

    warn_upscale(clips, dims, args.aspect)

    def one(c):
        if c.get("type") == "image":
            print(f"Normalizing {Path(c['file']).name} "
                  f"(still -> {c['duration']:g}s Ken Burns, "
                  f"kb={c.get('kb', 'auto')})...", flush=True)
            return c, normalize_image(c, outdir, dims)
        print(f"Normalizing {Path(c['file']).name} "
              f"({c['profile']}, {c['orientation']}, "
              f"pan={c.get('pan','center')})...", flush=True)
        return c, normalize(c, outdir, args.dlog_lut, dims, not args.no_crop)

    todo = [c for c in clips if "error" not in c]
    if args.workers > 1:
        with ThreadPoolExecutor(max_workers=args.workers) as pool:
            results = list(pool.map(one, todo))
    else:
        results = [one(c) for c in todo]

    done, failed = [], []
    for c, out in results:
        (done if out else failed).append(c["file"])

    shape = (f"{dims[0]}x{dims[1]} ({args.aspect})" if not args.no_crop
             else f"full-frame at {dims[1]}px tall — crop per segment at render")
    print(f"\nDone: {len(done)} mezzanine files in {outdir}/ at {shape}")
    if failed:
        print(f"Failed: {failed}")
        sys.exit(1)


if __name__ == "__main__":
    main()