#!/usr/bin/env python3
"""R4 probe — integrated three-hop end-to-end latency.

Per PLAN.md SLICE-01 TASK-01-05: feed a sample ASR transcript → Ollama
gemma4:cloud streaming → Cartesia TTS streaming; measure end-to-end
(transcript-in → first-audio-out). Run 10 iterations. Also measure the same
path with Piper self-hosted (if Piper can be stood up locally; otherwise note
as pending and pre-stage in SLICE-02).

Exit code 0 in all cases:
  - If OLLAMA_API_KEY or CARTESIA_API_KEY is missing, print KEY_MISSING and
    exit 0 (the probe infrastructure is the deliverable).
  - If present, run the live integrated probe and print e2e latency.

The Piper leg is invoked only if PRAXIS_TTS=piper is set AND a Piper voice
model is available; otherwise it is documented as pre-staged (R4 mitigation).
"""

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 _missing(keys: list[str]) -> None:
    _banner(
        "KEY_MISSING — " + ", ".join(keys) + " not set.\n"
        "  Cannot run live integrated e2e probe. Probe infrastructure is built\n"
        "  and ready; live measurements are pending API key provisioning.\n"
        "  Set the missing key(s) in .env (see .env.example) and re-run."
    )


CHAT_URL = os.environ.get("OLLAMA_CHAT_URL", "https://ollama.com/api/chat")
CARTESIA_WS_URL = "wss://api.cartesia.ai/tts/websocket"
ROLEPLAY_MODEL = os.environ.get("OLLAMA_ROLEPLAY_MODEL", "gemma4:cloud")
DEFAULT_VOICE_ID = "a3536a36-1d18-4efb-a95a-7c44b7b5e384"

# The "transcript-in" — a realistic ASR final transcript from the learner.
SAMPLE_TRANSCRIPT = "Hi, I want to help you with your order. What happened?"
SYSTEM_PROMPT = (
    "You are Jordan, a customer who received a damaged product. "
    "You are frustrated but not abusive. Stay in character. Keep responses "
    "to 1-2 sentences."
)


async def _ollama_first_token(api_key: str) -> tuple[str | None, float | None, str | None]:
    """Stream Ollama gemma4:cloud, return (full_text, ttft_ms, error)."""
    import httpx

    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
    body = {
        "model": ROLEPLAY_MODEL,
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": SAMPLE_TRANSCRIPT},
        ],
        "stream": True,
    }
    t0 = time.perf_counter()
    ttft_ms: float | None = None
    chunks: list[str] = []

    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            async with client.stream("POST", CHAT_URL, headers=headers, json=body) as resp:
                if resp.status_code != 200:
                    text = await resp.aread()
                    return None, None, f"HTTP {resp.status_code}: {text[:200]!r}"
                async for line in resp.aiter_lines():
                    if not line:
                        continue
                    try:
                        chunk = json.loads(line)
                    except json.JSONDecodeError:
                        continue
                    content = chunk.get("message", {}).get("content", "")
                    if content:
                        if ttft_ms is None:
                            ttft_ms = (time.perf_counter() - t0) * 1000.0
                        chunks.append(content)
    except Exception as exc:  # pragma: no cover
        return None, None, f"connection error: {exc}"

    return "".join(chunks), ttft_ms, None


async def _cartesia_first_audio(
    api_key: str, text: str, voice_id: str, model_id: str
) -> tuple[float | None, str | None]:
    """Send text to Cartesia WS, return (first_audio_ms_from_t0, error)."""
    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": text,
                "voice": {"id": voice_id},
                "output_format": {
                    "container": "raw",
                    "encoding": "pcm_s16le",
                    "sample_rate": 24000,
                },
                "stream": True,
            }
            await ws.send(json.dumps(req))
            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
                if isinstance(msg, str):
                    data = json.loads(msg)
                    if data.get("type") == "done":
                        break
    except Exception as exc:  # pragma: no cover
        return None, f"cartesia error: {exc}"

    return first_audio_ms, None


