#!/usr/bin/env python3
"""R1 probe — Deepgram Nova-3 streaming ASR first-partial-transcript latency.

Per PLAN.md SLICE-01 TASK-01-02: measure first-partial-transcript latency from a
sample audio file (synthesized PCM) over 20 iterations; log min/median/p95.

Exit code 0 in all cases:
  - If DEEPGRAM_API_KEY is missing, print a clear KEY_MISSING banner and exit 0
    (the probe infrastructure is the deliverable; live numbers come when keys
    are provisioned).
  - If the key is present, run the live probe and print a latency table.

Usage:
    python scripts/probe_deepgram.py [--iterations N] [--model nova-3]
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import statistics
import sys
import time
from pathlib import Path

# Make the project importable when run from the repo root.
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

try:
    from dotenv import load_dotenv

    load_dotenv()
except ImportError:  # pragma: no cover - dotenv is a declared dep
    pass


def _banner(msg: str) -> None:
    print("\n" + "=" * 72)
    print(msg)
    print("=" * 72 + "\n")


def _require_key() -> str | None:
    """Return the Deepgram API key or None (with a printed banner if missing)."""
    key = os.environ.get("DEEPGRAM_API_KEY", "").strip()
    if not key:
        _banner(
            "KEY_MISSING — DEEPGRAM_API_KEY not set.\n"
            "  Cannot run live Deepgram probe. Probe infrastructure is built\n"
            "  and ready; live measurements are pending API key provisioning.\n"
            "  Set DEEPGRAM_API_KEY in .env (see .env.example) and re-run."
        )
        return None
    return key


def _synth_pcm(duration_s: float = 2.0, sample_rate: int = 16000) -> bytes:
    """Synthesize a short mono 16-bit PCM buffer (silence + a low tone).

    Deepgram needs real audio frames; we generate a recognizable signal so the
    streaming endpoint returns a partial. The exact transcript content is not
    the point — the *latency to first partial* is.
    """
    import math
    import struct

    n = int(duration_s * sample_rate)
    frames = bytearray()
    for i in range(n):
        # 220 Hz tone for the first 1.5s, then silence — a clearly voiced segment.
        if i < int(1.5 * sample_rate):
            sample = int(16000 * math.sin(2 * math.pi * 220 * i / sample_rate))
        else:
            sample = 0
        frames += struct.pack("<h", sample)
    return bytes(frames)


DEEPGRAM_WS_URL = "wss://api.deepgram.com/v1/listen"


async def _probe_once(api_key: str, model: str, pcm: bytes, sample_rate: int) -> float | None:
    """Open a Deepgram streaming WebSocket, send PCM, return ms-to-first-partial.

    Uses the raw Deepgram streaming WebSocket API (not the SDK) so the probe is
    independent of SDK version churn and measures the actual network path.
    """
    import websockets

    params = (
        f"?model={model}&language=en&encoding=linear16&channels=1"
        f"&sample_rate={sample_rate}&interim_results=true&endpointing=300"
    )
    headers = [("Authorization", f"Token {api_key}")]
    t0 = time.perf_counter()
    first_partial_ms: float | None = None

    try:
        async with websockets.connect(
            DEEPGRAM_WS_URL + params, additional_headers=headers, open_timeout=10
        ) as ws:
            # Send in small chunks to mimic real streaming.
            chunk = 3200  # 100ms of 16kHz mono 16-bit
            for i in range(0, len(pcm), chunk):
                await ws.send(pcm[i : i + chunk])
                await asyncio.sleep(0.02)
            # Wait for the first transcript message.
            try:
                while True:
                    msg = await asyncio.wait_for(ws.recv(), timeout=5)
                    if isinstance(msg, str):
                        data = json.loads(msg)
                        if data.get("type") == "Results":
                            channel = data.get("channel", {})
                            alts = channel.get("alternatives", [])
                            if alts and alts[0].get("transcript", "").strip():
                                first_partial_ms = (time.perf_counter() - t0) * 1000.0
                                break
            except asyncio.TimeoutError:
                pass
            # Signal close.
            try:
                await ws.send(json.dumps({"type": "CloseStream"}))
            except Exception:
                pass
    except Exception as exc:  # pragma: no cover - network/auth errors
        print(f"  [probe] Deepgram connection failed: {exc}")
        return None

    return first_partial_ms


async def run_live(api_key: str, iterations: int, model: str) -> list[float]:
    sample_rate = 16000
    pcm = _synth_pcm(duration_s=2.0, sample_rate=sample_rate)
    samples: list[float] = []
    print(f"  Running {iterations} Deepgram Nova-3 iterations (model={model})...")
    for i in range(iterations):
        ms = await _probe_once(api_key, model, pcm, sample_rate)
        if ms is not None:
            samples.append(ms)
            print(f"  [{i + 1:2d}/{iterations}] first-partial: {ms:6.1f} ms")
        else:
            print(f"  [{i + 1:2d}/{iterations}] no partial received (skipped)")
        await asyncio.sleep(0.3)
    return samples


def _summarize(samples: list[float], label: str) -> dict:
    if not samples:
        print(f"\n  {label}: no samples collected.\n")
        return {"label": label, "n": 0}
    s = sorted(samples)
    p95 = s[int(0.95 * (len(s) - 1))]
    row = {
        "label": label,
        "n": len(s),
        "min_ms": round(min(s), 1),
        "median_ms": round(statistics.median(s), 1),
        "p95_ms": round(p95, 1),
        "mean_ms": round(statistics.mean(s), 1),
    }
    print(
        f"  {label}: n={row['n']} min={row['min_ms']:.1f} "
        f"median={row['median_ms']:.1f} p95={row['p95_ms']:.1f} "
        f"mean={row['mean_ms']:.1f} (ms)"
    )
    return row


async def amain() -> int:
    parser = argparse.ArgumentParser(description="R1 Deepgram Nova-3 latency probe")
    parser.add_argument("--iterations", type=int, default=20)
    parser.add_argument("--model", default=os.environ.get("DEEPGRAM_MODEL", "nova-3"))
    parser.add_argument("--out", default=None, help="optional JSON results path")
    args = parser.parse_args()

    _banner("R1 PROBE — Deepgram Nova-3 first-partial-transcript latency")
    api_key = _require_key()
    if api_key is None:
        return 0

    samples = await run_live(api_key, args.iterations, args.model)
    summary = _summarize(samples, "deepgram_nova3_first_partial")
    print()
    if args.out:
        Path(args.out).write_text(json.dumps(summary, indent=2))
        print(f"  Wrote {args.out}")
    return 0


def main() -> int:
    return asyncio.run(amain())


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