"""Live-path test of server.py over a real WebSocket (real-time paced 32 ms frames).
python ws_test.py <ws_url> <user1.wav> [<barge.wav>]  -> prints events; saves ws_test_agent.wav"""
import asyncio, json, sys, time, numpy as np, soundfile as sf, librosa, websockets
async def main(url, w1, w2=None):
    def load(p):
        y, sr = sf.read(p, dtype="float32"); y = y.mean(1) if y.ndim > 1 else y
        return (librosa.resample(y, orig_sr=sr, target_sr=16000) * 32767).astype(np.int16)
    u1 = load(w1); u2 = load(w2) if w2 else None
    t0 = time.time(); ev = []; audio = []; first_audio = [None]; got_reply = [False]; got_reply = [False]
    async with websockets.connect(url, max_size=2 ** 24) as ws:
        async def rx():
            async for m in ws:
                t = round(time.time() - t0, 3)
                if isinstance(m, bytes):
                    audio.append(np.frombuffer(m, np.int16)); first_audio[0] = first_audio[0] or (t if got_reply[0] else None); ev.append((t, "audio", len(m) // 2 / 24000))
                else:
                    d = json.loads(m); got_reply[0] = got_reply[0] or d.get("type") == "reply"; ev.append((t, d.get("type"), {k: v for k, v in d.items() if k != "type"})); print(t, d, flush=True)
        task = asyncio.create_task(rx())
        async def send(y):
            for i in range(0, len(y) - 511, 512):
                await ws.send(y[i:i + 512].tobytes()); await asyncio.sleep(0.032)
        z = np.zeros(512, np.int16)
        await send(u1); print(round(time.time() - t0, 3), "user1 sent", flush=True)
        for _ in range(int(6.0 / 0.032)):
            await ws.send(z.tobytes()); await asyncio.sleep(0.032)
            if u2 is not None and first_audio[0] and time.time() - t0 > first_audio[0] + 1.0: break
        if u2 is not None:
            print(round(time.time() - t0, 3), "barge-in user2 starts", flush=True); await send(u2)
        for _ in range(int(8 / 0.032)):
            await ws.send(z.tobytes()); await asyncio.sleep(0.032)
        task.cancel()
    if audio: sf.write("ws_test_agent.wav", np.concatenate(audio).astype(np.float32) / 32768, 24000)
    print("WS_TEST_DONE", "first_audio_at", first_audio[0])
asyncio.run(main(*sys.argv[1:]))