async def _piper_first_audio(text: str) -> tuple[float | None, str | None]:
    """Synthesize via Piper self-hosted, return (first_audio_ms, error).

    Piper pre-staging note (R4 mitigation): Piper is pre-staged as the
    production v0.1 TTS fallback per ARCHITECTURE.md. The voice model must be
    downloaded separately (see docs/latency-report.md). If not available,
    returns an error string that the caller documents as pending.
    """
    try:
        from piper import PiperVoice  # type: ignore
    except ImportError:
        return None, "piper-tts not installed (pre-staged for SLICE-02)"

    model_path = os.environ.get("PIPER_VOICE_MODEL", "")
    if not model_path or not Path(model_path).exists():
        return None, "PIPER_VOICE_MODEL not set or file missing (pre-staged for SLICE-02)"

    import io

    t0 = time.perf_counter()
    try:
        voice = PiperVoice.load(model_path)
        wav_bytes = io.BytesIO()
        for chunk in voice.synthesize(text):
            wav_bytes.write(chunk.audio_int16_bytes)
        first_audio_ms = (time.perf_counter() - t0) * 1000.0
        return first_audio_ms, None
    except Exception as exc:  # pragma: no cover
        return None, f"piper error: {exc}"


async def _e2e_once_cartesia(ollama_key: str, cartesia_key: str, voice_id: str, model_id: str) -> dict:
    """Run the integrated ASR-transcript → Ollama → Cartesia path once."""
    t_start = time.perf_counter()
    text, ttft_ms, llm_err = await _ollama_first_token(ollama_key)
    if llm_err or not text:
        return {"ok": False, "error": llm_err or "empty LLM output", "ttft_ms": None}
    tts_ms, tts_err = await _cartesia_first_audio(cartesia_key, text, voice_id, model_id)
    if tts_err or tts_ms is None:
        return {"ok": False, "error": tts_err or "no TTS audio", "ttft_ms": ttft_ms}
    e2e_ms = (time.perf_counter() - t_start) * 1000.0
    return {
        "ok": True,
        "ttft_ms": ttft_ms,
        "tts_first_audio_ms": tts_ms,
        "e2e_ms": e2e_ms,
        "llm_text": text[:80],
    }


async def _e2e_once_piper(ollama_key: str) -> dict:
    """Run the integrated ASR-transcript → Ollama → Piper path once."""
    t_start = time.perf_counter()
    text, ttft_ms, llm_err = await _ollama_first_token(ollama_key)
    if llm_err or not text:
        return {"ok": False, "error": llm_err or "empty LLM output", "ttft_ms": None}
    tts_ms, tts_err = await _piper_first_audio(text)
    if tts_err or tts_ms is None:
        return {"ok": False, "error": tts_err or "no TTS audio", "ttft_ms": ttft_ms, "piper_pending": True}
    e2e_ms = (time.perf_counter() - t_start) * 1000.0
    return {
        "ok": True,
        "ttft_ms": ttft_ms,
        "tts_first_audio_ms": tts_ms,
        "e2e_ms": e2e_ms,
        "llm_text": text[:80],
    }


