#!/usr/bin/env python
"""Latency benchmark for the MentorCore voice pipeline.

Measures p50/p95 per stage for a full voice turn and compares against
the ARCHITECTURE.md latency budget. Exits nonzero if p50 first-audio > 2.0s.

Usage:
    uv run python scripts/latency-bench.py --persona marcus-aurelius [--runs 10]

Requirements: compiled persona + ANTHROPIC_API_KEY in env (uses real pipeline).
For dry-run with FakeProvider: set FAKE_PROVIDER=1 in env.
"""

from __future__ import annotations

import argparse
import asyncio
import os
import sys
import time
import wave
from pathlib import Path

# ── Budget table (seconds) from ARCHITECTURE.md ───────────────────────────────
BUDGET = {
    "stt": 0.250,
    "retrieval": 0.030,
    "llm_first_sentence": 0.700,
    "tts_first_chunk": 0.250,
    "first_audio_total": 2.000,
}

_REPO_ROOT = Path(__file__).resolve().parent.parent
_FIXTURE_WAV = _REPO_ROOT / "tests" / "fixtures" / "audio" / "fixture.wav"


def _load_fixture_wav() -> bytes:
    with wave.open(str(_FIXTURE_WAV), "rb") as wf:
        return wf.readframes(wf.getnframes())


def _percentile(data: list[float], pct: int) -> float:
    if not data:
        return float("nan")
    sorted_data = sorted(data)
    idx = (pct / 100) * (len(sorted_data) - 1)
    lo, hi = int(idx), min(int(idx) + 1, len(sorted_data) - 1)
    return sorted_data[lo] + (sorted_data[hi] - sorted_data[lo]) * (idx - lo)


async def _run_one(slug: str, audio_bytes: bytes) -> dict[str, float]:
    """Run one full voice turn and return per-stage latencies in seconds."""
    from runtime.api.chat import _index_dir, _package_dir
    from runtime.conversation.prompt import assemble
    from runtime.conversation.retrieval import retrieve
    from runtime.speech.tts import VoiceConfig, synthesize_stream
    from shared.config import MODELS
    from shared.package import load_package

    persona = load_package(slug)
    package_dir = _package_dir(slug)
    index_dir = _index_dir(slug)

    # ── STT ──────────────────────────────────────────────────────────────────
    if os.environ.get("FAKE_PROVIDER") == "1":
        from runtime.speech.stt import FakeSTT

        stt_fn = FakeSTT().transcribe
    else:
        from runtime.speech.stt import transcribe as stt_fn  # type: ignore[assignment]

    t0 = time.perf_counter()
    transcript = stt_fn(audio_bytes)
    stt_elapsed = time.perf_counter() - t0

    if not transcript:
        transcript = "Tell me about virtue."  # fallback for silence fixture

    # ── Retrieval ─────────────────────────────────────────────────────────────
    t0 = time.perf_counter()
    retrieved = retrieve(transcript, index_dir, top_k=5)
    retrieval_elapsed = time.perf_counter() - t0

    # ── Prompt assembly ───────────────────────────────────────────────────────
    system_prompt, _ = assemble(persona, package_dir, retrieved_chunks=retrieved)

    # ── LLM streaming ─────────────────────────────────────────────────────────
    if os.environ.get("FAKE_PROVIDER") == "1":
        from runtime.providers.fake import FakeProvider

        provider = FakeProvider.single("Be just. Seek wisdom. Endure hardship with equanimity.")
    else:
        from runtime.providers.factory import get_provider

        provider = get_provider()

    messages = [{"role": "user", "content": transcript}]

    llm_first_sentence: float = 0.0
    tts_first_chunk: float = 0.0
    first_audio_total: float = 0.0

    # Buffer for sentence detection
    sentence_buf = ""
    _SENTENCE_END = __import__("re").compile(r"(?<=[.!?])\s+")

    t_llm_start = time.perf_counter()
    t_first_sentence: float | None = None
    t_first_audio: float | None = None

    token_queue: asyncio.Queue[str | None] = asyncio.Queue()

    async def _stream_tokens() -> None:
        nonlocal t_first_sentence
        async for chunk in provider.stream_chat(
            system=system_prompt,
            messages=messages,
            model=MODELS.runtime_conversation,
        ):
            await token_queue.put(chunk)
        await token_queue.put(None)

    async def _token_iter():
        nonlocal sentence_buf, t_first_sentence
        while True:
            tok = await token_queue.get()
            if tok is None:
                if sentence_buf.strip():
                    if t_first_sentence is None:
                        t_first_sentence = time.perf_counter()
                    yield sentence_buf.strip()
                break
            sentence_buf += tok
            parts = _SENTENCE_END.split(sentence_buf)
            if len(parts) > 1:
                for s in parts[:-1]:
                    if s.strip():
                        if t_first_sentence is None:
                            t_first_sentence = time.perf_counter()
                        yield s.strip()
                sentence_buf = parts[-1]

    fake = os.environ.get("FAKE_PROVIDER") == "1"
    voice_config = VoiceConfig(engine="fake") if fake else VoiceConfig()

    stream_task = asyncio.create_task(_stream_tokens())
    tts_start: float | None = None

    async for _pcm in synthesize_stream(_token_iter(), voice_config):
        if t_first_audio is None:
            t_first_audio = time.perf_counter()
            tts_start = t_first_sentence or time.perf_counter()

    await stream_task

    now = time.perf_counter()
    llm_first_sentence = (t_first_sentence or now) - t_llm_start
    tts_first_chunk = (t_first_audio or now) - (tts_start or t_llm_start)
    first_audio_total = (t_first_audio or now) - t_llm_start + stt_elapsed + retrieval_elapsed

    return {
        "stt": stt_elapsed,
        "retrieval": retrieval_elapsed,
        "llm_first_sentence": llm_first_sentence,
        "tts_first_chunk": tts_first_chunk,
        "first_audio_total": first_audio_total,
    }


