#!/usr/bin/env python3
"""Mix the soundtrack onto the rendered video from one JSON config, then encode the music version.

  python3 mix.py mix.json out.mp4

mix.json (times are film seconds; everything except video, music and duration is optional):
  {"video": "film.mp4", "music": "score.mp3", "offset_ms": 32, "duration": 65.0,
   "impacts": [6.5, 59.5], "impact_file": "sfx/impact.mp3",
   "riser_end": 59.45, "riser_file": "sfx/riser.mp3", "riser_peak": 0.82,
   "whips": [18.0, 40.0], "pushes": [27.0], "whoosh_file": "sfx/whoosh.mp3", "whoosh_peak": 0.14,
   "clicks": [11.0, 42.0], "tick_file": "sfx/tick.wav",
   "vo": [[0, 12.23, 1.68], ...], "vo_dir": "vo-final", "vo_gain": 1.25, "vo_shift": -0.08, "duck": 0.55,
   "extras": [["path.wav", 30.0, 1.0]]}

What it does, and why:
  - music is delayed by offset_ms so its beats land on the cuts (audio-analyze.py gives the value)
  - under each voice line [i, start, length] the music ducks by `duck` (0.55 = -7 dB) with 150 ms ramps,
    and only there; every line is first matched to -15 LUFS, then scaled by vo_gain
  - each click gets the tick, levelled ~9 dB above the music's highs around it, so it reads over strings
  - one fixed gain brings the mix to -14 LUFS and a limiter catches the peaks. Not loudnorm: in its
    dynamic mode it pumps, pulling the music down for seconds after every loud hit
  - the music fades over the last 1.6 s; the video stream is copied untouched
"""
import json, os, subprocess, sys
import numpy as np
from common import ffmpeg, decode, lufs

C = json.load(open(sys.argv[1])); out = sys.argv[2]; FF = ffmpeg()
DUR, off = C["duration"], C.get("offset_ms", 0)
inputs, fc, parts = [C["video"], C["music"]], [], ["[mus]"]
duck = "0"
for i, t, end in C.get("vo", []):
    duck = f"max({duck},clip((t-{t - 0.12:.3f})/0.15,0,1)*clip(({t + end + 0.2:.3f}-t)/0.15,0,1))"
fc.append(f"[1:a]aformat=sample_fmts=fltp:channel_layouts=stereo,adelay={off:.0f}|{off:.0f},volume='1-{C.get('duck', 0.55)}*{duck}':eval=frame,"
          f"afade=t=out:st={DUR - 1.6}:d=1.6[mus]")


def place(path, times, gain=1.0, gains=None):
    if not path or not times: return
    inputs.append(path); idx = len(inputs) - 1; base = f"s{idx}"
    fc.append(f"[{idx}:a]aformat=sample_fmts=fltp:channel_layouts=stereo,volume={gain},asplit={len(times)}" + "".join(f"[{base}_{k}]" for k in range(len(times))))
    for k, at in enumerate(times):
        ms = max(0, int(round(at * 1000)))
        fc.append(f"[{base}_{k}]volume={(gains[k] if gains else 1):.3f},adelay={ms}|{ms}[{base}o{k}]"); parts.append(f"[{base}o{k}]")


place(C.get("impact_file"), C.get("impacts", []), 0.3)
if C.get("riser_end"): place(C.get("riser_file"), [C["riser_end"] - C.get("riser_peak", 0.82)], 0.6)
place(C.get("whoosh_file"), [t + 0.22 - C.get("whoosh_peak", 0.14) for t in C.get("whips", [])] + [t + 0.25 - C.get("whoosh_peak", 0.14) for t in C.get("pushes", [])], 0.28)
clicks = []
for t in sorted(C.get("clicks", [])):
    if not clicks or t - clicks[-1] > 0.15: clicks.append(t)
if clicks and C.get("tick_file"):
    SR = 44100
    hp = decode(C["music"], SR, af=f"adelay={off:.0f},highpass=f=2000,highpass=f=2000")
    tk = decode(C["tick_file"], SR, af="highpass=f=2000"); tkdb = 20 * np.log10(np.sqrt(np.mean(tk[: int(0.06 * SR)] ** 2)) + 1e-9)
    def level(t):
        seg = hp[max(0, int((t - 0.25) * SR)):int((t + 0.1) * SR)]
        local = 20 * np.log10(np.sqrt(np.mean(seg ** 2)) + 1e-9) if len(seg) else -90
        return float(np.clip(10 ** ((max(local, -45) + 9 - tkdb) / 20), 0.25, 1.6))
    place(C["tick_file"], clicks, 1.0, [level(t) for t in clicks])
for i, t, end in C.get("vo", []):
    p = next(q for q in (f"{C['vo_dir']}/m{i}.wav", f"{C['vo_dir']}/m{i}.mp3") if os.path.exists(q))
    place(p, [t + C.get("vo_shift", 0.0)], round(min(C.get("vo_gain", 1.0) * 10 ** ((-15.0 - lufs(p)) / 20), 4.0), 3))
for path, start, g in C.get("extras", []):
    place(path, [start], g)
fc.append("".join(parts) + f"amix=inputs={len(parts)}:normalize=0:duration=first[mix]")
pre = out + ".pre.wav"
args = [FF, "-v", "error", "-y"] + sum([["-i", p] for p in inputs], [])
subprocess.run(args + ["-filter_complex", ";".join(fc), "-map", "[mix]", "-t", str(DUR), "-ar", "48000", pre], check=True)
gain = -14.0 - lufs(pre)
subprocess.run([FF, "-v", "error", "-y", "-i", C["video"], "-i", pre, "-filter_complex",
                f"[1:a]volume={gain:.2f}dB,alimiter=limit=0.84:attack=3:release=60:level=false[a]", "-map", "0:v", "-map", "[a]",
                "-c:v", "copy", "-c:a", "aac", "-b:a", "192k", "-t", str(DUR), "-movflags", "+faststart", out], check=True)
os.remove(pre)
print(f"wrote {out}: gain {gain:+.1f} dB to -14 LUFS, {len(clicks)} clicks, {len(C.get('vo', []))} voice lines")