"""
WebSocket Test Client

Tests the WebSocket server with sine wave audio generation.

Usage:
    python test_client.py
    python test_client.py --url ws://localhost:8000/ws
"""

import asyncio
import aiohttp
import json
import struct
import math
import argparse
from datetime import datetime

# Configuration
SERVER_URL = "ws://localhost:8000/ws"
SAMPLE_RATE = 16000
FREQUENCY = 440  # 440Hz sine wave
CHUNK_DURATION_MS = 20
CHUNK_SIZE_BYTES = int(SAMPLE_RATE * 2 * (CHUNK_DURATION_MS / 1000.0))  # 640 bytes


def generate_sine_wave(duration_ms=1000):
    """
    Generate sine wave audio data.

    Format: 16kHz, mono, 16-bit PCM
    """
    num_samples = int(SAMPLE_RATE * (duration_ms / 1000.0))
    audio_data = bytearray()

    for x in range(num_samples):
        # Generate sine wave sample
        value = int(32767.0 * math.sin(2 * math.pi * FREQUENCY * x / SAMPLE_RATE))
        # Pack as little-endian 16-bit signed integer
        audio_data.extend(struct.pack('<h', value))

    return audio_data


async def receive_loop(ws, session_id):
    """
    Listen for incoming messages from the server.
    """
    print("👂 Listening for server responses...")
    async for msg in ws:
        timestamp = datetime.now().strftime("%H:%M:%S")

        if msg.type == aiohttp.WSMsgType.TEXT:
            try:
                data = json.loads(msg.data)
                event_type = data.get('event', 'Unknown')
                print(f"[{timestamp}] 📨 Event: {event_type}")
                print(f"           {json.dumps(data, indent=2)}")
            except json.JSONDecodeError:
                print(f"[{timestamp}] 📨 Text: {msg.data[:100]}...")

        elif msg.type == aiohttp.WSMsgType.BINARY:
            # Received audio chunk back
            print(f"[{timestamp}] 🔊 Audio: {len(msg.data)} bytes")

        elif msg.type == aiohttp.WSMsgType.CLOSED:
            print(f"\n[{timestamp}] ❌ Connection closed")
            break

        elif msg.type == aiohttp.WSMsgType.ERROR:
            print(f"\n[{timestamp}] ⚠️ WebSocket error")
            break


async def send_audio_loop(ws):
    """
    Stream sine wave audio to the server.
    """
    print("🎙️  Starting audio stream (sine wave)...")

    # Generate 5 seconds of audio
    audio_buffer = generate_sine_wave(5000)
    cursor = 0

    while cursor < len(audio_buffer):
        chunk = audio_buffer[cursor:cursor + CHUNK_SIZE_BYTES]
        if not chunk:
            break

        await ws.send_bytes(chunk)
        print(f"📤 Sent audio chunk: {len(chunk)} bytes", end="\r")

        cursor += len(chunk)

        # Sleep to simulate real-time (20ms per chunk)
        await asyncio.sleep(CHUNK_DURATION_MS / 1000.0)

    print("\n✅ Finished streaming audio")


async def run_client(url):
    """
    Run the WebSocket test client.
    """
    session = aiohttp.ClientSession()

    try:
        print(f"🔌 Connecting to {url}...")

        async with session.ws_connect(url) as ws:
            print("✅ Connected!")
            print()

            # Send invite command
            invite_cmd = {
                "command": "invite",
                "option": {
                    "codec": "pcm",
                    "samplerate": SAMPLE_RATE
                }
            }
            await ws.send_json(invite_cmd)
            print("📤 Sent invite command")
            print()

            # Send a ping command
            ping_cmd = {"command": "ping"}
            await ws.send_json(ping_cmd)
            print("📤 Sent ping command")
            print()

            # Wait a moment for responses
            await asyncio.sleep(1)

            # Run audio streaming and receiving in parallel
            await asyncio.gather(
                receive_loop(ws, None),
                send_audio_loop(ws)
            )

    except aiohttp.ClientConnectorError:
        print(f"❌ Connection failed. Is the server running at {url}?")
        print(f"   Start server with: python main.py")

    except Exception as e:
        print(f"❌ Error: {e}")

    finally:
        await session.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="WebSocket Audio Test Client")
    parser.add_argument("--url", default=SERVER_URL, help="WebSocket URL")
    args = parser.parse_args()

    try:
        asyncio.run(run_client(args.url))
    except KeyboardInterrupt:
        print("\n👋 Client stopped.")