#!/usr/bin/env python3
"""build_judge_set — precompute the set-judge SPA's data for a whole gig.

WHY THIS EXISTS
---------------
Before every release we ask the same question, by hand, in a music player:
"of the N tracks this set split into, which ones ship?" It takes an evening,
the answers live in a chat scrollback, and three weeks later nobody can say why
track 11 was cut. PLN, 2026-08-16: *"its a reusable tool tbh we do this flow
often"*.

So: same evidence the A/B Judge already shows (orbit lanes labelled from the
SCORE, activity measured from the STEMS), but the unit of judgement is a LIST —
15 tracks, each with a verdict and a comment, exported as one decisions.json
the release tooling reads directly.

WHAT IS NOT GUESSED
-------------------
* **Orbit labels come from the .tidal score**, via tidal_score.orbit_sounds —
  the actual sound on each dN, not a role inferred from the orbit number.
* **Role family is validated against audio** (spectral centroid at that orbit's
  loudest moment in THAT track), never from the sample name alone. A name-only
  classification is exactly the mistake `meth_bass` is the standing example of:
  it is a noisy wobble, not a sub.
* **Activity is measured** from the per-orbit stems, on the time axis, so a
  sparse pattern and a fade are distinguishable.

THE CALIBRATION THAT MATTERS
----------------------------
Clip time is MASTER time; the stems are on TAKE time, and the master is an
edited subset of the take. `keeps` is that edit as explicit numbers — the same
list the mix plan carries — so master→stem is exact and piecewise, including
across a cut. Deriving it by content-detection instead would slide every
boundary; see MIXING.md. The sum of the keeps must equal the master duration,
and this script asserts it rather than trusting it.

USAGE
    python3 build_judge_set.py SPEC.json [--out PATH]

    SPEC.json — see judge_specs/opal26.json for a worked example.
"""
from __future__ import annotations

import argparse
import json
import re
import subprocess
import sys
from pathlib import Path

import numpy as np

HERE = Path(__file__).parent
sys.path.insert(0, str(HERE))
from audio_lens import classify_family, load_window, profile          # noqa: E402
from models import (JudgeSet, JudgeTrack, MasterInfo, OrbitActivity,  # noqa: E402
                    RoleGroup, Variant)
from tidal_score import orbit_sounds                                  # noqa: E402

# Role families (DESIGN.md). Colour reinforces label + glyph + lane; never hue alone.
ROLE_GROUPS = [
    RoleGroup(key="percs",   label="Percs",   color="#ff8c00", glyph="▰"),
    RoleGroup(key="bass",    label="Bass",    color="#7c5cff", glyph="▂"),
    RoleGroup(key="melodic", label="Melodic", color="#36c5f0", glyph="♪"),
    RoleGroup(key="tops",    label="Tops",    color="#2dd4bf", glyph="≈"),
    RoleGroup(key="atmos",   label="Atmos",   color="#8a93a6", glyph="◌"),
]
ACTIVE_DB = -38.0     # "strongly on / driving"
LIT_FLOOR = -52.0     # audible floor — an orbit dims toward this, it never vanishes


_EMPH = re.compile(r"(\*\*|\*|__|_)(?=\S)(.+?)(?<=\S)\1")


def clean_title(s: str) -> str:
    """Strip markdown emphasis the backlog parse leaves in track names.

    tracks.json carries titles like `There's **Something About Drums**` and
    `Am i _Doing it Right_` — backlog.md is prose and the emphasis is PLN's, not
    part of the name. These titles end up in decisions.json and from there in
    upload metadata, so the asterisks would ship. Only MARKUP is removed: "<3"
    and capitalisation are his, and stay.
    """
    prev = None
    while prev != s:
        prev, s = s, _EMPH.sub(r"\2", s)
    return s.strip()


def slugify(s: str) -> str:
    s = re.sub(r"[^\w\s-]", "", s.lower())
    return re.sub(r"[\s_-]+", "-", s).strip("-")


def ffprobe_dur(path: Path) -> float:
    r = subprocess.run(["ffprobe", "-v", "error", "-show_entries", "format=duration",
                        "-of", "default=nw=1:nk=1", str(path)],
                       capture_output=True, text=True)
    try:
        return round(float(r.stdout.strip()), 3)
    except ValueError:
        return 0.0


