#!/usr/bin/env python3
"""Render a cut plan through DaVinci Resolve Studio instead of ffmpeg.

Same plan, same gates, same audio design — a different picture engine.
render_reel.py trims and concatenates with ffmpeg and grades with filter
chains. This pushes the plan into a Resolve timeline, reframes each segment
with Resolve's own transform, applies a LUT on the clip's node graph, renders
the picture with Resolve's H.264 encoder, then hands that picture to the same
audio-spine mux render_reel.py uses. The sound design lives in one place
either way.

When to reach for this rather than ffmpeg:
  - the reframe needs to track a subject (SmartReframe, see --smart-reframe)
  - the footage wants Resolve's D-Log conversion rather than our approximation
  - the human is going to keep working the timeline in Resolve afterwards

What stays with ffmpeg on purpose: loudness (Resolve normalises per clip, and
our audio is a mix built at render — measured: near-silent drone clips pegged
to +30 dB), captions styling, transitions, QC and compliance.

Requires Resolve STUDIO running (free refuses external scripting). LUTs must
live under Resolve's LUT directory: the bundled looks are copied into
<LUT dir>/reel-builder/ on first use, since SetLUT refuses absolute paths.

Usage:
  resolve_render.py plan/cut_plan.json --out work/master_resolve.mp4 --look gritty
  resolve_render.py plan/cut_plan.json --out work/master.mp4 --lut my/looks/x.cube
  resolve_render.py plan/cut_plan.json --out work/m.mp4 --smart-reframe --keep-project
  resolve_render.py plan/cut_plan.json --out work/m.mp4 --stabilize
"""
import argparse, json, shutil, subprocess, sys, tempfile, time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from resolve_bridge import (connect, plan_to_appends, fill_zoom, project_fps,   # noqa: E402
                            DEFAULT_FPS)
from render_reel import lay_audio_spine, dims_for, ASPECTS                        # noqa: E402

LUT_DIR = Path("/Library/Application Support/Blackmagic Design/DaVinci Resolve/LUT")
LUT_SUBDIR = "reel-builder"
BUNDLED_LUTS = Path(__file__).resolve().parent.parent / "assets/luts"
LOOKS = {"gritty": "looks/reel-gritty.cube", "punchy": "looks/reel-punchy.cube",
         "faded": "looks/reel-faded.cube"}
POLL_S = 2.0
RENDER_TIMEOUT_S = 1800


def pan_px(src_w, src_h, tl_w, tl_h, pan):
    """Resolve Pan value (pixels) for a cut plan's `pan` on a fill-zoomed clip.

    After fill_zoom the source overflows the frame horizontally by
    (src_w * tl_h / src_h) - tl_w pixels. Pan +overflow/2 shows the LEFT edge
    (image shifted right), -overflow/2 the right edge, 0 the centre — the same
    left/center/right or 0..1 fraction render_reel.pan_expr accepts. A source
    that does not overflow has nowhere to pan and returns 0.
    """
    if not all((src_w, src_h, tl_w, tl_h)):
        return 0.0
    overflow = src_w * (tl_h / src_h) - tl_w
    if overflow <= 0:
        return 0.0
    if pan is None:
        return 0.0
    names = {"left": 0.0, "center": 0.5, "centre": 0.5, "right": 1.0}
    if isinstance(pan, (int, float)):
        f = min(1.0, max(0.0, float(pan)))
    else:
        f = names.get(str(pan).lower(), 0.5)
    return round((0.5 - f) * overflow, 2)


def lut_for(look, lut):
    """Relative LUT path inside Resolve's LUT dir for --look or --lut, or None.

    SetLUT refuses absolute paths (measured), so a bundled look is addressed
    by where install_bundled_luts() puts it, and a caller's --lut must already
    be given relative to the LUT dir.
    """
    if lut:
        return lut
    if look and look != "none":
        if look not in LOOKS:
            sys.exit(f"--look must be one of {', '.join(LOOKS)} (ffmpeg grades "
                     f"like cine/rapha are filter chains with no LUT form yet).")
        return f"{LUT_SUBDIR}/{LOOKS[look]}"
    return None


def install_bundled_luts(lut_dir=LUT_DIR, src=BUNDLED_LUTS):
    """Copy the bundled looks into Resolve's LUT dir. Returns paths written."""
    written = []
    for rel in LOOKS.values():
        s, d = src / rel, lut_dir / LUT_SUBDIR / rel
        if not s.exists():
            continue
        if d.exists() and d.stat().st_size == s.stat().st_size:
            continue
        d.parent.mkdir(parents=True, exist_ok=True)
        shutil.copyfile(s, d)
        written.append(str(d))
    return written


