#!/usr/bin/env python3
"""take-master — sum an Ardour take's orbit stems into one mastered stereo FLAC.

WHY THIS EXISTS RATHER THAN `tidal-ears master mix`

The sister tool's mix is the right recipe and this reuses it verbatim — uniform
trim, 20 Hz/20 kHz guard, gentle bus compression, edge trim, two-pass loudnorm.
What it does not know is how ARDOUR LAYS OUT A TAKE, and that produced two
defects that only show up on this material:

  1. MONO. Ardour writes each orbit as TWO mono files, `Take95_Tidal 04-1%L.wav`
     and `...%R.wav`. Globbing them all into `amix` sums 24 mono streams into ONE
     channel. The stereo image is destroyed — and ParVagues pans deliberately
     (`# pan 0.8` on the hats, `0.2` on the sleepwalker), so this is not a
     technicality, it silently discards a compositional layer. Here each orbit's
     pair is `join`ed into stereo FIRST, and the mix is 12 stereo streams.

  2. 192 kHz. `loudnorm` internally resamples to 192 kHz to measure true peak,
     and with no explicit output rate ffmpeg keeps it. A 48 kHz session came out
     as a 624 MB 192 kHz file carrying no extra information. This pins `-ar` to
     the source rate.

WHY THE TRIM IS UNIFORM

`gain_for_stem` in the sister tool normalises each stem TOWARD a target peak.
That is right for stems of unknown provenance and wrong for these: PLN balanced
these twelve orbits by hand, on faders, while performing. Normalising them
individually rewrites that balance — measured at 4.0 dB of shift on Take95,
mostly making the hats hotter. A single trim applied to every stem preserves the
mix he played and lets loudnorm set the absolute level, which is also what the
take-lens work concluded independently.

The default trim is derived from the hottest stem, so the sum has headroom
without anyone choosing a number by feel.

Usage:
    tools/take-master.py --take Take95 -o output/Take95/Take95_master.flac
    tools/take-master.py --take Take95 --trim -15 -o master.flac
    tools/take-master.py --take Take95 --dry-run
"""
from __future__ import annotations

import argparse
import glob
import json
import math
import os
import pathlib
import re
import subprocess
import sys

ARDOUR = pathlib.Path.home() / "Work/Sound/Ardour"
# "Take95_Tidal 04-1%L.wav" -> take, orbit, side
STEM = re.compile(r"^(?P<take>[^_]+)_Tidal\s+(?P<orbit>\d+)-\d+%(?P<side>[LR])\.wav$")


def sh(cmd: list[str], timeout: int = 14400) -> tuple[int, str]:
    p = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
    return p.returncode, (p.stdout or "") + (p.stderr or "")


def find_pairs(take: str) -> list[tuple[int, str, str]]:
    """[(orbit, left_path, right_path)] sorted by orbit. Refuses on a lone side."""
    found: dict[int, dict[str, str]] = {}
    for p in glob.glob(str(ARDOUR / "**" / f"{take}_Tidal *.wav"), recursive=True):
        m = STEM.match(os.path.basename(p))
        if m and m.group("take") == take:
            found.setdefault(int(m.group("orbit")), {})[m.group("side")] = p
    if not found:
        raise SystemExit(f"take-master: no stems matching {take}_Tidal NN-N%[LR].wav "
                         f"under {ARDOUR}")
    lonely = [o for o, s in found.items() if len(s) != 2]
    if lonely:
        # Silently mixing a half-pair would pan that orbit hard to one side.
        raise SystemExit(f"take-master: orbit(s) {sorted(lonely)} have only one side — "
                         f"refusing to guess the other channel")
    return [(o, found[o]["L"], found[o]["R"]) for o in sorted(found)]


def probe(path: str) -> tuple[int, float]:
    rc, out = sh(["ffprobe", "-v", "error", "-select_streams", "a:0",
                  "-show_entries", "stream=sample_rate:format=duration",
                  "-of", "json", path], timeout=120)
    if rc != 0:
        raise SystemExit(f"take-master: ffprobe failed on {path}")
    j = json.loads(out)
    return (int(j["streams"][0]["sample_rate"]), float(j["format"]["duration"]))


def peak_dbfs(path: str) -> float:
    """True peak level in dBFS, which for these stems is ABOVE zero.

    astats, NOT volumedetect. volumedetect reports max_volume against a
    fixed-point full scale and SATURATES at 0.0 dB, so every one of these float
    stems read as exactly +0.0 — which silently turned a 19 dB trim into a 6 dB
    one. astats reports the real value: Tidal 04 peaks at +13.3 dBFS.
    """
    rc, out = sh(["ffmpeg", "-hide_banner", "-nostats", "-i", path, "-af",
                  "astats=measure_overall=Peak_level:measure_perchannel=none",
                  "-f", "null", "-"], timeout=3600)
    m = re.search(r"Peak level dB:\s*(-?\d+(?:\.\d+)?|-?inf)", out)
    if not m:
        raise SystemExit(f"take-master: could not measure peak of {path}")
    v = m.group(1)
    return -math.inf if v.endswith("inf") else float(v)


