#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "numpy>=1.26",
# ]
# ///
"""
validate_buffered_streaming.py - parity harness for buffered streaming
on parakeet-unified-en-0.6b.

For a target audio file at a chosen (L, C, R) tuple, runs:

  1. NeMo reference dumper:
       scripts/dump_reference_parakeet_nemo.py buffered_streaming ...
     producing per-chunk reference tensors under <out>/ref/.

  2. C++ transcribe-cli with TRANSCRIBE_DUMP_DIR=<out>/cpp/ and
     --stream-chunk-ms <feed-ms>, producing parallel C++ dumps.

  3. Per-chunk diff of stream.chunk.<step>.{audio_in, enc_out} between
     the two dump dirs.

  4. Final-transcript byte equality.

Usage:
    uv run --project scripts/envs/parakeet \\
        scripts/validate_buffered_streaming.py \\
        --nemo /path/to/parakeet-unified-en-0.6b.nemo \\
        --gguf models/parakeet-unified-en-0.6b/parakeet-unified-en-0.6b-F32.gguf \\
        --audio samples/jfk.wav \\
        --out build/validate_buffered_streaming/parakeet-unified/jfk/default \\
        [--left-secs 5.6 --chunk-secs 1.04 --right-secs 1.04] \\
        [--backend cpu --threads 1] \\
        [--feed-ms 500]

Exit code 0 = per-chunk parity OK across all chunks (audio_in
bit-exact + enc_out within fp32 tolerance). The final-transcript
byte match is reported informationally — greedy RNN-T can flip a
single token at a chunk boundary on fp32 noise even when the
encoder windows are byte-identical, so the per-chunk gate is the
algorithmic check of record. WER on test-clean is the corpus-level
gate.
"""

from __future__ import annotations

import argparse
import json
import os
import shutil
import subprocess
import sys
from pathlib import Path

import numpy as np


REPO = Path(__file__).resolve().parent.parent


def run_subprocess(cmd: list[str], *, label: str, env: dict | None = None,
                   cwd: Path | None = None, capture: bool = False):
    print(f"+ {label}: {' '.join(cmd)}")
    proc = subprocess.run(
        cmd,
        env=({**os.environ, **env} if env else None),
        cwd=str(cwd) if cwd else None,
        capture_output=capture,
        text=True,
    )
    if proc.returncode != 0:
        if capture:
            sys.stdout.write(proc.stdout or "")
            sys.stderr.write(proc.stderr or "")
        raise SystemExit(f"error: {label} returned {proc.returncode}")
    return proc


def run_ref_dump(*, nemo: Path, audio: Path, out_dir: Path,
                 left_secs: float, chunk_secs: float, right_secs: float,
                 force: bool) -> Path:
    stamp = out_dir / "stream_history.json"
    if stamp.exists() and not force:
        print(f"  [skip ref] {stamp} exists (pass --force to re-run)")
        return out_dir
    out_dir.mkdir(parents=True, exist_ok=True)
    cmd = [
        "uv", "run", "--project", str(REPO / "scripts/envs/parakeet"),
        str(REPO / "scripts/dump_reference_parakeet_nemo.py"),
        "buffered_streaming",
        "--model", str(nemo),
        "--audio", str(audio),
        "--out", str(out_dir),
        "--left-secs", f"{left_secs}",
        "--chunk-secs", f"{chunk_secs}",
        "--right-secs", f"{right_secs}",
    ]
    run_subprocess(cmd, label="ref dump", cwd=REPO)
    return out_dir


