"""stempack — batch the Foundry over a folder of ALREADY-separated producer stems.

The Foundry's engine1 exists to *make* stems (yt-dlp → demucs). A producer stem
pack skips that entirely: the stems arrive labelled by the human who mixed them
(`MAREA MIX10 123BPM KIT STEM.wav`). This module is the missing batch driver
(TODO #20) for that case — walk a pack, cut every stem, grade, and land one
SuperDirt kit per track × role family.

    python3 -m engine.stempack "<pack dir>" --tracks MAREA ANGIE --prefix fred

Three things it does that `find_takes` could not:

1. **Resample to 44.1 kHz before analysis.** These packs ship at 96 kHz, and
   librosa's hop/window defaults are sample-rate-relative — at 96 k every analysis
   frame spans half the musical time and beat tracking wanders. Measured on
   `MAREA … 123BPM KIT`, whose filename is free ground truth: at 96 k the finder
   returned 123/123/124/**119**/**128.1**/**128.5** bpm with a 0.661 top score; at
   44.1 k it returned 123.0–123.1 on every candidate with a 0.859 top score, in a
   third of the time. Rate is a correctness input here, not a performance knob.

2. **Role names beyond the demucs four.** `find_takes` keys stems by
   drums/bass/other/vocals — a dict that collides on KEYS1+KEYS2 and an export
   `keep=` filter that would silently drop every stem in this pack. `roles.py`
   maps the pack's ~40 tokens to a family and then *verifies* the label by
   measurement before anything downstream trusts it.

3. **Loop-vs-chop by family, arbitrated by grade.** Composite rhythmic stems want
   bar loops; isolated hits and voices want chops; where the family is ambiguous
   both are run and `grade` picks. Every bar loop is exported at an exact bar
   multiple so `loopAt <bars>` is correct by construction.
"""
from __future__ import annotations

import json
import re
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Callable, Optional

import numpy as np
import soundfile as sf

from . import grade as G
from . import loops as L
from . import roles as R

ANALYSIS_SR = 44100          # see module docstring §1 — a correctness constant
MIN_TIER = "B"               # below this a loop does not ship
TIER_ORDER = ["D", "C", "B", "A", "S"]


@dataclass
class CutLoop:
    """One exported loop, with everything needed to play and to justify it."""
    kit: str
    name: str
    path: str
    track: str
    stem_role: str            # the producer's own token ("all drums")
    family: str               # our effective family, post-verification
    mode: str                 # "loop" | "chop"
    bars: int
    bpm: float
    dur_s: float
    start_s: float
    finder_score: float
    tier: str
    grade: float
    flags: list
    tags: dict                # CLAP multi-axis tags, filled in later (may be {})


def _load_44k(path: Path) -> tuple[np.ndarray, np.ndarray, int]:
    """Read a stem → (stereo (n,ch) @44.1k, mono (n,) @44.1k, source_sr).

    Everything downstream — the grid, the candidate times, the zero-crossing snap
    and the written audio — lives at one rate, so the file that ships is exactly
    the signal that was measured.
    """
    import librosa
    y, sr = sf.read(str(path), always_2d=True, dtype="float32")
    if sr != ANALYSIS_SR:
        # soxr_mq, not the soxr_hq default: ~3x faster and its passband is already
        # far cleaner than anything the downstream onset/MFCC features can resolve.
        y = librosa.resample(y.T, orig_sr=sr, target_sr=ANALYSIS_SR,
                             res_type="soxr_mq").T
    return np.ascontiguousarray(y), y.mean(axis=1), sr


def _slice_export(y_st: np.ndarray, cand, *, rotate: bool = True,
                  mono_sum: bool = False, bar_bpm: float = 0.0) -> tuple[np.ndarray, int]:
    """Cut a candidate out of the stereo stem honouring the B-rules.

    Reuses `loops`' own primitives rather than restating them: the length-preserving
    zero-crossing snap (so the duration stays an exact bar multiple), the downbeat
    rotation, and the DC removal. A chop (bars=0) gets per-edge snapping instead,
    since there is no bar length to preserve.

    Rotation is per-loop here, not per-shared-window as in `export_take`. That is the
    right call for kits split by family: a loop rotated to its own downbeat starts on
    its 1, so ANY two loops in the pack layer in phase — which is exactly what you
    want when `fred_marea_drums` and `fred_marea_bass` are cut from different moments
    of the track and still have to sit together under `loopAt`.
    """
    sr = ANALYSIS_SR
    mono_ref = y_st[:, 0]
    a = int(cand.start_s * sr)
    if cand.bars > 0:
        # the TRACK tempo, not this window's local mean — see `track_tempo`
        length = L.bar_len_samples(cand.bars, bar_bpm or cand.bpm, sr)
        a, b = L._snap_length_preserving(mono_ref, a, sr, length)
    else:
        b = L._snap_zc(mono_ref, int(cand.end_s * sr), sr)
        a = L._snap_zc(mono_ref, a, sr)
    clip = y_st[a:b]
    if cand.bars > 0 and rotate:
        clip, _, _ = L._rotate_to_downbeat(clip, sr, cand.bars * L.BEATS_PER_BAR)
    if mono_sum and clip.shape[1] > 1:
        clip = clip.mean(axis=1, keepdims=True)
    return clip - clip.mean(axis=0, keepdims=True), sr