def build_filter(n_pairs: int, trim_db: float, sr: int, loudnorm: str,
                 trim_edges: bool) -> str:
    parts, labels = [], []
    for i in range(n_pairs):
        l, r = 2 * i, 2 * i + 1
        # join, not amerge: amerge's channel layout for two mono inputs is
        # version-dependent and has been known to produce a 2-channel stream that
        # downstream filters treat as mono-duplicated.
        parts.append(f"[{l}:a][{r}:a]join=inputs=2:channel_layout=stereo[s{i}]")
        parts.append(f"[s{i}]volume={trim_db:.2f}dB,highpass=f=20,lowpass=f=20000[a{i}]")
        labels.append(f"[a{i}]")
    parts.append(f"{''.join(labels)}amix=inputs={n_pairs}"
                 f":duration=longest:normalize=0[mix]")
    chain = ("[mix]acompressor=threshold=-22dB:ratio=1.5"
             ":attack=30:release=100:makeup=1")
    if trim_edges:
        # Trim the pre/post-set silence BEFORE loudnorm's gate sees it, so the
        # measured integrated loudness describes the music and not the room.
        # Never fill it back with noise — PLN's standing rule is trim, not fill.
        sr_f = ("silenceremove=start_periods=1:start_duration=0.5"
                ":start_threshold=-60dB:detection=peak")
        chain += f",{sr_f},areverse,{sr_f},areverse"
    chain += f",{loudnorm},aresample={sr}[out]"
    parts.append(chain)
    return ";".join(parts)


def main() -> int:
    ap = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--take", required=True)
    ap.add_argument("-o", "--output", type=pathlib.Path)
    ap.add_argument("--trim", type=float,
                    help="uniform per-stem trim in dB (default: derived from the "
                         "hottest stem so the sum has headroom)")
    ap.add_argument("--headroom", type=float, default=-6.0,
                    help="where the hottest stem should land (default -6 dBFS)")
    ap.add_argument("--lufs", type=float, default=-14.0)
    ap.add_argument("--tp", type=float, default=-1.0)
    ap.add_argument("--lra", type=float, default=11.0)
    ap.add_argument("--no-trim-edges", action="store_true")
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    pairs = find_pairs(args.take)
    sr, dur = probe(pairs[0][1])
    print(f"take-master: {args.take} — {len(pairs)} orbits "
          f"({', '.join('d%d' % o for o, _, _ in pairs)}), "
          f"{dur/60:.1f} min @ {sr} Hz\n")

    trim = args.trim
    if trim is None:
        print("  measuring peaks (one pass per stem)…")
        peaks = []
        for o, l, r in pairs:
            pk = max(peak_dbfs(l), peak_dbfs(r))
            peaks.append((o, pk))
            print(f"    d{o:<3} peak {pk:+6.1f} dBFS")
        hottest = max(p for _, p in peaks)
        trim = args.headroom - hottest
        print(f"\n  hottest stem {hottest:+.1f} dBFS -> uniform trim {trim:+.1f} dB "
              f"(lands it at {args.headroom:.1f})")
    else:
        print(f"  uniform trim {trim:+.1f} dB (given)")

    if args.dry_run:
        print("\n  --dry-run: nothing rendered")
        return 0
    if not args.output:
        raise SystemExit("take-master: -o/--output is required unless --dry-run")

    inputs: list[str] = []
    for _, l, r in pairs:
        inputs += ["-i", l, "-i", r]

    # Pass 1: measure. loudnorm needs real measurements to hit the target without
    # pumping; a single-pass loudnorm on an hour of dynamic material drifts.
    print("\n  loudnorm pass 1 (measuring)…")
    f1 = build_filter(len(pairs), trim, sr,
                      f"loudnorm=I={args.lufs}:TP={args.tp}:LRA={args.lra}"
                      f":print_format=json", not args.no_trim_edges)
    rc, out = sh(["ffmpeg", "-y", "-hide_banner", "-nostats", *inputs,
                  "-filter_complex", f1, "-map", "[out]", "-f", "null", "-"])
    if rc != 0:
        print(out[-3000:], file=sys.stderr)
        raise SystemExit("take-master: measuring pass failed")
    blocks = re.findall(r"\{[^{}]*input_i[^{}]*\}", out, re.S)
    if not blocks:
        print(out[-3000:], file=sys.stderr)
        raise SystemExit("take-master: could not parse loudnorm measurements")
    m = json.loads(blocks[-1])
    print(f"    measured: I={m['input_i']} TP={m['input_tp']} "
          f"LRA={m['input_lra']} thresh={m['input_thresh']}")

    ln2 = (f"loudnorm=I={args.lufs}:TP={args.tp}:LRA={args.lra}"
           f":measured_I={m['input_i']}:measured_TP={m['input_tp']}"
           f":measured_LRA={m['input_lra']}:measured_thresh={m['input_thresh']}"
           f":offset={m['target_offset']}:linear=true:print_format=summary")

    args.output.parent.mkdir(parents=True, exist_ok=True)
    print(f"\n  rendering -> {args.output}")
    rc, out = sh(["ffmpeg", "-y", "-hide_banner", "-nostats", *inputs,
                  "-filter_complex",
                  build_filter(len(pairs), trim, sr, ln2, not args.no_trim_edges),
                  "-map", "[out]",
                  "-ar", str(sr), "-ac", "2",       # the two defects, pinned shut
                  "-c:a", "flac", "-compression_level", "5", str(args.output)])
    if rc != 0:
        print(out[-3000:], file=sys.stderr)
        raise SystemExit("take-master: render failed")

    rsr, rdur = probe(str(args.output))
    rc, o2 = sh(["ffprobe", "-v", "error", "-select_streams", "a:0",
                 "-show_entries", "stream=channels", "-of", "csv=p=0",
                 str(args.output)], timeout=120)
    ch = o2.strip()
    size = args.output.stat().st_size / 1e6
    print(f"\ntake-master: done — {args.output}")
    print(f"    {ch} ch @ {rsr} Hz, {rdur/60:.1f} min, {size:.0f} MB")
    if ch != "2" or rsr != sr:
        print(f"    WARNING: expected 2 ch @ {sr} Hz", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())