def wait_for_job(status_fn, timeout_s=RENDER_TIMEOUT_S, poll_s=POLL_S, sleep=time.sleep):
    """Poll a render job to a terminal state. Returns the final status dict."""
    waited = 0.0
    while waited < timeout_s:
        st = status_fn() or {}
        if st.get("JobStatus") in ("Complete", "Failed", "Cancelled",
                                   "Background Render Cancelled",
                                   "Remote Render Cancelled"):
            return st
        sleep(poll_s)
        waited += poll_s
    return {"JobStatus": "Timeout", "Error": f"no terminal state after {timeout_s}s"}


def probe_duration(path):
    r = subprocess.run(["ffprobe", "-v", "error", "-show_entries", "format=duration",
                        "-of", "csv=p=0", str(path)], capture_output=True, text=True)
    return float(r.stdout.strip() or 0)


def main():
    ap = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("plan")
    ap.add_argument("--out", required=True)
    ap.add_argument("--project", help="Resolve project to create/load (default: plan folder name)")
    ap.add_argument("--timeline", default="reel", help="Timeline name (default reel)")
    ap.add_argument("--look", default="none", choices=["none", *LOOKS],
                    help="Bundled LUT look applied on each clip's node graph")
    ap.add_argument("--lut", help="LUT path RELATIVE to Resolve's LUT dir (overrides --look)")
    ap.add_argument("--smart-reframe", action="store_true",
                    help="Run Resolve's subject-tracking SmartReframe on every clip "
                         "instead of the plan's static pan (Studio AI, seconds per clip)")
    ap.add_argument("--stabilize", action="store_true",
                    help="Run Resolve's stabilizer on every clip before rendering "
                         "(seconds per clip; the analysis is cached on the project)")
    ap.add_argument("--aspect", choices=sorted(ASPECTS), default="9:16")
    ap.add_argument("--resolution", choices=("1080", "4k"), default="1080")
    ap.add_argument("--fps", type=float, default=None)
    ap.add_argument("--vo-overlay", metavar="FILE")
    ap.add_argument("--vo-start", type=float, default=0.0)
    ap.add_argument("--music-gain-db", type=float, default=-3.0)
    ap.add_argument("--natural-audio", type=float, default=None, metavar="DB")
    ap.add_argument("--keep-project", action="store_true",
                    help="Leave the Resolve project in place for hand editing")
    a = ap.parse_args()

    plan = json.loads(Path(a.plan).read_text())
    segs = plan.get("segments") or []
    if not segs:
        sys.exit(f"{a.plan} has no segments.")
    tl_w, tl_h = dims_for(a.aspect, a.resolution)
    want_fps = a.fps or plan.get("fps") or DEFAULT_FPS
    lut_rel = lut_for(a.look, a.lut)
    if lut_rel and lut_rel.startswith(LUT_SUBDIR):
        install_bundled_luts()

    resolve = connect()
    pm = resolve.GetProjectManager()
    # Whatever the human has open gets saved before this switches projects;
    # CreateProject/LoadProject replace the current project without asking.
    if pm.GetCurrentProject():
        pm.SaveProject()
    name = a.project or f"reel-builder_{Path(a.plan).resolve().parent.parent.name}"
    proj = pm.CreateProject(name) or pm.LoadProject(name)
    if proj is None:
        sys.exit(f"Could not create or load Resolve project {name!r}.")
    if not proj.SetSetting("timelineFrameRate", str(float(want_fps))):
        have = project_fps(proj)
        if abs(have - float(want_fps)) > 1e-6:
            sys.exit(f"Project {name!r} is locked at {have}fps (it holds a timeline); "
                     f"the plan needs {want_fps}. Use --project for a fresh one.")
    proj.SetSetting("timelineResolutionWidth", str(tl_w))
    proj.SetSetting("timelineResolutionHeight", str(tl_h))
    fps = project_fps(proj)
    if lut_rel:
        proj.RefreshLUTList()

    paths = []
    for s in segs:
        rp = Path(s["clip"]).resolve()
        if not rp.exists():
            sys.exit(f"Clip not found: {s['clip']} — run from the project directory.")
        if str(rp) not in paths:
            paths.append(str(rp))
    resolve.GetMediaStorage().AddItemListToMediaPool(paths)
    by_path = {}
    for item in proj.GetMediaPool().GetRootFolder().GetClipList() or []:
        fp = item.GetClipProperty("File Path")
        if fp:
            by_path.setdefault(str(Path(fp).resolve()), item)
    appends, missing = plan_to_appends(segs, by_path, fps)
    if missing:
        sys.exit("Resolve did not import: " + ", ".join(missing[:5]))

    mp = proj.GetMediaPool()
    tl_name = a.timeline
    n = 1
    while any(proj.GetTimelineByIndex(i).GetName() == tl_name
              for i in range(1, proj.GetTimelineCount() + 1)):
        n += 1
        tl_name = f"{a.timeline}_{n}"
    tl = mp.CreateEmptyTimeline(tl_name)
    if tl is None:
        sys.exit(f"Could not create timeline {tl_name!r}.")
    items = mp.AppendToTimeline(appends) or []
    if len(items) != len(segs):
        sys.exit(f"Resolve appended {len(items)} of {len(segs)} segments.")

    reframed, luts, stabilized = 0, 0, 0
    for seg, item in zip(segs, items):
        if a.stabilize and item.Stabilize():
            stabilized += 1
        mpi = item.GetMediaPoolItem()
        try:
            sw, sh = (int(x) for x in mpi.GetClipProperty("Resolution").split("x"))
        except (AttributeError, ValueError):
            sw = sh = 0
        z = fill_zoom(sw, sh, tl_w, tl_h)
        props = {"ZoomX": z, "ZoomY": z} if z > 1.0 else {}
        if a.smart_reframe:
            if item.SmartReframe():
                reframed += 1
        else:
            props["Pan"] = pan_px(sw, sh, tl_w, tl_h, seg.get("pan"))
        if props and item.SetProperties(props):
            reframed += 0 if a.smart_reframe else 1
        if lut_rel:
            g = item.GetNodeGraph()
            if g and g.SetLUT(1, lut_rel):
                luts += 1
    if lut_rel and luts != len(items):
        sys.exit(f"LUT {lut_rel!r} applied to {luts}/{len(items)} clips — is it "
                 f"under {LUT_DIR}? Run with --lut relative to that directory.")

    out = Path(a.out).resolve()
    out.parent.mkdir(parents=True, exist_ok=True)
    music = plan.get("music")
    needs_mux = bool(music or a.vo_overlay or a.natural_audio is not None)
    with tempfile.TemporaryDirectory() as td:
        target_dir = Path(td) if needs_mux else out.parent
        custom = "picture" if needs_mux else out.stem
        proj.SetCurrentRenderFormatAndCodec("mp4", "H264")
        ok = proj.SetRenderSettings({
            "SelectAllFrames": True, "TargetDir": str(target_dir), "CustomName": custom,
            "FormatWidth": tl_w, "FormatHeight": tl_h, "FrameRate": float(fps),
            "ExportVideo": True, "ExportAudio": not needs_mux or a.natural_audio is not None,
            "AudioCodec": "aac", "AudioSampleRate": 48000, "VideoQuality": 0})
        if not ok:
            sys.exit("Resolve rejected the render settings.")
        job = proj.AddRenderJob()
        if not job or not proj.StartRendering([job], False):
            sys.exit("Resolve did not start the render job.")
        st = wait_for_job(lambda: proj.GetRenderJobStatus(job))
        if st.get("JobStatus") != "Complete":
            sys.exit(f"Resolve render {st.get('JobStatus')}: {st.get('Error', '')}")
        picture = target_dir / f"{custom}.mp4"
        if not picture.exists():
            sys.exit(f"Render reported Complete but {picture} is missing.")

        if needs_mux:
            total = sum(float(s["duration"]) for s in segs)
            tmp = out.with_name(out.stem + ".tmp" + out.suffix)
            lay_audio_spine(picture, tmp, plan, total, [], None if a.vo_overlay else music,
                            a.natural_audio, a.vo_overlay, a.vo_start, a.music_gain_db,
                            False, "18", "192k")
            tmp.rename(out)

    if not a.keep_project:
        pm.CloseProject(proj)
        pm.DeleteProject(name)

    print(json.dumps({
        "out": str(out), "duration": round(probe_duration(out), 2), "engine": "resolve",
        "project": name if a.keep_project else None, "timeline": tl_name,
        "fps": fps, "size": f"{tl_w}x{tl_h}", "lut": lut_rel,
        "smart_reframe": a.smart_reframe, "stabilized": stabilized if a.stabilize else None,
        "render_ms": st.get("TimeTakenToRenderInMs"),
    }, indent=2))
    print("\nQC it like any master: check_cuts.py, then check_export.py after "
          "export_variants.py.", file=sys.stderr)


if __name__ == "__main__":
    main()