"""Pre-flight checks for the acoustic-space dataset. Run before every training run.

Every check here exists because something slipped past a weaker one:

  grid        a source missing from one space silently shrinks that space's
              coverage; the manifest still builds and training still runs
  duration    a failed rsync left a file TRUNCATED, not missing. Presence
              checks passed. Only a duration sweep caught it, and preprocessing
              silently skipped the item (71 of 72) while reporting exit 0
  integrity   the same failure again, 2026-08-10, in the place this script did
              not look: an interrupted write left a TARGET holding 0.54s of
              audio under a header still declaring 6.000s. Byte-for-byte it was
              9% present. Every header-based reader called it healthy; only the
              VAE refused it, months of "mysterious VAE skip" later. Bytes are
              the only honest witness, and the BUILT dataset needs checking,
              not just the slices -- the corruption was written by the copy step
  format      mixed sample rate / channel count reaches the VAE as garbage
  level       if one source is much louder than the rest, the model can learn
              loudness as a cue for space instead of learning the space
  depth       how far each render sits from its dry source. Catches a space
              that was rendered near-dry by mistake -- the failure mode that
              would teach contradictory things under one caption

Exit code is non-zero if anything fails, so it can gate a pipeline.

  python verify_dataset.py <sliced_dir> [--expect-seconds 6.0] [--manifest path]
                                        [--dataset-dir <.../audio>]

Pass --dataset-dir to also byte-check the built references/ and targets/. Do it
on every run: the slices can be perfect while the copy of them is not.
"""

import argparse
import subprocess
import sys
import wave
from pathlib import Path

import numpy as np

# Spaces that add an environmental bed rather than pure reverb. Level variation
# across sources is DELIBERATE here (it stops the model memorising one exact
# ambient recording), so depth outliers are not flagged as errors.
AMBIENCE_SPACES = {"outdoor_day_birds", "outdoor_night"}


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


def probe_format(p: Path) -> tuple:
    r = subprocess.run(
        ["ffprobe", "-v", "error", "-select_streams", "a",
         "-show_entries", "stream=sample_rate,channels", "-of", "csv=p=0", str(p)],
        capture_output=True, text=True,
    )
    parts = r.stdout.strip().split(",")
    try:
        return int(parts[0]), int(parts[1])
    except (IndexError, ValueError):
        return (0, 0)


def probe_integrity(p: Path) -> tuple:
    """Compare the frame count a WAV declares against the audio data really there.

    A truncated WAV keeps its original header, so anything reading length from
    metadata -- wave.getnframes(), soundfile.info(), most "duration" fields --
    reports the ORIGINAL length for a file that is mostly gone. ffprobe's
    format=duration is the exception, deriving duration from real size.

    Returns (ok, real_bytes, expected_bytes, real_seconds).
    """
    try:
        with wave.open(str(p), "rb") as w:
            n, ch, sw, sr = w.getnframes(), w.getnchannels(), w.getsampwidth(), w.getframerate()
            real = len(w.readframes(n))
    except Exception:
        return False, 0, 0, 0.0
    expected = n * ch * sw
    bps = ch * sw * sr
    # tolerate a padding byte, not a missing chunk
    return real >= expected - 4, real, expected, (real / bps if bps else 0.0)


