"""kitcheck — the post-export audit: what `grade` cannot see.

`grade.py` answers *is this file mechanically a clean loop* — seam, zero crossings,
DC, level, bar self-consistency. Every one of those can be perfect on a sample that
is musically useless, and this module exists for exactly the failures that survive a
clean grade. It is the substitute for an ear when nobody is listening yet.

Three checks, each closing a hole the rubric has by construction:

**1 · dead bars.** A four-bar loop whose last two bars are the tail of a fade grades
   S: the seam is silent-to-silent, the length is exact. It is still a two-bar loop
   with two bars of nothing after it. Per-bar RMS, flagged when any bar sits far
   below the loudest.

**2 · bar periodicity.** The finder cuts on a beat grid, so its windows are always
   *bar-aligned*; nothing checks they are a *bar-length musical unit*. Onset-envelope
   autocorrelation at the half/whole-bar lag says whether the material actually
   repeats at the claimed rate, or whether four bars of through-composed movement got
   labelled a loop.

**3 · near-duplicates.** A producer pack ships overlapping stems on purpose — ALL
   DRUMS, KIT, MAIN DRUM LOOP and DRUM BREAKS are four views of one groove. Cut the
   best window from each and you ship the same bar four times under four names. A kit
   whose 8 slots hold 3 distinct sounds is worse than a kit of 3.

Plus the two cheap invariants worth asserting rather than assuming: every bar-loop's
duration is an exact bar multiple at its own BPM, and every file is 44.1 kHz.

    python3 -m engine.kitcheck fred_kits.json
"""
from __future__ import annotations

import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

import numpy as np
import soundfile as sf

BAR_TOL_MS = 1.0          # a bar multiple is exact or it is not
DEAD_BAR_DB = 18.0        # a bar this far under the loudest is not carrying material
DUPE_CORR = 0.90          # normalised cross-correlation above this ⇒ same sound


@dataclass
class KitCheck:
    path: str
    name: str
    kit: str
    ok: bool = True
    problems: list = field(default_factory=list)
    bar_dev_ms: Optional[float] = None
    bar_rms_db: list = field(default_factory=list)
    periodicity: Optional[float] = None
    dupe_of: Optional[str] = None


def _mono(y):
    return y.mean(axis=1) if y.ndim > 1 else y


def bar_levels(y: np.ndarray, sr: int, bars: int) -> list[float]:
    """RMS in dBFS of each bar of the loop."""
    if bars <= 0:
        return []
    n = len(y) // bars
    return [round(float(20 * np.log10(np.sqrt(np.mean(y[i * n:(i + 1) * n] ** 2)) + 1e-12)), 1)
            for i in range(bars)]


def periodicity(y: np.ndarray, sr: int, bars: int) -> float:
    """How strongly the onset envelope repeats at the loop's own bar period (0..1).

    Autocorrelation of the onset strength envelope, read at the lag of one bar (and
    of half the loop, whichever the loop is long enough to have). A rhythmic loop
    peaks there; a through-composed window does not.
    """
    import librosa
    if bars < 2:
        return 1.0                       # a 1-bar loop has no internal period to check
    oenv = librosa.onset.onset_strength(y=y, sr=sr, hop_length=512)
    oenv = oenv - oenv.mean()
    if not np.any(oenv) or len(oenv) < 8:
        return 0.0
    ac = np.correlate(oenv, oenv, mode="full")[len(oenv) - 1:]
    ac = ac / (ac[0] + 1e-12)
    per_bar = len(oenv) / bars
    best = 0.0
    for k in (1, 2):                      # one bar, and two bars
        lag = int(round(per_bar * k))
        if 0 < lag < len(ac):
            best = max(best, float(ac[max(0, lag - 2):lag + 3].max()))
    return round(max(0.0, min(1.0, best)), 3)


def _fingerprint(y: np.ndarray, sr: int, n: int = 256) -> np.ndarray:
    """A short, tempo-agnostic signature: log-mel energy resampled to n columns.

    Two cuts of the same groove at different lengths must still compare equal, so the
    envelope is normalised in time (resampled to a fixed width) before comparison —
    a raw waveform correlation would call a 2-bar and a 4-bar cut of one break
    unrelated.
    """
    import librosa
    m = librosa.feature.melspectrogram(y=y, sr=sr, n_mels=24, hop_length=512)
    m = librosa.power_to_db(m + 1e-12)
    if m.shape[1] < 2:
        return np.zeros(24 * n)
    idx = np.linspace(0, m.shape[1] - 1, n)
    m = np.vstack([np.interp(idx, np.arange(m.shape[1]), row) for row in m])
    m = m - m.mean()
    return (m / (np.linalg.norm(m) + 1e-12)).ravel()


