#!/usr/bin/env python3
"""restage — apply kit-wide gain staging to sample kits that are already on disk.

`engine.stempack` stages a kit's level as it exports, but kits cut before that (or by
hand, in Audacity) carry whatever level their source had. Producer stems come off a loud
commercial master and land at full scale, which SuperDirt cannot use: amp = gain^4, so
`# gain 1.2` is +7.6 dB and the orbit clips the moment you touch it.

This applies the SAME function `stempack` uses at export (`kit_headroom_gain` — one gain
per kit, attenuate only), so a kit restaged here is bit-identical to one cut with the
fix in place. Reads and rewrites in place, 24-bit, and reports what it changed.

    python3 restage.py <kit dir> [<kit dir> …]        # explicit kits
    python3 restage.py --cuts fred_kits.json          # every kit named in a cut manifest
    python3 restage.py --all-under ~/Work/Sound/Samples --dry-run
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np
import soundfile as sf

sys.path.insert(0, str(Path(__file__).resolve().parent))
from engine.stempack import HEADROOM_CEILING_DB, kit_headroom_gain  # noqa: E402


def _peak_db(path: Path) -> float:
    y, _ = sf.read(str(path), always_2d=True, dtype="float32")
    return float(20 * np.log10(np.max(np.abs(y)) + 1e-12))


def restage_kit(kit_dir: Path, *, ceiling_db: float = HEADROOM_CEILING_DB,
                dry_run: bool = False) -> dict:
    """Bring one kit's hottest file to the ceiling, moving every file by the same gain."""
    wavs = sorted(p for p in kit_dir.glob("*.wav") if not p.name.startswith("."))
    if not wavs:
        return {"kit": kit_dir.name, "files": 0, "gain_db": 0.0, "changed": False}
    peaks = [_peak_db(p) for p in wavs]
    g = kit_headroom_gain(peaks, ceiling_db)
    gain_db = 20 * np.log10(g) if g > 0 else 0.0
    if g >= 1.0:
        return {"kit": kit_dir.name, "files": len(wavs), "gain_db": 0.0,
                "changed": False, "peak_before": max(peaks)}
    if not dry_run:
        for p in wavs:
            y, sr = sf.read(str(p), always_2d=True, dtype="float32")
            sf.write(str(p), y * g, sr, subtype="PCM_24")
    return {"kit": kit_dir.name, "files": len(wavs), "gain_db": round(float(gain_db), 2),
            "changed": True, "peak_before": round(max(peaks), 2),
            "peak_after": round(max(peaks) + float(gain_db), 2)}


def _regrade_manifest(cuts_path: Path, kits: set[str]) -> None:
    """Re-grade the restaged files and write the new grades back into the manifest."""
    from engine import grade as G
    cuts = json.loads(cuts_path.read_text())
    n = 0
    for c in cuts:
        if c["kit"] not in kits or not Path(c["path"]).exists():
            continue
        g = G.grade(c["path"], role=c.get("family"))
        c["grade"], c["tier"] = g.grade, g.tier
        c["flags"] = list(g.flags)
        n += 1
    cuts_path.write_text(json.dumps(cuts, indent=2))
    print(f"  re-graded {n} entries in {cuts_path.name}")


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(prog="restage", description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("kits", nargs="*", type=Path)
    ap.add_argument("--cuts", type=Path, help="a stempack cuts JSON; restage every kit in it")
    ap.add_argument("--all-under", type=Path, help="every immediate subdirectory of this path")
    ap.add_argument("--ceiling", type=float, default=HEADROOM_CEILING_DB)
    ap.add_argument("--dry-run", action="store_true")
    a = ap.parse_args(argv)

    dirs: list[Path] = list(a.kits)
    if a.cuts:
        cuts = json.loads(a.cuts.read_text())
        dirs += sorted({Path(c["path"]).parent for c in cuts})
    if a.all_under:
        dirs += sorted(d for d in a.all_under.iterdir() if d.is_dir())
    dirs = sorted({d.resolve() for d in dirs if d.is_dir()})
    if not dirs:
        ap.error("no kit directories given")

    rows = [restage_kit(d, ceiling_db=a.ceiling, dry_run=a.dry_run) for d in dirs]
    changed = [r for r in rows if r["changed"]]

    # Whoever changes the audio owns the record of it. Restaging moves every grade that
    # has a level term in it, so a manifest left alone would describe bytes that no longer
    # exist — and `kitgate` compares the two precisely to catch that kind of drift.
    if a.cuts and changed and not a.dry_run:
        _regrade_manifest(a.cuts, {r["kit"] for r in changed})
    for r in changed:
        print(f"  {r['kit']:34.34} {r['files']:3} files  "
              f"{r['peak_before']:+6.1f} → {r['peak_after']:+6.1f} dBFS "
              f"({r['gain_db']:+.1f} dB)")
    verb = "would restage" if a.dry_run else "restaged"
    print(f"{verb} {len(changed)} of {len(rows)} kits "
          f"(ceiling {a.ceiling:+.1f} dBFS; kits already below it are left alone)")
    return 0


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