"""
Locust load test for the STT / LLM / TTS stack.

GPU layout under test:
  GPU0 — LLM (:8000) + STT (:8002)
  GPU1 — TTS (:8003) dedicated (stream+pcm by default)

Metrics (appear as separate rows in the Locust Statistics table):
  llm_ttfb   – time to first SSE token (ms)   [= TTFT]
  llm_e2e    – full streaming completion (ms)
  tts_ttfa   – time to first PCM/WAV byte (ms)
  tts_e2e    – full audio download (ms)
  stt_ttfb   – time to first response byte (ms)
  stt_e2e    – full transcription (ms)
  pipe_e2e   – STT → LLM → TTS pipeline (ms)

Usage:
  ./run_locust_ui.sh              # web UI
  ./run_loadtest.sh               # headless concurrency ramp (recommended)
  TTS_STREAM=0 ./run_loadtest.sh  # compare against full-WAV path
"""

from __future__ import annotations

import io
import json
import math
import os
import struct
import time
import wave
from typing import Any

import requests
from locust import HttpUser, between, events, task

# ---------------------------------------------------------------------------
# Endpoints / models (override via env)
# ---------------------------------------------------------------------------
LLM_BASE = os.getenv("LLM_BASE", "http://127.0.0.1:8000")
STT_BASE = os.getenv("STT_BASE", "http://127.0.0.1:8002")
TTS_BASE = os.getenv("TTS_BASE", "http://127.0.0.1:8003")

LLM_MODEL = os.getenv("LLM_MODEL", "ibm-granite/granite-4.1-8b")
TTS_MODEL = os.getenv("TTS_MODEL", "Rabe3/Moss-Saudi-3")
# Empty → resolve from STT /v1/models on first use (local snapshot path).
STT_MODEL = os.getenv("STT_MODEL", "")

LLM_PROMPT = os.getenv(
    "LLM_PROMPT",
    "Reply in one short Arabic sentence saying hello.",
)
TTS_TEXT = os.getenv("TTS_TEXT", "Marhaba, kayf halak?")
LLM_MAX_TOKENS = int(os.getenv("LLM_MAX_TOKENS", "64"))

# Wait between tasks per user (think time). Lower = more aggressive.
WAIT_MIN = float(os.getenv("WAIT_MIN", "0.5"))
WAIT_MAX = float(os.getenv("WAIT_MAX", "1.5"))

_stt_model_cache: str | None = None
_fixture_wav: bytes | None = None


def _fire(
    name: str,
    response_time_ms: float,
    response_length: int = 0,
    exception: BaseException | None = None,
    request_type: str = "METRIC",
) -> None:
    events.request.fire(
        request_type=request_type,
        name=name,
        response_time=response_time_ms,
        response_length=response_length,
        exception=exception,
        context={},
    )


def _resolve_stt_model() -> str:
    global _stt_model_cache
    if STT_MODEL:
        return STT_MODEL
    if _stt_model_cache:
        return _stt_model_cache
    r = requests.get(f"{STT_BASE}/v1/models", timeout=30)
    r.raise_for_status()
    _stt_model_cache = r.json()["data"][0]["id"]
    return _stt_model_cache


