#!/usr/bin/env python3
"""
Quick WebSocket test - sends a few tokens and receives audio
"""
import asyncio
import websockets
import json

async def quick_test():
    uri = "ws://localhost:8765/ws/stream"
    
    print(f"Connecting to {uri}...")
    async with websockets.connect(uri) as websocket:
        print("✓ Connected!")
        
        # Send a simple sentence token by token
        tokens = ["Hello", "!", " ", "This", " ", "is", " ", "a", " ", "test", "."]
        
        for token in tokens:
            print(f"→ Sending: '{token}'")
            await websocket.send(json.dumps({"token": token, "pitch_shift": 0}))
            
            # Try to receive audio (non-blocking)
            try:
                response = await asyncio.wait_for(websocket.recv(), timeout=5.0)  # Increased timeout
                if isinstance(response, bytes):
                    print(f"  ← Received {len(response)} bytes of audio")
                else:
                    print(f"  ← Received: {response}")
            except asyncio.TimeoutError:
                print(f"  (no response yet)")
        
        print("\n→ Flushing buffer...")
        await websocket.send(json.dumps({"flush": True}))
        
        # Get remaining audio
        try:
            while True:
                response = await asyncio.wait_for(websocket.recv(), timeout=10.0)  # Increased timeout
                if isinstance(response, bytes):
                    print(f"  ← Received {len(response)} bytes of audio (flush)")
                else:
                    print(f"  ← Received: {response}")
        except asyncio.TimeoutError:
            print("  (flush complete)")
        
        print("\n✓ Test complete!")

if __name__ == "__main__":
    asyncio.run(quick_test())