#!/usr/bin/env python3
"""R2 probe — Cartesia Sonic TTS first-audio-byte latency.

Per PLAN.md SLICE-01 TASK-01-03: WebSocket to Cartesia Sonic, send a sample text
chunk, measure first-audio-byte latency over 20 iterations; log min/median/p95.

Exit code 0 in all cases:
  - If CARTESIA_API_KEY is missing, print KEY_MISSING and exit 0.
  - If present, run the live probe and print a latency table.
"""

from __future__ import annotations

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

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

try:
    from dotenv import load_dotenv

    load_dotenv()
except ImportError:  # pragma: no cover
    pass


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


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


SAMPLE_TEXT = (
    "Hi, I received my order yesterday and the item is cracked. "
    "I want my money back."
)

CARTESIA_WS_URL = "wss://api.cartesia.ai/tts/websocket"
DEFAULT_VOICE_ID = "a3536a36-1d18-4efb-a95a-7c44b7b5e384"


async def _probe_once(api_key: str, voice_id: str, model_id: str) -> float | None:
    """Open Cartesia WS, request TTS, return ms-to-first-audio-byte."""
    import websockets

    headers = [("x-api-key", api_key), ("cartesia-version", "2024-06-10")]
    t0 = time.perf_counter()
    first_audio_ms: float | None = None

    try:
        async with websockets.connect(
            CARTESIA_WS_URL, additional_headers=headers, open_timeout=10
        ) as ws:
            req = {
                "model_id": model_id,
                "transcript": SAMPLE_TEXT,
                "voice": {"id": voice_id},
                "output_format": {
                    "container": "raw",
                    "encoding": "pcm_s16le",
                    "sample_rate": 24000,
                },
                "stream": True,
            }
            await ws.send(json.dumps(req))
            # Read frames until we get the first audio chunk.
            while True:
                msg = await asyncio.wait_for(ws.recv(), timeout=10)
                if isinstance(msg, (bytes, bytearray)):
                    first_audio_ms = (time.perf_counter() - t0) * 1000.0
                    break
                # JSON control messages (e.g. done) — ignore until audio.
                if isinstance(msg, str):
                    data = json.loads(msg)
                    if data.get("type") == "done":
                        break
    except Exception as exc:  # pragma: no cover - network/auth errors
        print(f"  [probe] Cartesia connection failed: {exc}")
        return None

    return first_audio_ms


async def run_live(api_key: str, iterations: int, voice_id: str, model_id: str) -> list[float]:
    samples: list[float] = []
    print(f"  Running {iterations} Cartesia Sonic iterations (voice={voice_id})...")
    for i in range(iterations):
        ms = await _probe_once(api_key, voice_id, model_id)
        if ms is not None:
            samples.append(ms)
            print(f"  [{i + 1:2d}/{iterations}] first-audio: {ms:6.1f} ms")
        else:
            print(f"  [{i + 1:2d}/{iterations}] no audio 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="R2 Cartesia Sonic latency probe")
    parser.add_argument("--iterations", type=int, default=20)
    parser.add_argument("--voice-id", default=os.environ.get("CARTESIA_VOICE_ID", DEFAULT_VOICE_ID))
    parser.add_argument("--model-id", default="sonic-2")
    parser.add_argument("--out", default=None)
    args = parser.parse_args()

    _banner("R2 PROBE — Cartesia Sonic first-audio-byte latency")
    api_key = _require_key()
    if api_key is None:
        return 0

    samples = await run_live(api_key, args.iterations, args.voice_id, args.model_id)
    summary = _summarize(samples, "cartesia_sonic_first_audio")
    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())