def _sine_wav(seconds: float = 1.0, sr: int = 16000, hz: float = 440.0) -> bytes:
    """Small mono PCM WAV used as STT input."""
    n = int(sr * seconds)
    buf = io.BytesIO()
    with wave.open(buf, "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(sr)
        for i in range(n):
            sample = int(8000 * math.sin(2 * math.pi * hz * i / sr))
            w.writeframes(struct.pack("<h", sample))
    return buf.getvalue()


def _fixture() -> bytes:
    global _fixture_wav
    if _fixture_wav is None:
        path = os.getenv("STT_WAV", "")
        if path and os.path.isfile(path):
            with open(path, "rb") as f:
                _fixture_wav = f.read()
        else:
            _fixture_wav = _sine_wav(1.0)
    return _fixture_wav


def measure_llm_stream(
    session: requests.Session, prompt: str | None = None
) -> dict[str, Any]:
    """Streaming chat/completions → llm_ttfb + llm_e2e."""
    payload = {
        "model": LLM_MODEL,
        "messages": [{"role": "user", "content": prompt or LLM_PROMPT}],
        "max_tokens": LLM_MAX_TOKENS,
        "temperature": 0,
        "stream": True,
    }
    t0 = time.perf_counter()
    ttfb_ms: float | None = None
    tokens = 0
    text_parts: list[str] = []
    exc: BaseException | None = None
    try:
        with session.post(
            f"{LLM_BASE}/v1/chat/completions",
            json=payload,
            stream=True,
            timeout=(10, 300),
        ) as resp:
            if resp.status_code >= 400:
                body = resp.text[:300]
                raise RuntimeError(f"LLM HTTP {resp.status_code}: {body}")
            for raw in resp.iter_lines(decode_unicode=True):
                if not raw:
                    continue
                if ttfb_ms is None:
                    ttfb_ms = (time.perf_counter() - t0) * 1000
                line = raw.strip()
                if not line.startswith("data:"):
                    continue
                data = line[5:].strip()
                if data == "[DONE]":
                    break
                try:
                    chunk = json.loads(data)
                except json.JSONDecodeError:
                    continue
                delta = (chunk.get("choices") or [{}])[0].get("delta") or {}
                content = delta.get("content") or ""
                if content:
                    tokens += 1
                    text_parts.append(content)
    except BaseException as e:
        exc = e
        if ttfb_ms is None:
            ttfb_ms = (time.perf_counter() - t0) * 1000

    e2e_ms = (time.perf_counter() - t0) * 1000
    _fire("llm_ttfb", ttfb_ms or e2e_ms, exception=exc)
    _fire("llm_e2e", e2e_ms, response_length=tokens, exception=exc)
    if exc:
        raise exc
    return {
        "ttfb_ms": ttfb_ms,
        "e2e_ms": e2e_ms,
        "tokens": tokens,
        "text": "".join(text_parts),
    }


# Default: real Moss streaming (PCM). Set TTS_STREAM=0 for full-WAV path.
TTS_STREAM = os.getenv("TTS_STREAM", "1").lower() not in ("0", "false", "no")


def measure_tts(session: requests.Session, text: str | None = None) -> dict[str, Any]:
    """TTS /v1/audio/speech → tts_ttfa (first audio byte) + tts_e2e.

    With TTS_STREAM=1 (default): stream=true + response_format=pcm.
    With TTS_STREAM=0: full WAV (previous load-test path).
    """
    payload: dict[str, Any] = {"model": TTS_MODEL, "input": text or TTS_TEXT}
    if TTS_STREAM:
        payload["stream"] = True
        payload["response_format"] = "pcm"
    t0 = time.perf_counter()
    ttfa_ms: float | None = None
    audio = bytearray()
    exc: BaseException | None = None
    try:
        with session.post(
            f"{TTS_BASE}/v1/audio/speech",
            json=payload,
            stream=True,
            timeout=(10, 300),
        ) as resp:
            if resp.status_code >= 400:
                body = resp.text[:300]
                raise RuntimeError(f"TTS HTTP {resp.status_code}: {body}")
            for chunk in resp.iter_content(chunk_size=4 * 1024):
                if not chunk:
                    continue
                if ttfa_ms is None:
                    ttfa_ms = (time.perf_counter() - t0) * 1000
                audio.extend(chunk)
            if ttfa_ms is None:
                ttfa_ms = (time.perf_counter() - t0) * 1000
            if TTS_STREAM:
                if len(audio) < 1024:
                    raise RuntimeError(f"TTS PCM too short ({len(audio)} bytes)")
            elif len(audio) < 44 or audio[:4] != b"RIFF":
                raise RuntimeError(f"TTS did not return WAV (got {len(audio)} bytes)")
    except BaseException as e:
        exc = e
        if ttfa_ms is None:
            ttfa_ms = (time.perf_counter() - t0) * 1000

    e2e_ms = (time.perf_counter() - t0) * 1000
    _fire("tts_ttfa", ttfa_ms or e2e_ms, response_length=len(audio), exception=exc)
    _fire("tts_e2e", e2e_ms, response_length=len(audio), exception=exc)
    if exc:
        raise exc
    return {
        "ttfa_ms": ttfa_ms,
        "e2e_ms": e2e_ms,
        "bytes": len(audio),
        "audio": bytes(audio),
    }


def measure_stt(session: requests.Session, wav: bytes | None = None) -> dict[str, Any]:
    """Multipart /v1/audio/transcriptions → stt_ttfb + stt_e2e."""
    model = _resolve_stt_model()
    audio = wav or _fixture()
    t0 = time.perf_counter()
    ttfb_ms: float | None = None
    body = b""
    text = ""
    exc: BaseException | None = None
    try:
        with session.post(
            f"{STT_BASE}/v1/audio/transcriptions",
            files={"file": ("input.wav", audio, "audio/wav")},
            data={"model": model},
            stream=True,
            timeout=(10, 300),
        ) as resp:
            if resp.status_code >= 400:
                raise RuntimeError(f"STT HTTP {resp.status_code}: {resp.text[:300]}")
            for chunk in resp.iter_content(chunk_size=4096):
                if not chunk:
                    continue
                if ttfb_ms is None:
                    ttfb_ms = (time.perf_counter() - t0) * 1000
                body += chunk
            if ttfb_ms is None:
                ttfb_ms = (time.perf_counter() - t0) * 1000
            parsed = json.loads(body.decode("utf-8"))
            text = parsed.get("text", "")
    except BaseException as e:
        exc = e
        if ttfb_ms is None:
            ttfb_ms = (time.perf_counter() - t0) * 1000

    e2e_ms = (time.perf_counter() - t0) * 1000
    _fire("stt_ttfb", ttfb_ms or e2e_ms, response_length=len(body), exception=exc)
    _fire("stt_e2e", e2e_ms, response_length=len(body), exception=exc)
    if exc:
        raise exc
    return {"ttfb_ms": ttfb_ms, "e2e_ms": e2e_ms, "text": text}


# ---------------------------------------------------------------------------
# Locust users — pick one (or several) in the UI via class picker
# ---------------------------------------------------------------------------


class LLMUser(HttpUser):
    """Isolate LLM concurrency (GPU0). Watch llm_ttfb / llm_e2e."""

    host = LLM_BASE
    wait_time = between(WAIT_MIN, WAIT_MAX)
    weight = 3

    def on_start(self) -> None:
        self.session = requests.Session()

    def on_stop(self) -> None:
        self.session.close()

    @task
    def chat(self) -> None:
        measure_llm_stream(self.session)


class TTSUser(HttpUser):
    """Isolate TTS concurrency (GPU1, shared with STT). Watch tts_ttfa / tts_e2e."""

    host = TTS_BASE
    wait_time = between(WAIT_MIN, WAIT_MAX)
    weight = 2

    def on_start(self) -> None:
        self.session = requests.Session()

    def on_stop(self) -> None:
        self.session.close()

    @task
    def speak(self) -> None:
        measure_tts(self.session)


class STTUser(HttpUser):
    """Isolate STT concurrency (GPU1, shared with TTS). Watch stt_ttfb / stt_e2e."""

    host = STT_BASE
    wait_time = between(WAIT_MIN, WAIT_MAX)
    weight = 2

    def on_start(self) -> None:
        self.session = requests.Session()
        _resolve_stt_model()
        _fixture()

    def on_stop(self) -> None:
        self.session.close()

    @task
    def transcribe(self) -> None:
        measure_stt(self.session)


class PipelineUser(HttpUser):
    """Realistic voice turn: STT → LLM → TTS. Watch pipe_e2e plus per-stage metrics."""

    host = LLM_BASE
    wait_time = between(WAIT_MIN, WAIT_MAX)
    weight = 1

    def on_start(self) -> None:
        self.session = requests.Session()
        _resolve_stt_model()
        _fixture()

    def on_stop(self) -> None:
        self.session.close()

    @task
    def voice_turn(self) -> None:
        t0 = time.perf_counter()
        exc: BaseException | None = None
        try:
            stt = measure_stt(self.session)
            prompt = (stt.get("text") or "").strip() or LLM_PROMPT
            llm = measure_llm_stream(self.session, prompt=prompt)
            speak = (llm.get("text") or "").strip() or TTS_TEXT
            measure_tts(self.session, text=speak[:200])
        except BaseException as e:
            exc = e
        e2e_ms = (time.perf_counter() - t0) * 1000
        _fire("pipe_e2e", e2e_ms, exception=exc)
        if exc:
            raise exc