"""
scripts/validate_final_runtime.py

FINAL RUNTIME VALIDATION — NO MODEL CHANGES.

Validates the frozen architecture end-to-end in a normal, network-enabled
environment:

    16kHz mono audio -> Whisper Tiny frozen encoder -> mean pooling
    -> Logistic Regression -> P(END) -> threshold/debounce -> END/CONTINUE

This script does not modify the model, the classifier, the thresholds, or
the feature representation. It only loads and runs what already exists
(models/whisper_classifier.joblib, src/turn_detector/inference.py) and
reports real, measured numbers.

Usage:
    python scripts/validate_final_runtime.py
    python scripts/validate_final_runtime.py --audio path/to/clip.wav

If no --audio is given, it looks for a real sample under
data/raw/phase3_sample/audio/ (this project's own real-audio validation
clips) and uses the first one found.

Environment note: this script requires torch, transformers, and network
access to Hugging Face (to load openai/whisper-tiny) — none of which are
available in the sandbox this project was otherwise developed in. If run
there, it will fail with a clear, real error message rather than faking a
result (see the STOP CONDITION section at the bottom of this file). It is
otherwise ready to run as-is in Google Colab or a Hugging Face Space.
"""

from __future__ import annotations

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

PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT / "src"))

import numpy as np

PINNED_SKLEARN_VERSION = "1.6.1"  # must match requirements.txt and the version
                                    # models/whisper_classifier.joblib was pickled with


def check_sklearn_version() -> None:
    """Do NOT suppress the InconsistentVersionWarning by hiding it — the
    correct fix is running with the pinned dependency version. This check
    surfaces a clear, loud message if that pin isn't actually satisfied at
    runtime, rather than letting the warning fire and continuing quietly.
    """
    import sklearn
    installed = sklearn.__version__
    print(f"scikit-learn version installed: {installed}")
    print(f"scikit-learn version pinned in requirements.txt: {PINNED_SKLEARN_VERSION}")
    if installed != PINNED_SKLEARN_VERSION:
        print(
            f"WARNING: installed scikit-learn ({installed}) does not match the pinned "
            f"version ({PINNED_SKLEARN_VERSION}) that models/whisper_classifier.joblib "
            f"was pickled with. This is the actual root cause of any "
            f"InconsistentVersionWarning you may see below — the fix is "
            f"`pip install scikit-learn=={PINNED_SKLEARN_VERSION}`, not suppressing the "
            f"warning. Continuing anyway so you can see the real behavior, but results "
            f"should not be trusted until this is fixed.",
            file=sys.stderr,
        )
    else:
        print("scikit-learn version matches the pin exactly — no version-mismatch risk.")


def find_real_audio_sample(explicit_path: str | None) -> Path:
    if explicit_path:
        p = Path(explicit_path)
        if not p.exists():
            raise FileNotFoundError(f"--audio path does not exist: {p}")
        return p

    candidates_dir = PROJECT_ROOT / "data" / "raw" / "phase3_sample" / "audio"
    if candidates_dir.exists():
        wavs = sorted(candidates_dir.glob("*.wav"))
        if wavs:
            print(f"No --audio given; using real sample from this project's own "
                  f"Phase 3 validation set: {wavs[0].name}")
            return wavs[0]

    raise FileNotFoundError(
        "No --audio path given and no real WAV samples found under "
        "data/raw/phase3_sample/audio/. Pass --audio path/to/clip.wav explicitly."
    )


def load_wav_as_float32(path: Path) -> tuple[np.ndarray, int]:
    """Load a WAV file without requiring soundfile (uses the stdlib `wave`
    module, which handles PCM WAV natively) — real audio, not synthesized.
    """
    import wave

    with wave.open(str(path), "rb") as wf:
        sr = wf.getframerate()
        n_channels = wf.getnchannels()
        sampwidth = wf.getsampwidth()
        n_frames = wf.getnframes()
        raw = wf.readframes(n_frames)

    if sampwidth == 2:
        audio = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0
    elif sampwidth == 4:
        audio = np.frombuffer(raw, dtype=np.int32).astype(np.float32) / 2147483648.0
    elif sampwidth == 1:
        audio = (np.frombuffer(raw, dtype=np.uint8).astype(np.float32) - 128.0) / 128.0
    else:
        raise ValueError(f"Unsupported WAV sample width: {sampwidth} bytes")

    if n_channels > 1:
        audio = audio.reshape(-1, n_channels).mean(axis=1)

    return audio, sr


