"""
Cliente de teste para versão otimizada (WebSocket binário).
"""
import asyncio
import aiohttp
import json
import struct
import sys

# Tipos de mensagem binária
MSG_FRAME = 0x01
MSG_AUDIO = 0x02
MSG_AUDIO_CHUNK = 0x03

async def test_optimized(text: str, voice: str = "tara"):
    """Testa o servidor otimizado com WebSocket binário."""

    ws_url = "ws://localhost:8080/ws"

    print(f"Conectando a {ws_url}...")

    timeout = aiohttp.ClientTimeout(total=300)
    async with aiohttp.ClientSession(timeout=timeout) as session:
        async with session.ws_connect(ws_url, max_msg_size=50*1024*1024) as ws:
            print("Conectado!")

            # Enviar requisição
            request = {
                "action": "generate",
                "text": text,
                "voice": voice
            }
            print(f"Enviando: {text[:50]}...")
            await ws.send_json(request)

            frames_received = 0
            total_bytes = 0
            audio_received = False

            async for msg in ws:
                if msg.type == aiohttp.WSMsgType.TEXT:
                    data = json.loads(msg.data)
                    msg_type = data.get("type", "")

                    if msg_type == "status":
                        print(f"  Status: {data.get('message')}")

                    elif msg_type == "first_frame":
                        print(f"  Primeiro frame: {data.get('latency_ms')}ms")

                    elif msg_type == "done":
                        print(f"\n  Concluído!")
                        print(f"    Frames: {data.get('total_frames')}")
                        print(f"    Tempo: {data.get('elapsed_ms')}ms")
                        print(f"    Bytes enviados: {data.get('bytes_sent', 0)/1024:.1f}KB")
                        break

                    elif msg_type == "error":
                        print(f"  ERRO: {data.get('message')}")
                        break

                elif msg.type == aiohttp.WSMsgType.BINARY:
                    # Mensagem binária
                    buffer = msg.data
                    view = memoryview(buffer)

                    msg_type = buffer[0]

                    if msg_type == MSG_FRAME:
                        # [tipo:1][index:4][tamanho:4][dados]
                        frame_index, data_size = struct.unpack_from('<II', buffer, 1)
                        frame_data = buffer[9:9+data_size]
                        frames_received += 1
                        total_bytes += len(buffer)

                        if frames_received == 1:
                            print(f"  Recebendo frames binários...")
                        if frames_received % 25 == 0:
                            print(f"    Frame {frame_index}: {len(frame_data)} bytes")

                    elif msg_type == MSG_AUDIO:
                        # [tipo:1][sample_rate:4][tamanho:4][dados]
                        sample_rate, data_size = struct.unpack_from('<II', buffer, 1)
                        audio_data = buffer[9:9+data_size]
                        audio_received = True
                        total_bytes += len(buffer)
                        print(f"  Audio recebido: {len(audio_data)} bytes @ {sample_rate}Hz")

                elif msg.type in (aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.ERROR):
                    print("Conexão fechada")
                    break

            print(f"\nResumo:")
            print(f"  Frames recebidos: {frames_received}")
            print(f"  Total bytes binários: {total_bytes/1024:.1f}KB")
            print(f"  Audio recebido: {'Sim' if audio_received else 'Não'}")


if __name__ == "__main__":
    text = "Hello! This is a test of the optimized binary WebSocket streaming."
    if len(sys.argv) > 1:
        text = sys.argv[1]

    voice = "tara"
    if len(sys.argv) > 2:
        voice = sys.argv[2]

    asyncio.run(test_optimized(text, voice))