"""Drive the server's WebSocket like the browser does: stream a wav in real time in 2048-sample chunks, 3 turns."""

import asyncio
import json
import sys

import numpy as np
import soundfile as sf
import websockets


async def main(path, url="ws://127.0.0.1:8000/ws"):
    audio, rate = sf.read(path, dtype="float32")
    async with websockets.connect(url, max_size=None) as ws:
        print("<", json.loads(await ws.recv())["type"])
        for turn in range(3):  # later turns exercise the cached history
            await ws.send(json.dumps({"type": "start", "rate": rate}))
            for i in range(0, len(audio), 2048):
                await ws.send(audio[i:i + 2048].tobytes())
                await asyncio.sleep(2048 / rate)  # real-time pacing, like a mic
            await ws.send(json.dumps({"type": "stop"}))
            pcm_bytes = 0
            while True:
                msg = await ws.recv()
                if isinstance(msg, bytes):
                    pcm_bytes += len(msg)
                    continue
                m = json.loads(msg)
                if m["type"] in ("token", "audio_chunk", "log"):
                    continue
                if m["type"] == "done":
                    keys = ("first_text_ms", "first_frame_ms", "first_audio_ms", "flush_ms", "prefill_calls",
                            "prefill_busy_ms", "context_tokens", "context_tokens_after", "rtf", "gen_ms", "codec_decode_ms", "audio_out_s")
                    print("< done", m["text"][:80], {k: m["stats"].get(k) for k in keys})
                elif m["type"] != "prompt":
                    print("<", m["type"], {k: v for k, v in m.items() if k != "type"})
                if m["type"] in ("done", "error"):
                    break
            print(f"  received {pcm_bytes / 4 / 24000:.2f}s of PCM")


asyncio.run(main(sys.argv[1]))