def load_mono(p: Path, sr: int = 24000) -> np.ndarray:
    tmp = Path(f"/tmp/_vd_{p.parent.name[:12]}_{p.stem[:16]}.wav")
    subprocess.run(
        ["ffmpeg", "-y", "-v", "error", "-i", str(p), "-ac", "1", "-ar", str(sr), str(tmp)],
        capture_output=True,
    )
    try:
        with wave.open(str(tmp)) as w:
            return np.frombuffer(w.readframes(w.getnframes()), dtype=np.int16).astype(float) / 32768.0
    except Exception:
        return np.array([])


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("sliced_dir")
    ap.add_argument("--expect-seconds", type=float, default=6.0)
    ap.add_argument("--tolerance", type=float, default=0.05)
    ap.add_argument("--manifest", default=None)
    ap.add_argument("--dataset-dir", default=None,
                    help="built dataset audio/ dir; byte-checks references/ and targets/ too")
    ap.add_argument("--allow-longer", action="store_true",
                    help="permit renders longer than expected (cathedral tails); still flags SHORT files")
    args = ap.parse_args()

    D = Path(args.sliced_dir)
    if not (D / "dry").is_dir():
        print(f"no dry/ directory under {D}")
        return 2

    spaces = sorted(d.name for d in D.iterdir() if d.is_dir() and d.name != "dry")
    sources = sorted(p.stem for p in (D / "dry").glob("*.wav"))
    fails = []

    print(f"sources: {len(sources)}   spaces: {len(spaces)}")
    print()

    # --- 1. grid completeness -------------------------------------------------
    missing = [(s, sp) for s in sources for sp in spaces if not (D / sp / f"{s}.wav").exists()]
    print(f"[grid]     {'FAIL' if missing else 'ok  '}  {len(sources)*len(spaces)} expected, {len(missing)} missing")
    for s, sp in missing[:10]:
        print(f"             missing: {sp}/{s}.wav")
    if missing:
        fails.append("grid")

    # --- 2. duration ----------------------------------------------------------
    bad_dur = []
    for p in sorted(D.rglob("*.wav")):
        d = probe_duration(p)
        short = d < args.expect_seconds - args.tolerance
        long_ = d > args.expect_seconds + args.tolerance
        if short or (long_ and not args.allow_longer):
            bad_dur.append((p.relative_to(D), d))
    print(f"[duration] {'FAIL' if bad_dur else 'ok  '}  expecting {args.expect_seconds}s"
          f"{' (longer allowed)' if args.allow_longer else ''}, {len(bad_dur)} wrong")
    for rel, d in bad_dur[:10]:
        print(f"             {rel}  {d:.2f}s")
    if bad_dur:
        fails.append("duration")

    # --- 3. byte-level integrity ----------------------------------------------
    # The duration check above would catch truncation, but only where it is
    # pointed, and it was only ever pointed at the slices. The one corruption
    # this project has had was written into the BUILT dataset by the copy step,
    # which nothing verified at all.
    roots = [("sliced", D)]
    if args.dataset_dir:
        ds = Path(args.dataset_dir)
        for sub in ("references", "targets"):
            if (ds / sub).is_dir():
                roots.append((sub, ds / sub))
            else:
                print(f"[integrity] warning: no {sub}/ under {ds}")

    truncated, n_checked = [], 0
    for label, root in roots:
        for p in sorted(root.rglob("*.wav")):
            n_checked += 1
            ok, real, expected, secs = probe_integrity(p)
            if not ok:
                truncated.append((label, p.name, real, expected, secs))
    print(f"[integrity] {'FAIL' if truncated else 'ok  '}  {n_checked} files byte-checked"
          f"{'' if args.dataset_dir else ' (slices only -- pass --dataset-dir)'}"
          f", {len(truncated)} truncated")
    for label, name, real, expected, secs in truncated[:10]:
        pct = 100.0 * real / expected if expected else 0.0
        print(f"             {label}/{name}")
        print(f"               {real:,} of {expected:,} bytes ({pct:.1f}%) = {secs:.2f}s of real audio")
    if truncated:
        fails.append("integrity")

    # --- 4. format ------------------------------------------------------------
    fmts = {}
    for p in sorted(D.rglob("*.wav")):
        fmts.setdefault(probe_format(p), []).append(p.relative_to(D))
    print(f"[format]   {'FAIL' if len(fmts) > 1 else 'ok  '}  {len(fmts)} distinct (sample_rate, channels)")
    for f, ps in fmts.items():
        print(f"             {f}: {len(ps)} files" + (f"  e.g. {ps[0]}" if len(fmts) > 1 else ""))
    if len(fmts) > 1:
        fails.append("format")

    # --- 5. dry source levels -------------------------------------------------
    levels = {}
    for s in sources:
        x = load_mono(D / "dry" / f"{s}.wav")
        if x.size:
            levels[s] = 20 * np.log10(np.sqrt((x ** 2).mean()) + 1e-12)
    if levels:
        lo, hi = min(levels.values()), max(levels.values())
        spread = hi - lo
        # a wide spread is fine when it is transient material (claps): check peaks too
        print(f"[level]    ok    dry RMS spread {spread:.1f} dB  "
              f"({min(levels, key=levels.get)} {lo:.1f} .. {max(levels, key=levels.get)} {hi:.1f})")
        if spread > 15:
            print("             note: >15 dB spread. Fine for sparse transients (claps sit low on RMS,")
            print("             high on peak); worth checking if it is a genuinely quiet recording.")

    # --- 6. processing depth --------------------------------------------------
    print("[depth]    rel_diff = RMS(wet-dry)/RMS(dry) per space, across sources")
    for sp in spaces:
        vals = []
        for s in sources:
            wp, dp = D / sp / f"{s}.wav", D / "dry" / f"{s}.wav"
            if not wp.exists():
                continue
            w, d = load_mono(wp), load_mono(dp)
            if not w.size or not d.size:
                continue
            n = min(len(w), len(d))
            vals.append(float(np.sqrt(((w[:n] - d[:n]) ** 2).mean()) / (np.sqrt((d[:n] ** 2).mean()) + 1e-12)))
        if not vals:
            continue
        med, lo, hi = float(np.median(vals)), min(vals), max(vals)
        note = ""
        if sp in AMBIENCE_SPACES:
            note = "  (ambience bed - level variation is deliberate)"
        elif lo < 0.25 * med:
            note = "  <- OUTLIER: some source rendered near-dry"
            fails.append(f"depth:{sp}")
        print(f"             {sp:24} median {med:.3f}  range {lo:.3f}-{hi:.3f}{note}")

    # --- 7. manifest ----------------------------------------------------------
    # Two manifest shapes reach this script: the PAIRS grid (id / reference_audio
    # / audio / split) and the SLICES report (output / source_id / environment).
    # Reading pair columns off a slices.csv raised KeyError and exited 1 *after*
    # every check had already passed -- a green run reported as a failure.
    if args.manifest:
        import csv
        rows = list(csv.DictReader(open(args.manifest)))
        base = Path(args.manifest).parent
        cols = set(rows[0]) if rows else set()
        print()
        if {"reference_audio", "audio"} <= cols:
            path_cols, id_col = ("reference_audio", "audio"), "id"
        elif "output" in cols:
            path_cols, id_col = ("output",), "output"
        else:
            path_cols, id_col = (), ""
            print(f"[manifest] skip  {len(rows)} rows, unrecognised columns: {sorted(cols)}")
        if path_cols:
            broken = [r.get(id_col, "?") for r in rows for k in path_cols
                      if not (base / r[k]).resolve().exists()]
            n_train = sum(1 for r in rows if r.get("split") == "train")
            print(f"[manifest] {'FAIL' if broken else 'ok  '}  {len(rows)} rows, {n_train} train, "
                  f"{len(broken)} broken paths, {n_train // max(len(spaces),1)} train sources/space")
            for b in broken[:6]:
                print(f"             broken: {b}")
            if broken:
                fails.append("manifest")

    print()
    if fails:
        print(f"FAILED: {', '.join(sorted(set(fails)))}")
        return 1
    print("all checks passed - safe to preprocess")
    return 0


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