#!/usr/bin/env python3
"""ASR-judge synthesized wavs with Soniox stt-async-v5 (same engine as the
training transcripts) -> WER/CER vs reference text.

Env: SONIOX_API_KEY required.
Usage: python asr_eval_soniox.py --wavdir /opt/work/eval/step200
Writes <wavdir>/asr_report_soniox.json (same shape as the whisper report).
"""

import argparse
import json
import os
import re
import time
import unicodedata
from pathlib import Path

import jiwer
import requests

API = "https://api.soniox.com"
KEY = os.environ["SONIOX_API_KEY"]
HDR = {"Authorization": f"Bearer {KEY}"}

TASHKEEL_RE = re.compile(
    "[ؐ-ًؚ-ٰٟۖ-ۜ۟-۪ۨ-ۭـ]"
)
PUNCT_RE = re.compile(r"[^\w\s]|[_]", re.UNICODE)
WS_RE = re.compile(r"\s+")


def norm(t: str) -> str:
    t = unicodedata.normalize("NFC", t)
    t = TASHKEEL_RE.sub("", t)
    t = t.replace("أ", "ا").replace("إ", "ا").replace("آ", "ا")
    t = t.replace("ى", "ي").replace("ة", "ه")
    t = PUNCT_RE.sub(" ", t)
    return WS_RE.sub(" ", t).strip()


def transcribe(wav: Path) -> str:
    with open(wav, "rb") as f:
        r = requests.post(f"{API}/v1/files", headers=HDR,
                          files={"file": (wav.name, f, "audio/wav")}, timeout=60)
    r.raise_for_status()
    file_id = r.json()["id"]
    try:
        r = requests.post(f"{API}/v1/transcriptions", headers=HDR, json={
            "file_id": file_id,
            "model": "stt-async-v5",
            "language_hints": ["ar"],
        }, timeout=30)
        r.raise_for_status()
        tid = r.json()["id"]
        try:
            for _ in range(120):
                s = requests.get(f"{API}/v1/transcriptions/{tid}", headers=HDR,
                                 timeout=30).json()
                if s.get("status") == "completed":
                    break
                if s.get("status") == "error":
                    raise RuntimeError(f"soniox error: {s.get('error_message')}")
                time.sleep(1)
            else:
                raise TimeoutError("soniox poll timeout")
            tr = requests.get(f"{API}/v1/transcriptions/{tid}/transcript",
                              headers=HDR, timeout=30).json()
            return tr.get("text", "")
        finally:
            requests.delete(f"{API}/v1/transcriptions/{tid}", headers=HDR, timeout=30)
    finally:
        requests.delete(f"{API}/v1/files/{file_id}", headers=HDR, timeout=30)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--wavdir", required=True)
    args = ap.parse_args()
    wavdir = Path(args.wavdir)

    refs = {}
    for line in (wavdir / "timing.jsonl").read_text(encoding="utf-8").splitlines():
        r = json.loads(line)
        if "text" in r:
            refs[r["id"]] = r["text"]

    rows = []
    for sid, ref in refs.items():
        wav = wavdir / f"{sid}.wav"
        if not wav.exists():
            rows.append({"id": sid, "error": "missing_wav"})
            continue
        try:
            hyp = transcribe(wav)
        except Exception as e:
            rows.append({"id": sid, "error": str(e)[:200]})
            print(f"{sid:12s} ERROR {e}")
            continue
        r_n, h_n = norm(ref), norm(hyp)
        wer = jiwer.wer(r_n, h_n) if r_n else 1.0
        cer = jiwer.cer(r_n, h_n) if r_n else 1.0
        rows.append({"id": sid, "ref": ref, "hyp": hyp,
                     "wer": round(wer, 3), "cer": round(cer, 3)})
        print(f"{sid:12s} WER={wer:.2f} CER={cer:.2f} | {hyp[:70]}")

    ok = [r for r in rows if "wer" in r]
    summary = {
        "engine": "soniox-stt-async-v5",
        "n": len(rows), "n_ok": len(ok),
        "mean_wer": round(sum(r["wer"] for r in ok) / max(len(ok), 1), 4),
        "mean_cer": round(sum(r["cer"] for r in ok) / max(len(ok), 1), 4),
    }
    print("SUMMARY", json.dumps(summary))
    (wavdir / "asr_report_soniox.json").write_text(
        json.dumps({"summary": summary, "rows": rows}, ensure_ascii=False, indent=1),
        encoding="utf-8")


if __name__ == "__main__":
    main()