def loudness(path: Path) -> dict:
    """Integrated LUFS / LRA / true peak for one file."""
    r = subprocess.run(["ffmpeg", "-hide_banner", "-nostats", "-i", str(path),
                        "-af", "ebur128=peak=true", "-f", "null", "-"],
                       capture_output=True, text=True)
    summ = r.stderr.split("Summary:")[-1]

    def g(pat):
        m = re.search(pat, summ)
        return round(float(m.group(1)), 1) if m else None
    return {"I": g(r"I:\s*(-?[\d.]+)\s*LUFS"),
            "LRA": g(r"LRA:\s*(-?[\d.]+)\s*LU"),
            "TP": g(r"Peak:\s*(-?[\d.]+)\s*dBFS")}


def peaks(path: Path, n: int = 900) -> list[float]:
    """Waveform overview: `n` buckets of absolute peak, normalised to 0..1.

    Decoded at full rate and bucket-MAXed rather than resampled — resampling
    smooths transients away, and a waveform whose peaks are averaged out is a
    picture of a different, quieter track. Streamed through a pipe so an 80 min
    file never has to be resident.
    """
    dur = ffprobe_dur(path)
    if dur <= 0:
        return []
    proc = subprocess.Popen(
        ["ffmpeg", "-v", "error", "-i", str(path), "-ac", "1", "-ar", "48000",
         "-f", "f32le", "-"], stdout=subprocess.PIPE)
    per = max(1, int(dur * 48000 / n))
    out, buf = [], b""
    try:
        while True:
            chunk = proc.stdout.read(per * 4 * 8)
            if not chunk:
                break
            buf += chunk
            usable = (len(buf) // (per * 4)) * (per * 4)
            if usable:
                a = np.frombuffer(buf[:usable], dtype=np.float32).reshape(-1, per)
                out.extend(np.abs(a).max(axis=1).tolist())
                buf = buf[usable:]
        if buf:
            a = np.frombuffer(buf[:len(buf) // 4 * 4], dtype=np.float32)
            if a.size:
                out.append(float(np.abs(a).max()))
    finally:
        proc.stdout.close()
        proc.wait()
    hi = max(out) if out else 0.0
    return [round(float(v / hi), 4) for v in out] if hi > 0 else []


def make_master_to_stem(keeps, master_dur=None):
    """master-time -> stem-time, piecewise across the edit.

    `keeps` is [[a,b], ...] in STEM time; the master is their concatenation. A
    bin landing inside a cut would otherwise read the wrong orbit's activity —
    which is how a judge UI ends up confidently lighting a lane that is not
    playing.
    """
    total = sum(b - a for a, b in keeps)
    if master_dur is not None and abs(total - master_dur) > 0.05:
        raise SystemExit(
            f"keeps sum to {total:.3f}s but the master is {master_dur:.3f}s — "
            "the edit map is wrong, and every orbit lane would be shifted. "
            "Fix `keeps` in the spec (see MIXING.md: edits are explicit numbers).")

    def f(t: float) -> float:
        acc = 0.0
        for a, b in keeps:
            span = b - a
            if t < acc + span:
                return a + (t - acc)
            acc += span
        return keeps[-1][1]
    return f


def track_orbits(tidal_path, stemmap, m2s, start_s, dur_s, bin_s, stems_dir, glob):
    """Per-orbit lanes for one track: label from the score, activity from the stems,
    family validated by the measured centroid at that orbit's loudest moment."""
    if not tidal_path or not Path(tidal_path).exists():
        return []
    sounds = orbit_sounds(tidal_path)              # {orbit:int -> {sound, raw_all, ...}}
    rms = stemmap["rms_db"]
    names = stemmap["orbits"]
    sm_bin = stemmap.get("bin_s", bin_s)
    n_bins = int(round(dur_s / bin_s))

    stem_paths = sorted(Path(stems_dir).glob(glob))
    out = []
    for i, name in enumerate(names):
        # Stem-map rows are named after the FILE, and those names differ per
        # export ("orbit-03", "Opal26_clean_Tidal 03"). The orbit is the LAST
        # integer in the name — "Opal26" would otherwise win and every lane
        # would be attributed to orbit 26.
        m = re.search(r"(\d+)\s*$", str(name))
        if not m:
            continue
        o = int(m.group(1))
        if o not in sounds:                        # not triggered in this score
            continue
        row = rms[i]
        act = []
        for b in range(n_bins):
            st = m2s(start_s + b * bin_s)          # master -> stem, per bin
            idx = int(round(st / sm_bin))
            act.append(round(float(row[idx]), 1) if 0 <= idx < len(row) else -240.0)
        if not act or max(act) <= LIT_FLOOR:       # never audible in this track
            continue

        sound = sounds[o]["sound"]
        # Measure, don't infer: profile this orbit's stem at its loudest bin here.
        prof = None
        if i < len(stem_paths):
            t_peak = m2s(start_s + int(np.argmax(act)) * bin_s)
            try:
                prof = profile(load_window(stem_paths[i], max(0, t_peak - 8), t_peak + 8))
            except Exception:
                prof = None
        group, cen = classify_family(sound, prof)
        out.append(OrbitActivity(
            orbit=f"{o:02d}", label=sound, sound=sound, group=group,
            centroid=round(cen) if cen else None, activity=act))
    out.sort(key=lambda x: x.orbit)
    return out


def build(spec: dict) -> JudgeSet:
    root = Path(spec.get("liveRoot", "/home/pln/Work/Sound/Tidal"))
    clips_dir = Path(spec["clipsDir"])
    tracks_dir = Path(spec["tracksDir"]) if spec.get("tracksDir") else None
    stemmap = json.loads(Path(spec["stemmap"]).read_text())
    meta = json.loads(Path(spec["tracksJson"]).read_text())
    bin_s = float(spec.get("binS", stemmap.get("bin_s", 2.0)))
    variant = Variant(spec.get("variant", "stream"))

    master_dur = ffprobe_dur(Path(spec["master"])) if spec.get("master") else None
    m2s = make_master_to_stem(spec["keeps"], master_dur)

    tracks = []
    for i, t in enumerate(meta["tracks"], start=1):
        title = clean_title(t["name"])
        # Clips are named "NN - Title.flac" by `master split`; match on the index,
        # not on the title — titles get punctuation-normalised and would miss.
        clip = next((p for p in sorted(clips_dir.glob("*.flac"))
                     if p.name.startswith(f"{i:02d} ")), None)
        full = next((p for p in sorted(tracks_dir.glob("*.flac"))
                     if p.name.startswith(f"{i:02d} ")), None) if tracks_dir else None
        if clip is None:
            print(f"  ! no clip for #{i} {title} — skipped", file=sys.stderr)
            continue
        clip_dur = ffprobe_dur(clip)
        # Activity spans the WHOLE track, never just the head. Measured: the head
        # is bit-exact with the track's first 16 s (r=1.0000), i.e. it is the
        # INTRO — 1 to 4 orbits where the track proper runs nine. Lanes built from
        # the head would show a track as thin when it is only starting quietly,
        # and that is a ship/cut call made on the wrong evidence. One activity
        # array, indexed by track time, serves both transports.
        track_dur = t.get("duration_s", t["end_s"] - t["start_s"])
        tidal = str(root / t["file"]) if t.get("file") else None

        masters = {}
        if full is not None:
            ld = loudness(full)
            masters[variant] = MasterInfo(file=full.name, **ld)

        orbits = track_orbits(tidal, stemmap, m2s, t["start_s"], track_dur,
                              bin_s, spec["stemsDir"], spec.get("stemsGlob", "*.wav"))
        tracks.append(JudgeTrack(
            n=i, id=slugify(title), title=title,
            audio=f"{spec['audioBase']}/{clip.name}",
            fullAudio=f"{spec['fullAudioBase']}/{full.name}" if full else None,
            startS=t["start_s"], endS=t["end_s"], durS=clip_dur,
            trackDurS=track_dur,
            bpm=t.get("bpm"), section=t.get("section", ""), style=t.get("style", ""),
            tidal=tidal, binS=bin_s, orbits=orbits, masters=masters,
            peaks=peaks(full) if full else [],
            samples=t.get("samples", []),
        ))
        lanes = ", ".join(f"{o.label}·{o.orbit}[{o.group}]" for o in orbits)
        i_db = masters[variant].I if variant in masters else None
        print(f"  #{i:>2} {title[:30]:<30} {clip_dur:>5.1f}s "
              f"I={i_db if i_db is not None else '—':>6} "
              f"{len(orbits)} orbits → {lanes}", flush=True)

    return JudgeSet(
        gig=spec["gig"], title=spec["title"], date=spec["date"],
        calibration=spec.get("calibration", ""),
        activeDb=ACTIVE_DB, litFloorDb=LIT_FLOOR, roleGroups=ROLE_GROUPS,
        note=spec.get("note", ""), variant=variant, tracks=tracks,
    )


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("spec", help="path to a judge-set spec JSON")
    ap.add_argument("--out", help="output path (default: ui/public/judge-<gig>.json)")
    args = ap.parse_args()

    spec = json.loads(Path(args.spec).read_text())
    print(f"building judge set for {spec['gig']}…", flush=True)
    data = build(spec)

    out = Path(args.out) if args.out else (
        HERE.parent / "ui" / "public" / f"judge-{data.gig}.json")
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(data.model_dump_json())        # validated on the way out
    print(f"✓ {out}  ({out.stat().st_size / 1024:.0f} KB, {len(data.tracks)} tracks)")
    return 0


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