def _summarize(samples: list[float], label: str) -> dict:
    if not samples:
        print(f"  {label}: no samples collected.")
        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="R4 integrated e2e latency probe")
    parser.add_argument("--iterations", type=int, default=10)
    parser.add_argument("--voice-id", default=os.environ.get("CARTESIA_VOICE_ID", DEFAULT_VOICE_ID))
    parser.add_argument("--cartesia-model", default="sonic-2")
    parser.add_argument("--out", default=None)
    parser.add_argument("--piper", action="store_true", help="also run the Piper leg")
    args = parser.parse_args()

    _banner("R4 PROBE — integrated three-hop e2e (transcript → Ollama → TTS)")
    ollama_key = os.environ.get("OLLAMA_API_KEY", "").strip()
    cartesia_key = os.environ.get("CARTESIA_API_KEY", "").strip()

    missing = []
    if not ollama_key:
        missing.append("OLLAMA_API_KEY")
    if not cartesia_key:
        missing.append("CARTESIA_API_KEY")
    if missing:
        _missing(missing)
        return 0

    # ── Cartesia leg ────────────────────────────────────────────────────────
    print(f"\n  Cartesia leg — {args.iterations} iterations:")
    e2e_samples: list[float] = []
    ttft_samples: list[float] = []
    tts_samples: list[float] = []
    for i in range(args.iterations):
        r = await _e2e_once_cartesia(ollama_key, cartesia_key, args.voice_id, args.cartesia_model)
        if r.get("ok"):
            e2e_samples.append(r["e2e_ms"])
            ttft_samples.append(r["ttft_ms"])
            tts_samples.append(r["tts_first_audio_ms"])
            print(f"  [{i + 1:2d}/{args.iterations}] e2e={r['e2e_ms']:.1f}ms "
                  f"(llm_ttft={r['ttft_ms']:.1f}, tts={r['tts_first_audio_ms']:.1f})")
        else:
            print(f"  [{i + 1:2d}/{args.iterations}] error: {r.get('error')}")
        await asyncio.sleep(0.5)

    print()
    e2e_summary = _summarize(e2e_samples, "e2e_cartesia")
    ttft_summary = _summarize(ttft_samples, "e2e_cartesia_llm_ttft")
    tts_summary = _summarize(tts_samples, "e2e_cartesia_tts_first_audio")

    # ── Piper leg (optional / pre-staged) ───────────────────────────────────
    piper_summary: dict = {}
    if args.piper:
        print(f"\n  Piper leg — {args.iterations} iterations:")
        p_e2e: list[float] = []
        p_ttft: list[float] = []
        p_tts: list[float] = []
        for i in range(args.iterations):
            r = await _e2e_once_piper(ollama_key)
            if r.get("ok"):
                p_e2e.append(r["e2e_ms"])
                p_ttft.append(r["ttft_ms"])
                p_tts.append(r["tts_first_audio_ms"])
                print(f"  [{i + 1:2d}/{args.iterations}] e2e={r['e2e_ms']:.1f}ms")
            elif r.get("piper_pending"):
                print(f"  [{i + 1:2d}/{args.iterations}] Piper pre-staged (pending voice model) — skipping")
                break
            else:
                print(f"  [{i + 1:2d}/{args.iterations}] error: {r.get('error')}")
            await asyncio.sleep(0.5)
        print()
        piper_summary = _summarize(p_e2e, "e2e_piper")
    else:
        print("\n  Piper leg not requested (--piper). Piper is pre-staged as the R4 "
              "mitigation per ARCHITECTURE.md; live Piper measurement pending "
              "voice-model provisioning (see docs/latency-report.md).")

    # ── Budget comparison ───────────────────────────────────────────────────
    budget = 600.0
    print(f"\n  Latency budget: {budget:.0f}ms")
    if e2e_samples:
        med = statistics.median(e2e_samples)
        over = med > budget
        print(f"  Cartesia median e2e: {med:.1f}ms — {'OVER' if over else 'WITHIN'} budget "
              f"(delta {med - budget:+.1f}ms)")
    if piper_summary.get("n"):
        # type: ignore
        med = piper_summary.get("median_ms")
        if med:
            over = med > budget
            print(f"  Piper   median e2e: {med:.1f}ms — {'OVER' if over else 'WITHIN'} budget "
                  f"(delta {med - budget:+.1f}ms)")

    print()
    if args.out:
        result = {
            "e2e_cartesia": e2e_summary,
            "e2e_cartesia_llm_ttft": ttft_summary,
            "e2e_cartesia_tts_first_audio": tts_summary,
            "e2e_piper": piper_summary,
            "budget_ms": budget,
        }
        Path(args.out).write_text(json.dumps(result, indent=2))
        print(f"  Wrote {args.out}")
    return 0


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


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