#!/usr/bin/env python3
"""
benchmark_v3.py - Quick WER-free benchmark comparing Apple vs Whisper on file(s)

- For each audio file: run Apple (--bench) and optionally Whisper
- Measures wall time, RTF, segment counts
- For known reference texts, computes naive WER (word-level)
- Outputs table + JSON

Usage:
  python benchmark_v3.py /tmp/long.wav
  python benchmark_v3.py prayer.json --extract-audio  # future
  python benchmark_v3.py --test-say  # generates TTS samples + benches
"""

import subprocess
import json
import time
import os
import sys
import tempfile
from pathlib import Path

BIN = Path(__file__).parent / "apple_speech" / ".build" / "release" / "apple-speech-transcribe"

def run_apple_bench(wav_path: str, locale: str = "en-US"):
    cmd = [str(BIN), "--bench", wav_path, "--locale", locale]
    start = time.time()
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=20)
    wall = time.time() - start
    if result.returncode != 0:
        return {"ok": False, "stderr": result.stderr, "wall": wall}
    try:
        data = json.loads(result.stdout)
        data["ok"] = True
        data["bench_wall"] = wall
        return data
    except Exception as e:
        return {"ok": False, "error": str(e), "stdout": result.stdout[:500], "wall": wall}

def normalize_words(s: str):
    import re
    return re.findall(r"[a-z0-9']+", s.lower())

def wer(ref: str, hyp: str) -> float:
    """Simple word error rate via edit distance"""
    ref_words = normalize_words(ref)
    hyp_words = normalize_words(hyp)
    if not ref_words:
        return 0.0 if not hyp_words else 1.0
    # DP
    n, m = len(ref_words), len(hyp_words)
    dp = [[0]*(m+1) for _ in range(n+1)]
    for i in range(n+1): dp[i][0] = i
    for j in range(m+1): dp[0][j] = j
    for i in range(1, n+1):
        for j in range(1, m+1):
            cost = 0 if ref_words[i-1] == hyp_words[j-1] else 1
            dp[i][j] = min(dp[i-1][j]+1, dp[i][j-1]+1, dp[i-1][j-1]+cost)
    return dp[n][m] / n

def test_with_say():
    tests = [
        ("en-US", "Hello world, this is a test of Apple's new speech transcription API."),
        ("en-US", "The quick brown fox jumps over the lazy dog while pricing thirty five dollars and ninety nine cents."),
        ("en-US", "Please send an email to john at example dot com with the meeting notes."),
        ("es-ES", "Hola mundo, esto es una prueba de la nueva API de transcripción de Apple."),
        ("es-ES", "Buenos días, ¿cómo estás? Espero que tengas un excelente día."),
        ("fr-FR", "Bonjour le monde, ceci est un test de la nouvelle API de transcription d'Apple."),
        ("de-DE", "Hallo Welt, dies ist ein Test der neuen Apple Speech Transkriptions-API."),
    ]

    voice_map = {
        "en-US": "Samantha",
        "es-ES": "Paulina",
        "fr-FR": "Amelie",
        "de-DE": "Anna",
    }

    results = []
    for locale, ref_text in tests:
        voice = voice_map.get(locale, "Samantha")
        # Generate aiff
        tmp_aiff = tempfile.mktemp(suffix=".aiff")
        tmp_wav = tempfile.mktemp(suffix=".wav")
        try:
            subprocess.run(["say", "-v", voice, "-o", tmp_aiff, ref_text], timeout=10)
            subprocess.run(["afconvert", tmp_aiff, "-f", "WAVE", "-d", "LEI16@16000", tmp_wav], timeout=10)
            if not os.path.exists(tmp_wav):
                continue
            apple = run_apple_bench(tmp_wav, locale=locale)
            hyp = apple.get("text","") if apple.get("ok") else ""
            score = wer(ref_text, hyp) if apple.get("ok") else 1.0
            results.append({
                "locale": locale,
                "ref": ref_text,
                "hyp": hyp,
                "wer": score,
                "rtf": apple.get("rtf", 0),
                "processing_sec": apple.get("processing_sec", 0),
                "ok": apple.get("ok", False)
            })
            print(f"[{locale}] WER={score:.1%} RTF={apple.get('rtf',0):.1f}x")
            print(f"  REF: {ref_text}")
            print(f"  HYP: {hyp}")
            print()
        except Exception as e:
            print(f"[{locale}] error: {e}")
        finally:
            for p in [tmp_aiff, tmp_wav]:
                if os.path.exists(p):
                    os.unlink(p)

    print("\n=== Summary ===")
    for r in results:
        print(f"{r['locale']:6} WER={r['wer']:.1%} RTF={r['rtf']:.1f}x ok={r['ok']}")

    # Save
    out_path = Path(__file__).parent / "logs" / "apple_bench.json"
    out_path.parent.mkdir(exist_ok=True)
    with open(out_path, "w") as f:
        json.dump(results, f, indent=2, ensure_ascii=False)
    print(f"\nSaved to {out_path}")

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("files", nargs="*", help="WAV files to bench")
    parser.add_argument("--locale", default="en-US")
    parser.add_argument("--test-say", action="store_true", help="Generate TTS samples and bench")
    args = parser.parse_args()

    if not BIN.exists():
        print(f"Binary not found: {BIN}")
        print("Run: cd apple_speech && bash build.sh")
        sys.exit(1)

    if args.test_say or not args.files:
        print("=== Running TTS-based benchmark ===")
        test_with_say()
    else:
        for fp in args.files:
            print(f"\n=== {fp} ===")
            res = run_apple_bench(fp, locale=args.locale)
            print(json.dumps(res, indent=2, ensure_ascii=False))