#!/usr/bin/env python3
"""Measure Drama 3 / SGLang TTS streaming: TTFB, burst sizes, real-time ratio.

Streams /v1/audio/speech as raw PCM and reports how the audio arrives, so the
vocoder window schedule (stream_stride / ramp / followup) and the real-time
ratio can be checked. Run against the local SGLang-Omni server.

Two modes:
  (default)    one request for the whole text -> only the server's own vocoder
               window seams are present.
  --chunked    split the text on sentence boundaries and issue one request per
               sentence, concatenating the raw PCM like the engine does -> adds
               the engine's per-chunk request boundaries on top of the seams.
               Comparing the two WAVs tells whether an artifact comes from the
               vocoder (present in both) or from the engine's chunking (only in
               the chunked one).
"""
import array
import re
import sys
import time
import wave

import requests

URL = "http://localhost:8000/v1/audio/speech"
MODEL = "fishaudio/drama-3"
SR = 44100
BYTES_PER_SAMPLE = 2

# Voice-cloning reference (matches the engine's .env defaults).
REF_AUDIO = "/root/piperust/audio-agent-engine/voice_profiles/solis.wav"
REF_TEXT = "/root/piperust/audio-agent-engine/voice_profiles/solis.txt"


def load_ref_text(path: str) -> str:
    try:
        with open(path, "r", encoding="utf-8") as f:
            return f.read().strip()
    except OSError:
        return ""


def save_wav(path: str, pcm: bytes, sr: int) -> None:
    with wave.open(path, "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(BYTES_PER_SAMPLE)
        w.setframerate(sr)
        w.writeframes(pcm)


def analyze_seams(pcm: bytes, boundaries: list[int]) -> None:
    """Compare the sample-to-sample jump exactly at each boundary against the
    jumps inside the windows. A boundary jump no larger than the interior
    distribution means no amplitude discontinuity there (the vocoder crossfade /
    the engine's crossfade are doing their job)."""
    s = array.array("h")
    s.frombytes(pcm[: len(pcm) // BYTES_PER_SAMPLE * BYTES_PER_SAMPLE])
    n = len(s)
    if n < 2:
        return
    deltas = sorted(abs(s[i] - s[i - 1]) for i in range(1, n))
    p999 = deltas[int(len(deltas) * 0.999)]
    global_max = deltas[-1]
    bounds = [b for b in boundaries[:-1] if 0 < b < n]
    bj = [abs(s[b] - s[b - 1]) for b in bounds]
    max_b = max(bj) if bj else 0
    rank = deltas.index(max_b) / len(deltas) * 100 if bj else 0
    print(f"seam analysis: boundaries={len(bounds)} max_boundary_jump={max_b} "
          f"global_max_jump={global_max} global_p99.9={p999} "
          f"max_boundary_percentile={rank:.2f}")


def stream_pcm(text: str, use_ref: bool, t0: float):
    """POST one text and return (pcm_bytes, events, sample_rate).
    events = list of (t_seconds, bytes_cumulative, chunk_bytes)."""
    body = {
        "model": MODEL,
        "input": text,
        "voice": "default",
        "response_format": "pcm",
        "stream": True,
    }
    if use_ref:
        body["references"] = [
            {"audio_path": REF_AUDIO, "text": load_ref_text(REF_TEXT)}
        ]
    resp = requests.post(URL, json=body, stream=True, timeout=120)
    resp.raise_for_status()
    sr = int(resp.headers.get("x-sample-rate") or SR)
    pcm = bytearray()
    events = []
    total = 0
    with resp:
        for chunk in resp.iter_content(chunk_size=None):
            if not chunk:
                continue
            total += len(chunk)
            pcm += chunk
            events.append((time.perf_counter() - t0, total, len(chunk)))
    return bytes(pcm), events, sr


def summarise_events(events) -> None:
    ttfb = events[0][0]
    tlast = events[-1][0]
    audio_s = (events[-1][1] // BYTES_PER_SAMPLE) / SR
    print(f"chunks={len(events)} ttfb={ttfb*1000:.0f}ms "
          f"last={tlast*1000:.0f}ms audio={audio_s:.2f}s "
          f"wall={tlast:.2f}s ratio={tlast/audio_s:.3f}x")


def measure(text: str, use_ref: bool = True, out: str | None = None) -> None:
    print(f"[single] text={text!r} ref={use_ref}")
    t0 = time.perf_counter()
    pcm, events, sr = stream_pcm(text, use_ref, t0)
    print("x-sample-rate:", sr)
    summarise_events(events)
    boundaries = [cum // BYTES_PER_SAMPLE for _, cum, _ in events]
    analyze_seams(pcm, boundaries)
    if out:
        save_wav(out, pcm, sr)
        print(f"saved WAV -> {out}")


def measure_chunked(text: str, use_ref: bool = True, out: str | None = None) -> None:
    parts = [p for p in re.split(r"(?<=[.!?;])\s+", text.strip()) if p]
    print(f"[chunked] parts={len(parts)} ref={use_ref}")
    t0 = time.perf_counter()
    pcm = bytearray()
    events = []
    sr = SR
    total = 0
    for i, part in enumerate(parts):
        t_req = time.perf_counter()
        p, ev, sr = stream_pcm(part, use_ref, t0)
        ttfb = ev[0][0] - (t_req - t0)
        print(f"  part {i}: {len(part):3d} chars ttfb={ttfb*1000:4.0f}ms "
              f"audio={(len(p)//BYTES_PER_SAMPLE)/sr:.2f}s part={part[:40]!r}")
        pcm += p
        for t, cum, cb in ev:
            total += cb
            events.append((t, total, cb))
    print("x-sample-rate:", sr)
    summarise_events(events)
    boundaries = [cum // BYTES_PER_SAMPLE for _, cum, _ in events]
    analyze_seams(bytes(pcm), boundaries)
    if out:
        save_wav(out, bytes(pcm), sr)
        print(f"saved WAV -> {out}")


if __name__ == "__main__":
    args = sys.argv[1:]
    use_ref = "--no-ref" not in args

    def arg_value(flag):
        return args[args.index(flag) + 1] if flag in args else None

    out = arg_value("--out")
    text = arg_value("--text") or (
        "Ciao, questa e' una prova per misurare la latenza del motore vocale. "
        "Contiamo insieme da uno fino a venti: uno, due, tre, quattro, cinque, "
        "sei, sette, otto, nove, dieci, undici, dodici, tredici, quattordici, "
        "quindici, sedici, diciassette, diciotto, diciannove, venti.")
    if "--chunked" in args:
        measure_chunked(text, use_ref=use_ref, out=out)
    else:
        measure(text, use_ref=use_ref, out=out)