def _best_candidates(y_mono: np.ndarray, family: str, grid, *, per_stem: int,
                     bars: tuple) -> dict[str, list]:
    """Run the mode(s) this family calls for, keeping the modes SEPARATE.

    They must not compete on grade, because the rubric scores them on different weights
    and the comparison is rigged. `grade` classes anything under 0.75 s as a one-shot and
    scores it `0.5·level + 0.3·dc + 0.2·zc` — no seam term, no bar term, since a one-shot
    is never looped. Export already removes DC and snaps to zero crossings, so almost
    every chop lands near 1.0, while a loop pays 0.35 on seam and 0.25 on bar consistency
    and rarely maxes both.

    Measured on ME: every chop graded 0.993–0.999 (S), every loop 0.841–0.881 (A). Ranked
    head to head, chops take every slot on a "both" family and the vocal kits come out as
    stabs with no loops in them — precisely backwards. So each mode gets its own slots.
    """
    mode = R.FAMILIES.get(family, R.FAMILIES["tonal"])["mode"]
    out: dict[str, list] = {}
    if mode in ("loops", "both"):
        out["loop"] = L.analyze_stem(y_mono, ANALYSIS_SR, bars=bars, top_n=per_stem,
                                     grid=grid)
    if mode in ("chops", "both"):
        out["chop"] = L.analyze_chops(y_mono, ANALYSIS_SR, top_n=per_stem)
    return out


