#!/usr/bin/env python3
"""kitindex — build the sample-kit index the auditioner SPA browses.

One JSON describing every kit under the sample root: what is in it, how to loop it,
and a peak envelope small enough to draw without downloading the audio. Written to
`armada/ui/public/kits.json`, fetched by `kits.html`, and served by both
`vite.config.ts` and `armada/serve.py` through the `samples` mount — so the URLs in
here work identically in dev and on the LAN.

Two sources, merged, and the merge is the point:

* **The folder is the truth about what exists.** SuperDirt enumerates the directory,
  so the index has to as well, via `pvbanks.playable` (the same listing kitgate and
  bank-check use). A kit with no manifest — 163 of PLN's kits predate the Foundry —
  still has to be browsable, or the tool only works on what it built itself.
* **A Foundry manifest is the truth about what a file MEANS.** tier, grade, family,
  producer role, bars, bpm and CLAP tags all come from `fred_kits.json`-style cuts;
  nothing here re-derives them, so the SPA and the gate cannot disagree.

`n` comes from the folder listing, never from the manifest order, because `n` is
what SuperDirt will play. (kitgate proves the two agree; this file does not assume it.)

Kit dirs do not all sit at the same depth — a pack can be staged in its own subtree
(`Samples/Fred/output/fred_marea_drums/`) while the older kits sit directly under
`Samples/`. So `--root` is repeatable and says *where to look for kit folders*, while
`--mount-root` is the one directory the SPA's `samples` mount points at and the only
thing URLs are computed from. Keeping those two ideas apart is what lets a pack move
without every URL in the index changing.

    python3 kitindex.py --root ~/Work/Sound/Samples/Fred/output --cuts <cuts.json>
    python3 kitindex.py --all                    # every kit under the sample root
"""
from __future__ import annotations

import argparse
import fnmatch
import json
import sys
import time
from pathlib import Path

import numpy as np
import soundfile as sf

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
sys.path.insert(0, str(HERE.parent))            # tools/, for pvbanks

PEAKS = 128          # envelope columns per sample — enough for a 200 px row
OUT = HERE.parent.parent / "armada/ui/public/kits.json"
MOUNT = "samples"    # the /audio/<prefix> mount in armada/ui/audio-mounts.json
CACHE = HERE / ".kitindex-cache.json"   # envelope cache, keyed on size+mtime


def envelope(path: Path, cols: int = PEAKS) -> tuple[list[int], float, float, int, int]:
    """A tiny peak envelope, plus the facts the player needs before decoding.

    Read as one block: these files are seconds long, and a per-column seek would be
    slower than the read. Mono-summed — the SPA draws one lane per sample, and a
    two-lane waveform at 200 px wide is noise.
    """
    y, sr = sf.read(str(path), dtype="float32", always_2d=True)
    n, ch = y.shape
    mono = y.mean(axis=1)
    if n == 0:
        return [0] * cols, 0.0, 0.0, sr, ch
    # pad to a whole number of columns so the last one is not a short-window artefact
    per = max(1, int(np.ceil(n / cols)))
    pad = per * cols - n
    blocks = np.abs(np.pad(mono, (0, pad))).reshape(cols, per)
    peak = float(np.max(np.abs(mono)))
    env = blocks.max(axis=1)
    scaled = (env / (peak + 1e-12) * 100.0).round().astype(int).tolist()
    rms = float(np.sqrt(np.mean(mono ** 2)))
    return scaled, peak, rms, sr, ch


def db(x: float) -> float:
    return round(float(20 * np.log10(max(x, 1e-12))), 1)


def load_cuts(paths: list[Path]) -> dict[str, dict]:
    """Manifest rows keyed by absolute audio path."""
    by_path: dict[str, dict] = {}
    for p in paths:
        for c in json.loads(p.read_text()):
            by_path[str(Path(c["path"]).resolve())] = c
    return by_path


def _load_cache() -> dict:
    """Envelopes keyed on `path|size|mtime_ns`.

    Reading 1884 files to draw 1884 little waveforms takes 3.5 minutes, and the whole
    point of re-running this is that ONE pack changed. The key includes size and mtime
    so a restaged kit — same path, new bytes — misses the cache: a stale envelope drawn
    over fresh audio would be a lie that looks like data.
    """
    try:
        return json.loads(CACHE.read_text())
    except Exception:
        return {}