def time_call(fn, *args, **kwargs):
    t0 = time.perf_counter()
    result = fn(*args, **kwargs)
    t1 = time.perf_counter()
    return result, (t1 - t0) * 1000.0


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--audio", type=str, default=None, help="Path to a real WAV file. "
                     "Defaults to a real sample from data/raw/phase3_sample/audio/.")
    args = ap.parse_args()

    print("=" * 70)
    print("FINAL RUNTIME VALIDATION — architecture is frozen, not modified here")
    print("=" * 70)

    # --- dependency check (sklearn version pin) ---
    try:
        check_sklearn_version()
    except ImportError as e:
        print(f"FATAL: scikit-learn not installed: {e}", file=sys.stderr)
        sys.exit(1)

    # --- locate real audio ---
    try:
        audio_path = find_real_audio_sample(args.audio)
    except FileNotFoundError as e:
        print(f"FATAL: {e}", file=sys.stderr)
        sys.exit(1)

    audio, sr = load_wav_as_float32(audio_path)
    print(f"\nLoaded real audio: {audio_path}")
    print(f"  duration: {len(audio)/sr:.2f}s, sample_rate: {sr}, samples: {len(audio)}")

    # --- attempt to import torch/transformers and construct TurnDetector ---
    print("\n" + "-" * 70)
    print("Loading TurnDetector (classifier is loaded eagerly; Whisper is lazy)...")
    print("-" * 70)

    try:
        from turn_detector.inference import TurnDetector, TurnDetectorConfig, InferenceError
    except Exception as e:
        print(f"FATAL: could not import TurnDetector: {e}", file=sys.stderr)
        sys.exit(1)

    try:
        detector = TurnDetector(TurnDetectorConfig())
    except InferenceError as e:
        print(f"\nFATAL: could not construct TurnDetector (classifier stage failed): {e}", file=sys.stderr)
        sys.exit(1)

    print("Classifier loaded successfully.")
    print("Attempting to load openai/whisper-tiny (requires torch, transformers, and "
          "network access to Hugging Face)...")

    t_load_start = time.perf_counter()
    try:
        detector.load_whisper()
    except InferenceError as e:
        # ---------------------------------------------------------------
        # STOP CONDITION: do not fake a result. Report the exact blocker.
        # ---------------------------------------------------------------
        print("\n" + "=" * 70)
        print("STOPPED — Whisper could not be loaded in this environment.")
        print("=" * 70)
        print(f"\nExact error:\n{e}\n")
        print(
            "This script is otherwise complete and ready to run as-is in an environment "
            "with torch, transformers, and network access to huggingface.co — e.g.:\n"
            "  - Google Colab (Runtime -> Run all, after `pip install -r requirements.txt`)\n"
            "  - A Hugging Face Space (torch/transformers are part of the standard runtime)\n"
            "  - Any local/cloud machine with `pip install -r requirements.txt` and network access\n\n"
            "No result below this point is fabricated or simulated — the script exits here."
        )
        sys.exit(2)
    t_load_end = time.perf_counter()
    whisper_load_time_ms = (t_load_end - t_load_start) * 1000.0
    print(f"Whisper Tiny loaded in {whisper_load_time_ms:.1f}ms (device: {detector._device})")

    # =========================================================================
    # Run 1: COLD (Whisper just loaded above, but this is the first predict()
    # call — first-call framework/kernel warmup effects, e.g. CUDA context
    # init, still apply even though the model itself is already loaded)
    # =========================================================================
    print("\n" + "-" * 70)
    print("Run 1 (cold-ish: first predict() call after model load)")
    print("-" * 70)
    result_cold, total_cold_ms = time_call(detector.predict, audio, sr)
    print(json.dumps(result_cold, indent=2))
    print(f"Measured wall-clock for this call: {total_cold_ms:.2f}ms")

    # =========================================================================
    # Run 2: WARM (model fully warmed up, second call)
    # =========================================================================
    print("\n" + "-" * 70)
    print("Run 2 (warm: second predict() call, same audio)")
    print("-" * 70)
    result_warm, total_warm_ms = time_call(detector.predict, audio, sr)
    print(json.dumps(result_warm, indent=2))
    print(f"Measured wall-clock for this call: {total_warm_ms:.2f}ms")

    print(f"\nCold call latency: {total_cold_ms:.2f}ms")
    print(f"Warm call latency: {total_warm_ms:.2f}ms")
    print(f"Cold - warm delta: {total_cold_ms - total_warm_ms:+.2f}ms")

    # =========================================================================
    # Per-stage latency breakdown (preprocessing / encoder / classifier),
    # measured directly against the SAME loaded detector instance — reusing
    # its already-loaded feature_extractor/whisper_model/classifier, not
    # reloading anything, and not changing the model itself.
    # =========================================================================
    print("\n" + "-" * 70)
    print("Per-stage latency breakdown (warm, this loaded detector instance)")
    print("-" * 70)

    import torch

    processed_audio = np.asarray(audio, dtype=np.float32)

    _, preprocessing_ms = time_call(
        detector._feature_extractor, processed_audio, sampling_rate=sr, return_tensors="pt"
    )

    inputs = detector._feature_extractor(processed_audio, sampling_rate=sr, return_tensors="pt")
    input_features = inputs["input_features"].to(detector._device)

    def encoder_forward():
        with torch.no_grad():
            return detector._whisper_model.encoder(input_features)

    encoder_out, encoder_ms = time_call(encoder_forward)
    pooled = encoder_out.last_hidden_state.mean(dim=1).squeeze(0).cpu().numpy().reshape(1, -1)

    _, classifier_ms = time_call(detector._classifier.predict_proba, pooled)

    total_staged_ms = preprocessing_ms + encoder_ms + classifier_ms

    print(f"Preprocessing (feature extraction): {preprocessing_ms:.2f}ms")
    print(f"Whisper encoder forward pass:       {encoder_ms:.2f}ms")
    print(f"Classifier (predict_proba):         {classifier_ms:.2f}ms")
    print(f"Sum of stages:                      {total_staged_ms:.2f}ms")
    print(f"(compare to measured warm predict() call: {total_warm_ms:.2f}ms — should be in the same ballpark; "
          f"small differences are validation-overhead/measurement-boundary noise, not a bug by themselves)")

    # =========================================================================
    # Verification checks (Step 8)
    # =========================================================================
    print("\n" + "-" * 70)
    print("Verification checks")
    print("-" * 70)

    checks = []

    p_sum = result_warm["end_probability"] + result_warm["continue_probability"]
    checks.append(("probabilities sum to 1", abs(p_sum - 1.0) < 1e-6, f"sum={p_sum}"))

    checks.append(("decision is END or CONTINUE", result_warm["decision"] in ("END", "CONTINUE"),
                    f"decision={result_warm['decision']}"))

    all_numeric = [result_warm["end_probability"], result_warm["continue_probability"], result_warm["latency_ms"]]
    no_nan_inf = all(np.isfinite(v) for v in all_numeric)
    checks.append(("no NaN/Inf in output", no_nan_inf, f"values={all_numeric}"))

    expected_schema = {"decision", "end_probability", "continue_probability", "latency_ms"}
    actual_schema = set(result_warm.keys())
    checks.append(("output schema matches expected keys exactly", actual_schema == expected_schema,
                    f"expected={expected_schema}, actual={actual_schema}"))

    all_passed = True
    for name, passed, detail in checks:
        status = "PASS" if passed else "FAIL"
        if not passed:
            all_passed = False
        print(f"  [{status}] {name} ({detail})")

    print("\n" + "=" * 70)
    if all_passed:
        print("ALL VERIFICATION CHECKS PASSED.")
    else:
        print("ONE OR MORE VERIFICATION CHECKS FAILED — see above. Not overridden or hidden.")
    print("=" * 70)

    summary = {
        "audio_file": str(audio_path),
        "duration_sec": len(audio) / sr,
        "sklearn_version_installed": __import__("sklearn").__version__,
        "sklearn_version_pinned": PINNED_SKLEARN_VERSION,
        "device": detector._device,
        "whisper_load_time_ms": whisper_load_time_ms,
        "cold_call_latency_ms": total_cold_ms,
        "warm_call_latency_ms": total_warm_ms,
        "cold_minus_warm_delta_ms": total_cold_ms - total_warm_ms,
        "staged_latency_ms": {
            "preprocessing": preprocessing_ms,
            "encoder": encoder_ms,
            "classifier": classifier_ms,
            "sum": total_staged_ms,
        },
        "result_warm": result_warm,
        "verification_checks": [{"name": n, "passed": bool(p), "detail": d} for n, p, d in checks],
        "all_checks_passed": all_passed,
    }
    print("\nFull summary (JSON):")
    print(json.dumps(summary, indent=2))

    sys.exit(0 if all_passed else 1)


if __name__ == "__main__":
    main()