def _rank_by_grade(by_mode: dict[str, list], y_st: np.ndarray, family: str,
                   keep: int, *, bar_bpm: float = 0.0) -> list:
    """Materialise each candidate exactly as it would be written, then grade THAT.

    This is `analyze_stem`'s `verify` discipline taken one step further: the finder
    re-scores its seam on the post-snap window, and here the whole rubric is applied to
    the post-snap, post-rotation, DC-removed, mono-summed clip — the actual bytes. A
    candidate cannot score well on a signal that never ships.

    Ranking happens WITHIN a mode (see `_best_candidates`); the modes are then
    concatenated so both survive into the kit.
    """
    mono_sum = family == "bass"
    picked = []
    for mode, cands in by_mode.items():
        scored = []
        for c in cands:
            try:
                clip, sr = _slice_export(y_st, c, mono_sum=mono_sum,
                                         bar_bpm=bar_bpm)
            except Exception:
                continue
            if clip.shape[0] < sr // 8:                  # < 125 ms: not a sample
                continue
            g = G.grade_array(clip.T, sr, path=f"<{mode}@{c.start_s}>",
                              role=R.FAMILIES.get(family, {}).get("stem"))
            if c.bars > 0 and _has_dead_bars(clip, sr, c.bars):
                continue          # rejected here rather than reported later — see below
            scored.append((mode, c, g, clip))
        # grade first, finder score as the tiebreak — two clips can both be mechanically
        # perfect, and then musical structure is the only thing left to separate them.
        scored.sort(key=lambda x: (-x[2].grade, -x[1].score))
        # a chop is cheap and a loop is the point: chops get half the slots, minimum one.
        picked += scored[: keep if mode == "loop" else max(1, keep // 2)]
    return picked


def track_tempo(grid) -> tuple[float, float]:
    """One tempo for the whole track, plus how much to trust it. (bpm, confidence).

    `analyze_stem` reports each candidate's bpm as the MEAN of the local per-beat BPM
    inside its own window, which is the right number for detecting drift but the wrong
    one to derive a bar length from: every loop then gets its own idiosyncratic tempo
    and two loops from the same track no longer layer. Measured on ME, the track grid
    read 129.2 bpm while the two shipped loops were cut at 150.5.

    So the bar length comes from a single robust track tempo — the mode of the local-BPM
    histogram, not the mean, because a handful of dropped or doubled beats drag a mean
    and leave a mode alone. Confidence is the share of beats sitting within ±5 % of it,
    which is what tells us whether there is a pulse here at all.
    """
    if grid is None or len(grid[2]) < 8:
        return 0.0, 0.0
    bpm = np.asarray(grid[2], dtype=float)
    bpm = bpm[(bpm > 50) & (bpm < 220)]
    if bpm.size < 8:
        return 0.0, 0.0
    hist, edges = np.histogram(bpm, bins=np.arange(50, 221, 1.0))
    centre = edges[int(np.argmax(hist))] + 0.5
    near = np.abs(bpm - centre) <= centre * 0.05
    if near.sum() >= 4:                       # refine on the beats that agree
        centre = float(np.median(bpm[near]))
        near = np.abs(bpm - centre) <= centre * 0.05
    return round(float(centre), 2), round(float(near.mean()), 3)


TEMPO_TRUST = 0.35          # below this share of on-tempo beats there is no usable grid


def _has_dead_bars(clip: np.ndarray, sr: int, bars: int) -> bool:
    """Does this loop stop partway through? Rejected at export, not merely flagged.

    A four-bar loop whose last two bars are the tail of a fade is a two-bar loop with two
    bars of silence stapled on, and `grade` cannot object to it: silent-to-silent is a
    perfect seam and the length is still an exact bar multiple, so it grades S. Ten of the
    first 161 cuts were exactly this. It is an objective defect with an objective test, so
    it belongs in the gate rather than in a report somebody has to read.
    """
    from .kitcheck import DEAD_BAR_DB, bar_levels
    lv = bar_levels(clip.mean(axis=1) if clip.ndim > 1 else clip, sr, bars)
    return bool(lv) and (max(lv) - min(lv)) > DEAD_BAR_DB


def _pick_grid(stems: list[dict], progress) -> tuple[Optional[tuple], Optional[str]]:
    """One grid for the whole track, from its most rhythmic *verified* stem.

    `find_takes` picks by name (`drums` → `bass` → first). Here the pick is by
    measurement: the highest onset rate among stems that measured as drums, so a
    mislabelled or near-empty drum stem cannot set the grid every other stem is
    cut against.
    """
    rhythmic = [s for s in stems if s["probe"].measured_family == "drums"
                and s["probe"].usable]
    pool = rhythmic or [s for s in stems if s["probe"].usable]
    if not pool:
        return None, None
    lead = max(pool, key=lambda s: s["probe"].onset_rate)
    if progress:
        progress(f"  grid ← {lead['name'].role_token} "
                 f"(onset {lead['probe'].onset_rate:.1f}/s)")
    return L.beat_grid(lead["mono"], ANALYSIS_SR), lead["name"].role_token


def analyze_track(track_dir: Path, *, per_stem: int = 6, keep_per_stem: int = 3,
                  bars: tuple = (1, 2, 4, 8),
                  progress: Optional[Callable[[str], None]] = None) -> dict:
    """Load every stem of one track, verify it, and rank its best cuttable material."""
    wavs = sorted(p for p in track_dir.glob("*.wav"))
    stems = []
    for p in wavs:
        name = R.parse_stem_name(p.name)
        y_st, y_mono, src_sr = _load_44k(p)
        pr = R.probe(y_mono, ANALYSIS_SR, name.claimed_family)
        fam = R.effective_family(name, pr)
        if progress:
            mark = "·" if pr.agrees else "≠"
            skip = "" if pr.usable else "  SKIP"
            progress(f"  {mark} {name.role_token:18.18} {fam:6} "
                     f"rms {pr.rms_dbfs:6.1f}  act {pr.active_frac:4.0%}  "
                     f"cen {pr.centroid_hz:6.0f}Hz  on {pr.onset_rate:4.1f}/s"
                     f"{('  ' + pr.note) if pr.note else ''}{skip}")
        stems.append({"path": p, "name": name, "probe": pr, "family": fam,
                      "st": y_st, "mono": y_mono, "src_sr": src_sr})

    usable = [s for s in stems if s["probe"].usable]
    grid, grid_from = _pick_grid(usable, progress)
    grid_bpm, tempo_conf = track_tempo(grid)
    bar_ok = tempo_conf >= TEMPO_TRUST
    if progress:
        progress(f"  tempo {grid_bpm:.1f} bpm, confidence {tempo_conf:.0%}"
                 + ("" if bar_ok else "  ⇒ NO usable pulse: chops only, no bar-loops"))

    # the producer sometimes declares the tempo in the filename — free ground truth
    declared = next((s["name"].bpm_hint for s in stems if s["name"].bpm_hint), None)
    if declared and grid_bpm:
        err = abs(grid_bpm - declared) / declared * 100
        if progress:
            verdict = "✓" if err < 2 else "⚠"
            progress(f"  {verdict} grid {grid_bpm:.1f} bpm vs declared {declared:.0f} "
                     f"({err:.1f}% off)")

    for s in usable:
        if progress:
            progress(f"  cutting {s['name'].role_token}…")
        cands = _best_candidates(s["mono"], s["family"], grid,
                                 per_stem=per_stem, bars=bars)
        if not bar_ok:
            cands.pop("loop", None)          # a bar-loop off a tempo we do not believe
        s["picks"] = _rank_by_grade(cands, s["st"], s["family"], keep_per_stem,
                                    bar_bpm=grid_bpm if bar_ok else 0.0)

    return {"dir": track_dir, "stems": stems, "usable": usable,
            "grid_bpm": grid_bpm, "tempo_conf": tempo_conf, "bar_ok": bar_ok,
            "grid_from": grid_from, "declared_bpm": declared}


# ── export ───────────────────────────────────────────────────────────────────
def _tidal_token(s: str) -> str:
    """A sample-name token Tidal and the filesystem both like."""
    return re.sub(r"[^a-z0-9]+", "", s.lower())


def export_track(result: dict, *, prefix: str, samples_root: Optional[Path] = None,
                 min_tier: str = MIN_TIER, link: bool = True,
                 progress: Optional[Callable[[str], None]] = None) -> list[CutLoop]:
    """Write one kit per track × family and return what landed.

    Kit = `<prefix>_<track>_<family>`; file = `NN_<producer role>_<bars>b.wav`. The
    bar count is in the NAME because that is the number you need at the keyboard:
    `loopAt 4 $ s "fred_marea_drums" # n 0` is only right if that file really is
    four bars, and `kit_multiples_check` proves it after the fact.
    """
    from . import publish
    root = samples_root or publish.samples_root()
    floor = TIER_ORDER.index(min_tier)
    track_slug = _tidal_token(result["dir"].name)
    by_kit: dict[str, list] = {}
    for s in result["usable"]:
        for mode, cand, g, clip in s.get("picks", []):
            if TIER_ORDER.index(g.tier) < floor:
                continue
            by_kit.setdefault(f"{prefix}_{track_slug}_{s['family']}", []).append(
                (s, mode, cand, g, clip))

    out: list[CutLoop] = []
    for kit, items in sorted(by_kit.items()):
        items.sort(key=lambda t: (-t[3].grade,))
        items = _drop_duplicates(items)
        kit_dir = root / kit
        kit_dir.mkdir(parents=True, exist_ok=True)
        written: list[Path] = []
        # one gain for the kit, computed before anything is written
        head = kit_headroom_gain([float(20 * np.log10(np.max(np.abs(c[4])) + 1e-12))
                                  for c in items])
        for i, (s, mode, cand, g, clip) in enumerate(items):
            role_tok = _tidal_token(s["name"].role_token) or "loop"
            suffix = f"{cand.bars}b" if cand.bars > 0 else "chop"
            dest = kit_dir / f"{i:02d}_{role_tok}_{suffix}.wav"
            clip = clip * head
            if head != 1.0:
                # re-grade the STAGED clip: several sub-scores have a level term, so the
                # grade computed before staging describes audio that is not what shipped.
                # Selection above ran on the pre-staging grades, which is fine — the gain
                # is uniform across the kit, so it cannot reorder its members — but the
                # number recorded in the manifest has to be the one you get by grading
                # the file on disk, or `kitgate`'s cross-check is comparing two different
                # signals and calling the difference a defect.
                g = G.grade_array(clip.T, ANALYSIS_SR, path=str(dest),
                                  role=R.FAMILIES.get(s["family"], {}).get("stem"))
            sf.write(str(dest), clip, ANALYSIS_SR, subtype="PCM_24")
            written.append(dest)
            out.append(CutLoop(
                kit=kit, name=dest.stem, path=str(dest), track=result["dir"].name,
                stem_role=s["name"].role_token, family=s["family"], mode=mode,
                bars=cand.bars,
                bpm=round(result["grid_bpm"] or cand.bpm, 2),
                dur_s=round(clip.shape[0] / ANALYSIS_SR, 4),
                start_s=cand.start_s, finder_score=cand.score,
                tier=g.tier, grade=g.grade, flags=list(g.flags), tags={}))
        # only bar-loops have a bar multiple to be off; a chop is sub-bar by definition,
        # and checking it produces a warning that means nothing.
        bar_files = [w for w, it in zip(written, items) if it[2].bars > 0]
        chk = L.kit_multiples_check(bar_files, result["grid_bpm"] or 120.0)
        if progress:
            bad = "" if chk.get("ok", True) else "  ⚠ off-grid rows"
            hd = "" if head >= 1.0 else f"  −{-20 * np.log10(head):.1f} dB headroom"
            progress(f"  → {kit:34.34} {len(written):2} files{bad}{hd}")
        if link:
            try:
                publish.link_kit(kit, samples=root)
            except Exception as e:                        # never fail a cut on a link
                if progress:
                    progress(f"    (link skipped: {e})")
    return out


HEADROOM_CEILING_DB = -6.0     # kit-wide peak ceiling; see `kit_headroom_gain`


def kit_headroom_gain(peaks_dbfs: list[float], ceiling_db: float = HEADROOM_CEILING_DB
                      ) -> float:
    """ONE attenuation for a whole kit, so its loudest file sits at `ceiling_db`.

    These are stems from a loud commercial master — MAREA's ALL DRUMS and KIT both peak
    at +0.2 dBFS — so a loop cut from them arrives at full scale, which `grade` flags as
    clipping and which SuperDirt cannot use: amp = gain^4, so `# gain 1.2` is +7.6 dB and
    the orbit clips the moment you touch it (`reference_superdirt_gain_is_quartic`). The
    house rule is peaks at -6 dBFS at the source, and let the amplification happen later
    (`feedback_gig_gain_staging`).

    Two things this deliberately does NOT do:

    * **Not per file.** One scalar for the kit. Normalising each file individually would
      make a sparse hat loop as loud as a full drum bus and destroy the balance you would
      otherwise get for free when two loops from the same kit play together.
    * **Never boost.** A kit already quieter than the ceiling is left exactly where it is.
      A quiet stem is quiet on purpose; lifting it would amplify its noise floor and throw
      away real information about the arrangement.
    """
    if not peaks_dbfs:
        return 1.0
    hottest = max(peaks_dbfs)
    if hottest <= ceiling_db:
        return 1.0
    return float(10 ** ((ceiling_db - hottest) / 20.0))


def _drop_duplicates(items: list, thresh: float | None = None) -> list:
    """Remove near-identical cuts from a kit, keeping the best-graded of each group.

    A producer pack overlaps its stems on purpose — ALL DRUMS, KIT, MAIN DRUM LOOP and
    DRUM BREAKS are four views of one groove, and the finder will happily pick the same
    bar out of all four. Eight slots holding three distinct sounds is a worse kit than
    three. `kitcheck` reports duplicates after the fact; this stops them being written,
    reusing the same tempo-agnostic mel fingerprint so the two agree by construction.
    """
    from .kitcheck import DUPE_CORR, _fingerprint
    thresh = DUPE_CORR if thresh is None else thresh   # one definition of "same sound"
    kept, prints = [], []
    for it in items:                                    # already sorted best-first
        clip = it[4]
        fp = _fingerprint(clip.mean(axis=1) if clip.ndim > 1 else clip, ANALYSIS_SR)
        if any(float(fp @ q) > thresh for q in prints):
            continue
        kept.append(it); prints.append(fp)
    return kept


def tag_with_clap(cuts: list[CutLoop], *, progress=None) -> None:
    """Attach CLAP multi-axis tags in place. Optional — a failure never blocks a cut.

    The cheap probe in `roles.py` cannot tell a sung line from a rhodes chord; CLAP
    can, and it is what turns a kit listing into something searchable by vibe rather
    than by the producer's shorthand. It is advisory: `feedback_sample_identity` —
    CLAP proposes, PLN's ear confirms.
    """
    import sys
    tide = Path.home() / "Work/Sound/Tidal/armada/tide-table"
    if not tide.is_dir() or not cuts:
        return
    sys.path.insert(0, str(tide))
    try:
        import sample_semantics as SEM
    except Exception as e:
        if progress:
            progress(f"  (CLAP unavailable: {e})")
        return
    arrays = []
    for c in cuts:
        y, sr = sf.read(c.path, always_2d=True, dtype="float32")
        arrays.append(y.mean(axis=1))
    try:
        emb = np.asarray(SEM.embed_audio_arrays(arrays))
        for c, e in zip(cuts, emb):
            c.tags = SEM.tag_audio_embed(e, topk=2)
    except Exception as e:
        if progress:
            progress(f"  (CLAP tagging failed: {e})")


# ── report ───────────────────────────────────────────────────────────────────
def cheatsheet(cuts: list[CutLoop]) -> str:
    """A paste-ready `.tidal` block: every kit, every index, what it is, how to loop it.

    A kit you cannot address is a kit you will not use. Tidal indexes `n` by the
    sorted filename, which is exactly the `NN_` order written above — so the index
    printed here is the index that plays.
    """
    lines = ["-- fred kits — generated by engine.stempack; n indices are sort order", ""]
    for kit in sorted({c.kit for c in cuts}):
        items = sorted((c for c in cuts if c.kit == kit), key=lambda c: c.name)
        bpms = {round(c.bpm) for c in items if c.bpm}
        lines.append(f"-- {kit}  ({len(items)} samples"
                     + (f", {'/'.join(str(b) for b in sorted(bpms))} bpm)" if bpms else ")"))
        for i, c in enumerate(items):
            inst = (c.tags.get("instrument") or [("", 0)])[0][0]
            how = f"loopAt {c.bars}" if c.bars > 0 else f"one-shot {c.dur_s:.2f}s"
            lines.append(f"--   n {i:<2} {c.name:26.26} {c.tier}  {how:12.12} "
                         f"{c.stem_role}{('  · ' + inst) if inst else ''}")
        ex = items[0]
        if ex.bars > 0:
            lines.append(f'-- d1 $ loopAt {ex.bars} $ chop 8 $ s "{kit}" # n 0')
        else:
            lines.append(f'-- d1 $ s "{kit}*4" # n (irand {len(items)})')
        lines.append("")
    return "\n".join(lines)


def track_summary(r: dict) -> dict:
    """The JSON-safe half of an analysis, so a child process can hand it to the parent.

    The parent regenerates the whole report after CLAP tagging; concatenating the
    children's own report fragments would freeze the tables before the tags exist.
    """
    return {
        "track": r["dir"].name,
        "n_stems": len(r["stems"]),
        "usable": [s["name"].role_token for s in r["stems"] if s["probe"].usable],
        "skipped": [s["name"].role_token for s in r["stems"] if not s["probe"].usable],
        "disagreements": [(s["name"].role_token, s["name"].claimed_family,
                           s["probe"].measured_family)
                          for s in r["stems"] if not s["probe"].agrees],
        "grid_bpm": r["grid_bpm"], "grid_from": r["grid_from"],
        "tempo_conf": r.get("tempo_conf", 0.0), "bar_ok": r.get("bar_ok", True),
        "declared_bpm": r["declared_bpm"],
    }


# CLAP phrases → the family they imply. Keyword matching, not an exhaustive map: the
# point is only to notice a loud disagreement, never to reassign anything.
_CLAP_FAMILY_HINTS = (
    ("drums", ("drum", "hi-hat", "hihat", "percussion", "snare", "kick", "clap",
               "cymbal", "tom", "break")),
    ("bass", ("bass", "sub-bass", "808")),
    ("vox", ("voice", "singing", "vocal", "choir", "chant", "rap", "speech")),
)


def clap_family(tags: dict) -> Optional[str]:
    """The family CLAP's top instrument tag implies, or None if it says nothing useful."""
    inst = (tags or {}).get("instrument") or []
    if not inst:
        return None
    phrase = str(inst[0][0]).lower()
    for fam, keys in _CLAP_FAMILY_HINTS:
        if any(k in phrase for k in keys):
            return fam
    return "tonal"


def clap_disagreements(cuts: list[CutLoop]) -> list[tuple]:
    """Cuts whose CLAP tag contradicts the kit they landed in. Advisory, never applied.

    The cheap probe cannot separate a sung line from a rhodes chord, and for a stem whose
    filename token we do not recognise there is no producer label to fall back on either —
    so those land wherever the coarse measurement pointed. CLAP is the one thing here that
    can tell the difference, and it runs after the cut, so the honest use of it is to
    surface the cases worth a second listen rather than to silently re-file them
    (`feedback_sample_identity`: CLAP proposes, PLN's ear confirms).
    """
    out = []
    for c in cuts:
        cf = clap_family(c.tags)
        if not cf:
            continue
        # compare at the STEM level, not the family: a `hits` kit IS percussion, so CLAP
        # hearing "a kick drum" in it agrees. And `fx`/`tonal`/`vox` all map to `other`
        # and are genuinely not separable, so they never contradict each other.
        want = R.FAMILIES.get(c.family, {}).get("stem")
        got = R.FAMILIES.get(cf, {}).get("stem")
        if want == got:
            continue
        soft = {"fx", "tonal", "vox"}
        if c.family in soft and cf in soft:
            continue
        inst = c.tags["instrument"][0]
        out.append((c.kit, c.name, c.stem_role, c.family, cf, inst[0], inst[1]))
    return out


def report_md(results: list[dict], cuts: list[CutLoop]) -> str:
    from collections import Counter
    lines = ["# Fred stem-pack cut — report", ""]
    for r in results:
        lines += [f"## {r['track']}", "",
                  f"- stems: {r['n_stems']} ({len(r['usable'])} usable"
                  + (f", skipped near-silent: {', '.join(r['skipped'])}" if r["skipped"] else "") + ")",
                  f"- grid: **{r['grid_bpm']:.1f} bpm** from `{r['grid_from']}` "
                  f"· confidence {r['tempo_conf']:.0%}"
                  + ("" if r["bar_ok"] else " · **no usable pulse — chops only**")
                  + (f" · producer declared **{r['declared_bpm']:.0f}**"
                     f" ({abs(r['grid_bpm'] - r['declared_bpm']) / r['declared_bpm'] * 100:.1f}% off)"
                     if r["declared_bpm"] else " · no declared tempo to check against")]
        if r["disagreements"]:
            lines.append("- label vs measurement disagreements (measurement wins):")
            lines += [f"  - `{tok}` claimed *{c}*, measured **{m}**"
                      for tok, c, m in r["disagreements"]]
        mine = [c for c in cuts if c.track == r["track"]]
        tiers = Counter(c.tier for c in mine)
        lines += ["", f"- shipped {len(mine)} samples · tiers "
                  + " ".join(f"{t}×{tiers[t]}" for t in reversed(TIER_ORDER) if tiers[t]), ""]
        for kit in sorted({c.kit for c in mine}):
            ks = sorted((c for c in mine if c.kit == kit), key=lambda c: c.name)
            lines.append(f"### `{kit}` — {len(ks)}")
            lines.append("")
            lines.append("| n | file | tier | grade | bars | bpm | dur | from | CLAP |")
            lines.append("|--:|---|---|--:|--:|--:|--:|---|---|")
            for i, c in enumerate(ks):
                tg = "; ".join(f"{v[0][0]}" for k, v in c.tags.items()
                               if v and k in ("instrument", "timbre")) if c.tags else ""
                lines.append(
                    f"| {i} | `{c.name}` | {c.tier} | {c.grade:.3f} | "
                    f"{c.bars or '—'} | {c.bpm:.0f} | {c.dur_s:.2f}s | "
                    f"`{c.stem_role}` | {tg} |")
            lines.append("")

    dis = clap_disagreements(cuts)
    if dis:
        lines += ["## Worth a second listen", "",
                  "CLAP's instrument tag contradicts the kit these landed in. Nothing is "
                  "re-filed on this basis — CLAP is the one signal here that can separate "
                  "a sung line from a keyboard, so it is used to say what to audition "
                  "first, not to decide.", "",
                  "| kit | file | from stem | filed as | CLAP hears | conf |",
                  "|---|---|---|---|---|--:|"]
        for kit, name, role, fam, cf, phrase, conf in dis:
            lines.append(f"| `{kit}` | `{name}` | `{role}` | {fam} | {cf} "
                         f"({phrase}) | {conf:.2f} |")
        lines.append("")
    return "\n".join(lines)


def _run_track(track_dir: Path, opts: dict) -> tuple[dict, list, str]:
    """Analyze + export one track. Runs in a worker process, so it returns text.

    Progress is buffered and returned rather than printed: interleaving nine tracks'
    per-stem lines across processes produces a log nobody can read. One block per
    track, printed when it lands.
    """
    buf: list[str] = [f"\u25b6 {track_dir.name}"]
    say = buf.append
    r = analyze_track(track_dir, per_stem=opts["per_stem"], keep_per_stem=opts["keep"],
                      bars=opts["bars"], progress=say)
    cuts: list[CutLoop] = []
    if not opts["dry_run"]:
        cuts = export_track(r, prefix=opts["prefix"], samples_root=opts["samples_root"],
                            min_tier=opts["min_tier"], link=opts["link"], progress=say)
    # drop the audio before this crosses a process boundary
    r["stems"] = [{k: v for k, v in s.items() if k not in ("st", "mono", "picks")}
                  for s in r["stems"]]
    r["usable"] = [s for s in r["stems"] if s["probe"].usable]
    return r, cuts, "\n".join(buf)


# ── CLI ──────────────────────────────────────────────────────────────────────
def main(argv=None) -> int:
    import argparse
    ap = argparse.ArgumentParser(prog="stempack", description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("pack", type=Path, nargs="?",
                    help="folder of per-track stem folders (omit with --report-only)")
    ap.add_argument("--tracks", nargs="*", default=None, help="track folders (default: all)")
    ap.add_argument("--prefix", default="fred", help="kit name prefix")
    ap.add_argument("--per-stem", type=int, default=6, help="candidates generated per stem")
    ap.add_argument("--keep", type=int, default=3, help="candidates kept per stem")
    ap.add_argument("--bars", type=int, nargs="+", default=[1, 2, 4, 8])
    ap.add_argument("--min-tier", default=MIN_TIER, choices=TIER_ORDER)
    ap.add_argument("--samples-root", type=Path, default=None)
    ap.add_argument("--out", type=Path, default=None, help="write report/cheatsheet here")
    ap.add_argument("--no-clap", action="store_true")
    ap.add_argument("--no-link", action="store_true")
    ap.add_argument("--dry-run", action="store_true", help="analyze and report, write nothing")
    ap.add_argument("--report-only", action="store_true",
                    help="rebuild the report/cheatsheet from an existing --out, no cutting")
    ap.add_argument("-j", "--jobs", type=int, default=3,
                    help="tracks analyzed in parallel (each pins one core)")
    a = ap.parse_args(argv)

    if a.report_only:
        return _report_only(a.out or Path.cwd())

    names = a.tracks or sorted(d.name for d in a.pack.iterdir() if d.is_dir())
    todo = []
    for nm in names:
        d = a.pack / nm
        if d.is_dir():
            todo.append(d)
        else:
            print(f"! no such track folder: {d}", flush=True)

    opts = dict(per_stem=a.per_stem, keep=a.keep, bars=tuple(a.bars), prefix=a.prefix,
                samples_root=a.samples_root, min_tier=a.min_tier,
                link=not a.no_link, dry_run=a.dry_run)
    out = a.out or Path.cwd()
    out.mkdir(parents=True, exist_ok=True)

    results, cuts = [], []
    if a.jobs == 1 or len(todo) == 1:
        for d in todo:
            r, c, log = _run_track(d, opts)
            print(log, flush=True)
            results.append(track_summary(r)); cuts += c
        (out / "tracks.json").write_text(json.dumps(results, indent=2))
    else:
        # A track is the unit of parallelism: every stem in a track is cut against one
        # shared beat grid, so stems cannot be split across workers, and tracks are
        # fully independent. Fan out as SUBPROCESSES rather than a ProcessPoolExecutor
        # — the pool's forkserver is broken under this Python (BrokenProcessPool on a
        # trivial task), and one process per track also means a track that blows up
        # loses only itself.
        results, cuts = _fan_out(todo, a, out)

    if cuts and not a.no_clap:
        say = lambda m: print(m, flush=True)                                # noqa: E731
        say("\nCLAP tagging…")
        tag_with_clap(cuts, progress=say)      # once, in the parent: one model load

    (out / "fred_kits.md").write_text(report_md(results, cuts))
    (out / "fred_kits.tidal").write_text(cheatsheet(cuts))
    (out / "fred_kits.json").write_text(json.dumps([asdict(c) for c in cuts], indent=2))
    print(f"\n{len(cuts)} samples in {len({c.kit for c in cuts})} kits "
          f"\u2192 {out}/fred_kits.{{md,tidal,json}}", flush=True)
    return 0


def _report_only(out: Path) -> int:
    """Rebuild `fred_kits.{md,tidal}` from the manifest already in `out`.

    The manifest is the record; the report and the cheatsheet are views of it. Anything
    that legitimately changes the audio after a cut — `restage.py` staging a kit's level
    and re-grading it — leaves those views describing numbers that are no longer true, and
    re-cutting an entire pack to refresh a markdown table would be absurd. This also picks
    up report sections added since the cut was made.
    """
    cuts_p = out / "fred_kits.json"
    if not cuts_p.exists():
        raise SystemExit(f"no manifest at {cuts_p}")
    cuts = [CutLoop(**c) for c in json.loads(cuts_p.read_text())]
    results = []
    tj = out / "tracks.json"
    if tj.exists():
        results = json.loads(tj.read_text())
    else:                                  # fan-out runs leave them per child
        for f in sorted((out / "_tracks").glob("*/tracks.json")):
            results += json.loads(f.read_text())
    (out / "fred_kits.md").write_text(report_md(results, cuts))
    (out / "fred_kits.tidal").write_text(cheatsheet(cuts))
    print(f"rebuilt report + cheatsheet for {len(cuts)} samples in "
          f"{len({c.kit for c in cuts})} kits from {cuts_p.name}")
    return 0


def _fan_out(todo: list[Path], a, out: Path) -> tuple[list, list]:
    """Run one `--jobs 1` subprocess per track, then merge their results.

    Children are told `--no-clap`: tagging loads a ~500 MB model, so it happens once
    in the parent over the merged set rather than once per track.
    """
    import subprocess
    from concurrent.futures import ThreadPoolExecutor
    scratch = out / "_tracks"
    scratch.mkdir(parents=True, exist_ok=True)

    def run(d: Path):
        sub = scratch / _tidal_token(d.name)
        cmd = [__import__("sys").executable, "-u", "-m", "engine.stempack", str(d.parent),
               "--tracks", d.name, "--jobs", "1", "--prefix", a.prefix,
               "--per-stem", str(a.per_stem), "--keep", str(a.keep),
               "--bars", *[str(b) for b in a.bars], "--min-tier", a.min_tier,
               "--out", str(sub), "--no-clap"]
        if a.samples_root:
            cmd += ["--samples-root", str(a.samples_root)]
        if a.no_link:
            cmd.append("--no-link")
        if a.dry_run:
            cmd.append("--dry-run")
        # tee to a per-track log rather than only capturing: a 30-minute analysis whose
        # output only appears at exit is a run you cannot supervise.
        sub.mkdir(parents=True, exist_ok=True)
        with open(sub / "run.log", "w") as lf:
            r = subprocess.run(cmd, cwd=Path(__file__).resolve().parent.parent,
                               stdout=lf, stderr=subprocess.STDOUT, text=True)
        return d, sub, r

    results, cuts = [], []
    with ThreadPoolExecutor(max_workers=a.jobs) as ex:   # threads only wait on children
        for d, sub, r in ex.map(run, todo):
            log = (sub / "run.log").read_text() if (sub / "run.log").exists() else ""
            print(log.rstrip(), flush=True)
            if r.returncode != 0:
                print(f"! {d.name} failed (exit {r.returncode})", flush=True)
                continue
            j = sub / "fred_kits.json"
            if j.exists():
                cuts += [CutLoop(**c) for c in json.loads(j.read_text())]
            s = sub / "tracks.json"
            if s.exists():
                results += json.loads(s.read_text())
    return results, cuts


if __name__ == "__main__":
    raise SystemExit(main())