def build(roots: list[Path], mount_root: Path, cuts: dict[str, dict],
          only: list[str], *, use_cache: bool = True, progress=None) -> dict:
    from pvbanks import playable
    cache = _load_cache() if use_cache else {}
    fresh: dict[str, list] = {}
    hits = 0

    def url_for(f: Path) -> str:
        """`/audio/samples/<path under the mount>`. A kit outside the mount would be
        unplayable in the browser, so say so here rather than serve a 404."""
        try:
            return f"/audio/{MOUNT}/{f.resolve().relative_to(mount_root.resolve())}"
        except ValueError:
            return ""

    kits = []
    seen: set[str] = set()
    dirs = [x for r in roots if r.is_dir() for x in sorted(r.iterdir()) if x.is_dir()]
    for d in dirs:
        if only and not any(fnmatch.fnmatch(d.name, pat) for pat in only):
            continue
        files = playable(d)
        if not files:
            continue
        if d.name in seen:
            # Two dirs claiming one bank name is exactly the shadowing SuperDirt does
            # silently; the index must not pick a winner behind PLN's back.
            if progress:
                progress(f"  ! {d.name}: a second directory claims this bank name — skipped")
            continue
        seen.add(d.name)
        samples = []
        for n, f in enumerate(files):          # n IS the folder index SuperDirt uses
            c = cuts.get(str(f.resolve()), {})
            st = f.stat()
            key = f"{f}|{st.st_size}|{st.st_mtime_ns}"
            cached = cache.get(key)
            if cached:
                env, peak, rms, sr, ch = cached
                hits += 1
            else:
                try:
                    env, peak, rms, sr, ch = envelope(f)
                except Exception as e:         # a kit with one unreadable file stays usable
                    if progress:
                        progress(f"    ! {f.name}: {e}")
                    continue
            fresh[key] = [env, peak, rms, sr, ch]
            samples.append({
                "n": n, "name": f.stem, "file": f.name,
                "url": url_for(f),
                "sr": sr, "ch": ch, "peak_dbfs": db(peak), "rms_dbfs": db(rms),
                "env": env,
                # from the manifest when the Foundry cut it, absent otherwise
                "bars": c.get("bars"), "bpm": c.get("bpm"),
                # where in the source track this was cut from: three loops off one stem
                # are otherwise indistinguishable in a list, both named for the stem
                "start_s": c.get("start_s"),
                "dur_s": c.get("dur_s"), "tier": c.get("tier"),
                "grade": c.get("grade"), "family": c.get("family"),
                "stem_role": c.get("stem_role"), "track": c.get("track"),
                "mode": c.get("mode"), "flags": c.get("flags") or [],
                "tags": c.get("tags") or {},
            })
        if not samples:
            continue
        known = [s for s in samples if s["bpm"]]
        kits.append({
            "kit": d.name, "n_samples": len(samples),
            "source": "foundry" if known else "legacy",
            "track": next((s["track"] for s in samples if s["track"]), None),
            "family": next((s["family"] for s in samples if s["family"]), None),
            "bpm": known[0]["bpm"] if known else None,
            "samples": samples,
        })
        if progress:
            progress(f"  {d.name:38.38} {len(samples):3} samples")
    if use_cache:
        try:
            CACHE.write_text(json.dumps(fresh, separators=(",", ":")))
        except OSError:
            pass
    if progress and hits:
        progress(f"  ({hits} envelopes from cache)")
    return {"generated": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
            "roots": [str(r) for r in roots], "mount_root": str(mount_root),
            "peaks": PEAKS, "kits": kits}


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(prog="kitindex", description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--cuts", type=Path, action="append", default=[],
                    help="a stempack cuts manifest; repeatable")
    ap.add_argument("--only", action="append", default=[],
                    help="glob on kit name; repeatable (default: every kit)")
    ap.add_argument("--all", action="store_true", help="explicit form of no --only")
    ap.add_argument("--root", type=Path, action="append", default=[],
                    help="a directory holding kit folders; repeatable "
                         "(default: the sample root itself)")
    ap.add_argument("--mount-root", type=Path, default=None,
                    help="what /audio/samples points at (default: the sample root)")
    ap.add_argument("--out", type=Path, default=OUT)
    ap.add_argument("--no-cache", action="store_true",
                    help="re-read every file instead of reusing cached envelopes")
    a = ap.parse_args(argv)

    from engine import publish
    mount_root = a.mount_root or publish.samples_root()
    roots = a.root or [publish.samples_root()]
    cuts = load_cuts(a.cuts)
    print(f"kitindex — {', '.join(str(r) for r in roots)}"
          + (f", {len(cuts)} manifest rows from {len(a.cuts)} file(s)" if cuts else ""))
    idx = build(roots, mount_root, cuts, [] if a.all else a.only,
                use_cache=not a.no_cache, progress=print)
    a.out.parent.mkdir(parents=True, exist_ok=True)
    body = json.dumps(idx, separators=(",", ":"))
    a.out.write_text(body)
    # `npm run build` copies public/ into dist/, so a regenerated index would otherwise
    # be invisible to the built page that `serve.py` hands the LAN — the same shape as
    # the stale preload that debuted at a venue. Write both, or neither.
    dist = a.out.parent.parent / "dist" / a.out.name
    if dist.parent.is_dir():
        dist.write_text(body)
        print(f"  (also updated {dist})")
    n = sum(k["n_samples"] for k in idx["kits"])
    legacy = sum(1 for k in idx["kits"] if k["source"] == "legacy")
    unreachable = [k["kit"] for k in idx["kits"]
                   if any(not s["url"] for s in k["samples"])]
    if unreachable:
        print(f"  ! {len(unreachable)} kits are outside {mount_root} and cannot be "
              f"served: {', '.join(unreachable[:4])}")
    print(f"\n{len(idx['kits'])} kits ({legacy} without a manifest), {n} samples "
          f"→ {a.out}  ({a.out.stat().st_size / 1024:.0f} KB)")
    return 0


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