def check_cuts(cuts: list[dict]) -> list[KitCheck]:
    """Audit every exported sample, then cross-check for duplicates within each kit."""
    out: list[KitCheck] = []
    prints: dict[str, list[tuple[str, np.ndarray]]] = {}

    for c in cuts:
        r = KitCheck(path=c["path"], name=c["name"], kit=c["kit"])
        p = Path(c["path"])
        if not p.exists():
            r.ok = False; r.problems.append("missing file"); out.append(r); continue
        y, sr = sf.read(str(p), always_2d=True, dtype="float32")
        m = _mono(y)

        if sr != 44100:
            r.ok = False; r.problems.append(f"sample rate {sr}, expected 44100")

        bars, bpm = c["bars"], c["bpm"]
        if bars > 0 and bpm > 0:
            want = bars * 4 * 60.0 / bpm
            dev = (len(m) / sr - want) * 1000
            r.bar_dev_ms = round(dev, 3)
            if abs(dev) > BAR_TOL_MS:
                r.ok = False
                r.problems.append(f"not a bar multiple: {dev:+.1f} ms off {bars} bars @ {bpm:.1f}")

            r.bar_rms_db = bar_levels(m, sr, bars)
            if r.bar_rms_db:
                lo, hi = min(r.bar_rms_db), max(r.bar_rms_db)
                if hi - lo > DEAD_BAR_DB:
                    r.ok = False
                    dead = [i + 1 for i, v in enumerate(r.bar_rms_db) if hi - v > DEAD_BAR_DB]
                    r.problems.append(f"dead bar(s) {dead}: {hi - lo:.0f} dB spread across bars")

            r.periodicity = periodicity(m, sr, bars)
            if r.periodicity < 0.15:
                r.problems.append(f"weak bar periodicity ({r.periodicity:.2f}) — "
                                  "bar-aligned but maybe not a repeating unit")

        prints.setdefault(c["kit"], []).append((c["name"], _fingerprint(m, sr)))
        out.append(r)

    by_name = {r.name: r for r in out}
    for kit, items in prints.items():
        for i in range(len(items)):
            for j in range(i + 1, len(items)):
                corr = float(items[i][1] @ items[j][1])
                if corr > DUPE_CORR:
                    later = by_name[items[j][0]]
                    if later.dupe_of is None:
                        later.dupe_of = items[i][0]
                        later.problems.append(
                            f"near-duplicate of {items[i][0]} in the same kit (r={corr:.2f})")
    return out


def report(checks: list[KitCheck]) -> str:
    bad = [c for c in checks if c.problems]
    lines = [f"# kit check — {len(checks)} samples, "
             f"{len(checks) - len(bad)} clean, {len(bad)} flagged", ""]
    if not bad:
        lines.append("No problems found.")
    for kit in sorted({c.kit for c in bad}):
        lines.append(f"## `{kit}`")
        for c in [x for x in bad if x.kit == kit]:
            mark = "✗" if not c.ok else "⚠"
            lines.append(f"- {mark} `{c.name}` — " + "; ".join(c.problems))
        lines.append("")
    dupes = [c for c in checks if c.dupe_of]
    lines += ["", f"**Distinct sounds:** {len(checks) - len(dupes)} of {len(checks)} "
              f"({len(dupes)} near-duplicates)"]
    return "\n".join(lines)


def main(argv=None) -> int:
    import argparse
    ap = argparse.ArgumentParser(prog="kitcheck", description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("cuts_json", type=Path, help="fred_kits.json from engine.stempack")
    ap.add_argument("--out", type=Path, default=None)
    a = ap.parse_args(argv)
    cuts = json.loads(a.cuts_json.read_text())
    checks = check_cuts(cuts)
    txt = report(checks)
    print(txt)
    if a.out:
        a.out.write_text(txt)
        (a.out.parent / (a.out.stem + ".json")).write_text(
            json.dumps([c.__dict__ for c in checks], indent=2))
    return 0 if all(c.ok for c in checks) else 1


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