def _fmt(val: float, budget: float) -> str:
    flag = "✓" if val <= budget else "✗"
    return f"{val*1000:6.0f}ms  (budget {budget*1000:.0f}ms) {flag}"


async def _main(slug: str, runs: int) -> int:
    print(f"\nLatency benchmark — {slug}, {runs} runs")
    print("=" * 60)

    audio_bytes = _load_fixture_wav()

    results: dict[str, list[float]] = {k: [] for k in BUDGET}

    for i in range(runs):
        print(f"  Run {i+1}/{runs}...", end=" ", flush=True)
        r = await _run_one(slug, audio_bytes)
        for k in BUDGET:
            results[k].append(r[k])
        print(f"first_audio={r['first_audio_total']*1000:.0f}ms")

    print()
    print(f"{'Stage':<25} {'p50':>10}  {'p95':>10}")
    print("-" * 60)

    p50_first_audio = _percentile(results["first_audio_total"], 50)

    for stage, budget in BUDGET.items():
        p50 = _percentile(results[stage], 50)
        p95 = _percentile(results[stage], 95)
        p50_str = f"{p50*1000:6.0f}ms"
        p95_str = f"{p95*1000:6.0f}ms"
        flag_p50 = "✓" if p50 <= budget else "✗"
        flag_p95 = "✓" if p95 <= budget * 1.5 else "✗"
        bgt = f"{budget*1000:.0f}ms"
        print(f"{stage:<25} {p50_str} {flag_p50}   {p95_str} {flag_p95}   (budget {bgt})")

    print()
    bgt_ms = BUDGET["first_audio_total"] * 1000
    p50_ms = p50_first_audio * 1000
    if p50_first_audio <= BUDGET["first_audio_total"]:
        print(f"✓ p50 first-audio {p50_ms:.0f}ms ≤ {bgt_ms:.0f}ms budget — PASS")
        return 0
    else:
        print(f"✗ p50 first-audio {p50_ms:.0f}ms > {bgt_ms:.0f}ms budget — FAIL")
        return 1


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="MentorCore voice latency benchmark")
    parser.add_argument("--persona", required=True, help="Persona slug to benchmark")
    parser.add_argument("--runs", type=int, default=10, help="Number of runs (default: 10)")
    args = parser.parse_args()

    sys.exit(asyncio.run(_main(args.persona, args.runs)))