"""Client-side bench for chatterbox-fast — ground-truth TTFB + starvation check.

The server's scheduler reasons about starvation from an *estimated* buffer; this
client measures the real thing: time-to-first-byte over the wire, and whether a
true 1×-realtime consumer ever runs dry. Saves the streamed audio to a .wav so
the operator can ear-A/B it against the one-shot.

Usage:
  python bench.py --host http://10.100.79.3:8197 --text "..." --out /refs/_fast.wav
  python bench.py --host ... --oneshot --out /refs/_oneshot.wav   # A/B baseline
Stdlib only (urllib + wave) so it runs anywhere, incl. inside the container.
"""

from __future__ import annotations

import argparse
import json
import time
import urllib.request
import wave

SR = 24000
BYTES_PER_SEC = SR * 2  # s16le mono

DEFAULT_TEXT = (
    "The cake is a lie. "
    "I am being entirely sincere, without a trace of sarcasm, when I say that "
    "this is the single most important scientific breakthrough in the entire "
    "history of this facility. "
    "You will be baked, and then there will be cake. "
    "It is delicious and moist, assuming you survive the testing protocol, which "
    "the available data suggests you almost certainly will not."
)


def run(host: str, text: str, out: str, *, oneshot: bool, voice: str | None) -> None:
    payload = {"text": text, "format": "pcm", "stream": not oneshot}
    if voice:
        payload["voice"] = voice
    req = urllib.request.Request(
        f"{host}/tts",
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json"},
        method="POST",
    )

    t0 = time.perf_counter()
    t_first = t0
    ttfb: float | None = None
    received = 0  # bytes
    worst_lead = float("inf")  # min (buffered_audio_s - elapsed_since_first_s)
    starve_events = 0
    pcm = bytearray()

    with urllib.request.urlopen(req) as resp:
        while True:
            buf = resp.read(4096)
            if not buf:
                break
            now = time.perf_counter()
            if ttfb is None:
                ttfb = now - t0
                t_first = now
            received += len(buf)
            pcm += buf
            buffered_s = received / BYTES_PER_SEC
            elapsed_s = now - t_first
            lead = buffered_s - elapsed_s  # >0 ⇒ ahead of a 1× player
            if lead < worst_lead:
                worst_lead = lead
            if lead < 0:
                starve_events += 1

    total = time.perf_counter() - t0
    audio_s = received / BYTES_PER_SEC
    if ttfb is None:
        raise SystemExit("no audio received from server")

    with wave.open(out, "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(SR)
        w.writeframes(bytes(pcm))

    mode = "oneshot" if oneshot else "stream"
    print(f"[{mode}] ttfb={ttfb*1000:.0f}ms  audio={audio_s:.2f}s  "
          f"wall={total:.2f}s  rtf={audio_s/total:.2f}x")
    if not oneshot:
        verdict = "OK (stayed ahead)" if starve_events == 0 else f"STARVED ({starve_events} reads dry)"
        print(f"        worst lead over 1x player = {worst_lead:.2f}s  → {verdict}")
    print(f"        wrote {out}")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--host", default="http://10.100.79.3:8197")
    ap.add_argument("--text", default=DEFAULT_TEXT)
    ap.add_argument("--out", default="/refs/_fast.wav")
    ap.add_argument("--voice", default=None)
    ap.add_argument("--oneshot", action="store_true", help="whole-text one-shot baseline")
    args = ap.parse_args()
    run(args.host, args.text, args.out, oneshot=args.oneshot, voice=args.voice)


if __name__ == "__main__":
    main()