#!/usr/bin/env python3
"""WebSocket streaming test for Qwen3-TTS /v1/audio/speech/stream endpoint."""
import sys
import json
import time
import wave
import struct
import asyncio
import re

try:
    import websockets
except ImportError:
    print("Installing websockets...")
    import subprocess
    subprocess.check_call([sys.executable, "-m", "pip", "install", "websockets"])
    import websockets

WS_URL = "ws://localhost:8091/v1/audio/speech/stream"
OUTPUT_FILE = "test_stream_output.wav"
SAMPLE_RATE = 24000
WORD_RE = re.compile(r"\S+\s*")


def iter_words(text: str):
    """Yield complete words, preserving trailing whitespace.

    The WebSocket client must not forward raw LLM tokens such as partial
    subwords. It buffers to words and lets the server detect sentence ends.
    """
    for match in WORD_RE.finditer(text):
        yield match.group(0)


async def stream_tts(
    text: str,
    voice: str = "Ryan",
    language: str = "German",
    instructions: str = "Eine warme, ruhige deutsche Stimme, klar artikuliert, freundlich und natuerlich.",
    seed: int = 42,
):
    print(f"[→] WebSocket connect: {WS_URL}")
    print(f"    text:     {text!r}")
    print(f"    voice:    {voice}, language: {language}, seed: {seed}")
    print(f"    instruct: {instructions!r}")

    pcm_chunks = []
    total_bytes = 0
    t0 = time.perf_counter()
    first_chunk_time = None

    async with websockets.connect(WS_URL) as ws:
        # Send session config
        config = {
            "type": "session.config",
            "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
            "task_type": "VoiceDesign",
            "voice": voice,
            "language": language,
            "instructions": instructions,
            "seed": seed,
            "response_format": "pcm",
            "stream_audio": True,
            "split_granularity": "sentence",
        }
        print("[cfg]", json.dumps(config, ensure_ascii=False, sort_keys=True), flush=True)
        await ws.send(json.dumps(config))

        # Send word-by-word. Do not split into sentences client-side; the
        # server still decides when a complete sentence is ready for audio.
        for word in iter_words(text):
            await ws.send(json.dumps({
                "type": "input.text",
                "text": word,
            }))
        await ws.send(json.dumps({"type": "input.done"}))

        # Receive
        async for message in ws:
            if isinstance(message, bytes):
                if first_chunk_time is None:
                    first_chunk_time = time.perf_counter() - t0
                    print(f"[OK] First audio chunk in {first_chunk_time:.3f}s")
                pcm_chunks.append(message)
                total_bytes += len(message)
            elif isinstance(message, str):
                msg = json.loads(message)
                msg_type = msg.get("type", "")
                if msg_type == "session.done":
                    break
                elif msg_type == "error":
                    print(f"[FAIL] Server error: {msg.get('message')}")
                    sys.exit(1)
                elif msg_type in ("audio.start", "audio.done"):
                    print(f"    [{msg_type}] sentence {msg.get('sentence_index', '?')}")

    elapsed = time.perf_counter() - t0

    # Write PCM chunks as WAV
    raw_pcm = b"".join(pcm_chunks)
    with wave.open(OUTPUT_FILE, "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)  # int16
        wf.setframerate(SAMPLE_RATE)
        wf.writeframes(raw_pcm)

    audio_duration = (total_bytes / 2) / SAMPLE_RATE
    print(f"[OK] Done: {elapsed:.2f}s total, {total_bytes/1024:.1f} KB PCM")
    print(f"     Audio duration: {audio_duration:.2f}s → saved to {OUTPUT_FILE}")
    if first_chunk_time:
        print(f"     First-chunk latency: {first_chunk_time:.3f}s")


if __name__ == "__main__":
    asyncio.run(stream_tts(
        text="Herzlich willkommen! Dieses System nutzt Qwen3-TTS für deutsche Sprachausgabe.",
        voice="Ryan",
        language="German",
    ))