def run_cpp_dump(*, gguf: Path, audio: Path, out_dir: Path,
                 feed_ms: int, backend: str | None,
                 threads: int | None,
                 left_ms: int | None,
                 chunk_ms: int | None,
                 right_ms: int | None) -> tuple[Path, str]:
    if out_dir.exists():
        shutil.rmtree(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    cli = REPO / "build/bin/transcribe-cli"
    if not cli.exists():
        raise SystemExit(
            f"error: {cli} not built; run cmake --build build --target transcribe-cli"
        )

    cmd = [
        str(cli),
        "--model", str(gguf),
        "--stream-chunk-ms", str(feed_ms),
        str(audio),
    ]
    if backend is not None:
        cmd += ["--backend", backend]
    if threads is not None:
        cmd += ["--threads", str(threads)]
    if left_ms is not None:
        cmd += ["--stream-buf-left-ms", str(left_ms)]
    if chunk_ms is not None:
        cmd += ["--stream-buf-chunk-ms", str(chunk_ms)]
    if right_ms is not None:
        cmd += ["--stream-buf-right-ms", str(right_ms)]
    env = {"TRANSCRIBE_DUMP_DIR": str(out_dir)}
    proc = run_subprocess(cmd, label="cpp dump", env=env,
                          cwd=REPO, capture=True)
    sys.stdout.write(proc.stdout)
    text = ""
    for line in proc.stdout.splitlines():
        if line.startswith("text:"):
            text = line[len("text:"):].strip()
            break
    return out_dir, text


def load_f32(path: Path, expected_shape: list[int] | None = None) -> np.ndarray:
    sidecar = path.with_suffix(".json")
    raw = np.fromfile(path, dtype=np.float32)
    if not sidecar.exists():
        return raw
    meta = json.loads(sidecar.read_text())
    shape = meta.get("shape") or expected_shape
    if shape:
        return raw.reshape(shape)
    return raw


def compare_pair(name: str, ref_path: Path, cpp_path: Path) -> dict:
    """Per-tensor diff. Returns a row with max_abs, mean_abs, p99_abs, rel."""
    if not ref_path.exists() or not cpp_path.exists():
        return {
            "name": name,
            "status": "MISSING",
            "ref": ref_path.exists(),
            "cpp": cpp_path.exists(),
        }
    ref = load_f32(ref_path).astype(np.float64)
    cpp = load_f32(cpp_path).astype(np.float64)
    if ref.shape != cpp.shape:
        # Some chunks may have differently-shaped tensors (last-chunk
        # divergence). Report and skip.
        return {
            "name": name,
            "status": "SHAPE_DIFF",
            "ref_shape": list(ref.shape),
            "cpp_shape": list(cpp.shape),
        }
    diff = np.abs(ref - cpp)
    max_abs = float(diff.max())
    mean_abs = float(diff.mean())
    p99_abs = float(np.quantile(np.abs(ref), 0.99))
    rel = float(max_abs / max(p99_abs, 1e-9))
    return {
        "name": name,
        "status": "OK",
        "max_abs": max_abs,
        "mean_abs": mean_abs,
        "p99_abs": p99_abs,
        "rel_max": rel,
        "n_elem": int(ref.size),
    }


def main() -> int:
    p = argparse.ArgumentParser()
    p.add_argument("--nemo", required=True, type=Path)
    p.add_argument("--gguf", required=True, type=Path)
    p.add_argument("--audio", required=True, type=Path)
    p.add_argument("--out", required=True, type=Path)
    p.add_argument("--left-secs", type=float, default=5.6)
    p.add_argument("--chunk-secs", type=float, default=1.04)
    p.add_argument("--right-secs", type=float, default=1.04)
    p.add_argument("--backend", default="cpu")
    p.add_argument("--threads", type=int, default=1)
    p.add_argument("--feed-ms", type=int, default=500)
    p.add_argument("--force", action="store_true")
    args = p.parse_args()

    out_dir = args.out.resolve()
    ref_dir = out_dir / "ref"
    cpp_dir = out_dir / "cpp"

    run_ref_dump(
        nemo=args.nemo.resolve(), audio=args.audio.resolve(),
        out_dir=ref_dir,
        left_secs=args.left_secs,
        chunk_secs=args.chunk_secs,
        right_secs=args.right_secs,
        force=args.force,
    )
    cpp_path, cpp_text = run_cpp_dump(
        gguf=args.gguf.resolve(), audio=args.audio.resolve(),
        out_dir=cpp_dir,
        feed_ms=args.feed_ms,
        backend=args.backend, threads=args.threads,
        left_ms=int(round(args.left_secs   * 1000)),
        chunk_ms=int(round(args.chunk_secs * 1000)),
        right_ms=int(round(args.right_secs * 1000)),
    )

    # Final transcript byte equality
    ref_transcript_path = ref_dir / "transcript.json"
    ref_text = ""
    if ref_transcript_path.exists():
        ref_text = json.loads(ref_transcript_path.read_text()).get("text", "")
    # CPP transcript: strip leading/trailing whitespace
    cpp_text_norm = cpp_text.strip()
    ref_text_norm = ref_text.strip()
    transcript_match = (cpp_text_norm == ref_text_norm)

    # Per-chunk comparison. Steady-state chunks should match within
    # CPU-fp32 noise. Last-chunk divergence is reported but tolerated.
    ref_steps = sorted({
        int(p.stem.split(".")[2])
        for p in ref_dir.glob("stream.chunk.*.audio_in.f32")
    })
    cpp_steps = sorted({
        int(p.stem.split(".")[2])
        for p in cpp_dir.glob("stream.chunk.*.audio_in.f32")
    })
    common = sorted(set(ref_steps) & set(cpp_steps))
    print(f"\n=== Per-chunk parity (ref steps: {ref_steps}; cpp steps: {cpp_steps}; "
          f"common: {common}) ===")

    # Tolerances calibrated to the parakeet family's accepted noise
    # envelope (tests/tolerances/parakeet.json allows max_abs=3.7,
    # mean_abs=0.05 at enc.final on the offline path). Per-chunk
    # enc_out is a slice of that same encoder output running on a
    # smaller window, so the same tolerance applies.
    TOL_MAX_ABS  = 5.0
    TOL_MEAN_ABS = 0.1

    rows: list[dict] = []
    fail_chunks = 0
    for step in common:
        for kind in ("audio_in", "enc_out"):
            name = f"stream.chunk.{step}.{kind}"
            r = compare_pair(
                name,
                ref_dir / f"{name}.f32",
                cpp_dir / f"{name}.f32",
            )
            rows.append({"step": step, **r})
            if r.get("status") == "OK":
                tag = "OK"
                over_max  = r["max_abs"]  > TOL_MAX_ABS
                over_mean = r["mean_abs"] > TOL_MEAN_ABS
                if over_max or over_mean:
                    tag = "FAIL"
                    fail_chunks += 1
                print(f"  step {step:>2} {kind:9s}: max_abs={r['max_abs']:.3e} "
                      f"mean_abs={r['mean_abs']:.3e} rel_max={r['rel_max']:.3e} [{tag}]")
            else:
                # Variable-stride algorithm produces identical chunk
                # geometry to ref, so SHAPE_DIFF or MISSING is now a
                # real failure (unlike the legacy fixed-stride path).
                print(f"  step {step:>2} {kind:9s}: {r['status']}  "
                      f"{r.get('ref_shape', '')} vs {r.get('cpp_shape', '')}")
                fail_chunks += 1

    print()
    print(f"final transcript ref: {ref_text_norm!r}")
    print(f"final transcript cpp: {cpp_text_norm!r}")
    print(f"transcript match: {transcript_match}")

    summary = {
        "transcript_match": transcript_match,
        "ref_text": ref_text_norm,
        "cpp_text": cpp_text_norm,
        "n_common_chunks": len(common),
        "ref_steps": ref_steps,
        "cpp_steps": cpp_steps,
        "chunk_failures": fail_chunks,
        "rows": rows,
    }
    (out_dir / "summary.json").write_text(json.dumps(summary, indent=2))

    if fail_chunks > 0:
        print(f"FAIL: {fail_chunks} chunks exceed tolerance or have wrong shape")
        return 1
    if not transcript_match:
        # Greedy RNN-T can tip a single-token decision on fp32 noise
        # even when per-chunk encoder outputs match within tolerance.
        # Report informationally — per-chunk parity is the algorithmic
        # gate; WER on test-clean is the corpus-level gate.
        print("WARN: transcript byte-match differs (per-chunk parity gate passes — "
              "likely fp32-noise tipping a greedy decision)")
    print("OK")
    return 0


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