"""Exercise "Voice + tools" over the WebSocket: typed requests, then a spoken one (made with the server's TTS)."""

import asyncio
import json

import numpy as np
import websockets

URL = "ws://127.0.0.1:8000/ws"
PROMPTS = [
    "What's the weather like in Paris right now?",
    "What is 17 times 23, divided by 4?",
    "Please make the screen green.",
    "What time is it in Tokyo?",
    "Tell me a fun fact about octopuses.",  # should answer without a tool
]


async def turn(ws, send):
    await send()
    pcm, events = [], []
    while True:
        m = await ws.recv()
        if isinstance(m, bytes):
            pcm.append(np.frombuffer(m, dtype=np.float32))
            continue
        m = json.loads(m)
        if m["type"] in ("tool_call", "tool_result", "ui"):
            events.append({k: v for k, v in m.items() if k not in ("t_ms",)})
        if m["type"] in ("done", "error"):
            return m, events, np.concatenate(pcm) if pcm else np.zeros(0, np.float32)


async def main():
    async with websockets.connect(URL, max_size=None) as ws:
        await ws.recv()
        await ws.send(json.dumps({"type": "settings", "mode": "tools", "keep_history": False}))
        for p in PROMPTS:
            done, events, pcm = await turn(ws, lambda: ws.send(json.dumps({"type": "text", "text": p})))
            print(f"\n> {p}")
            for e in events:
                print("  ", e)
            s = done.get("stats", {})
            print(f"  reply: {done.get('text')!r}\n  rounds={s.get('rounds')} audio={s.get('audio_out_s')}s "
                  f"first_text={s.get('first_text_ms')}ms gen={s.get('gen_ms')}ms")

        # spoken request: synthesize it with TTS, then send it as mic audio
        q = "What's the weather in Stockholm?"
        _, _, speech = await turn(ws, lambda: ws.send(json.dumps({"type": "tts", "text": q})))

        async def speak():
            await ws.send(json.dumps({"type": "start", "rate": 24000}))
            for i in range(0, len(speech), 2048):
                await ws.send(speech[i:i + 2048].tobytes())
            await ws.send(json.dumps({"type": "stop"}))

        done, events, _ = await turn(ws, speak)
        print(f"\n> (spoken) {q}")
        for e in events:
            print("  ", e)
        print(f"  reply: {done.get('text')!r} rounds={done['stats'].get('rounds')}")


asyncio.run(main())