#!/usr/bin/env python3
"""
Quick test script to verify WebSocket STT connectivity
"""

import asyncio
import websockets
import json
import base64
import tempfile
import wave
import numpy as np

async def test_websocket_stt():
    """Test WebSocket STT connection to verify it's working"""
    
    # Create a simple test audio file (1 second of silence)
    sample_rate = 16000
    duration = 1.0
    samples = int(sample_rate * duration)
    audio_data = np.zeros(samples, dtype=np.int16)
    
    # Write to temporary WAV file
    with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_file:
        temp_path = temp_file.name
        
        with wave.open(temp_path, 'wb') as wav_file:
            wav_file.setnchannels(1)
            wav_file.setsampwidth(2)
            wav_file.setframerate(sample_rate)
            wav_file.writeframes(audio_data.tobytes())
    
    # Read the audio file
    with open(temp_path, 'rb') as f:
        audio_bytes = f.read()
        audio_base64 = base64.b64encode(audio_bytes).decode('utf-8')
    
    print(f"🔍 TEST: Created test audio file: {len(audio_bytes)} bytes")
    
    # Test WebSocket URLs
    urls_to_test = [
        "wss://pgits-stt-gpu-service.hf.space/ws/stt",
        "ws://localhost:7860/ws/stt"
    ]
    
    for url in urls_to_test:
        print(f"\n🧪 TESTING: {url}")
        try:
            import ssl
            ssl_context = ssl.create_default_context()
            ssl_context.check_hostname = False
            ssl_context.verify_mode = ssl.CERT_NONE
            
            connect_kwargs = {}
            if url.startswith("wss://"):
                connect_kwargs["ssl"] = ssl_context
                
            async with websockets.connect(url, **connect_kwargs) as websocket:
                print(f"✅ Connected to {url}")
                
                # First, wait for connection confirmation
                print(f"📥 Waiting for connection confirmation...")
                confirm_response = await asyncio.wait_for(websocket.recv(), timeout=10.0)
                confirm_result = json.loads(confirm_response)
                print(f"📥 Connection confirmation: {confirm_result}")
                
                if confirm_result.get("type") != "stt_connection_confirmed":
                    print(f"❌ Unexpected connection response: {confirm_result.get('type')}")
                    continue
                
                # Now send test message
                message = {
                    "type": "stt_audio_chunk",
                    "audio_data": audio_base64,
                    "language": "auto",
                    "model_size": "base",
                    "is_final": True
                }
                
                print(f"📤 Sending audio message...")
                await websocket.send(json.dumps(message))
                print(f"📤 Audio message sent")
                
                # Wait for transcription response
                print(f"📥 Waiting for transcription response...")
                response = await asyncio.wait_for(websocket.recv(), timeout=30.0)
                result = json.loads(response)
                
                print(f"📥 Transcription response: {result}")
                
                if result.get("type") == "stt_transcription":
                    print(f"✅ STT SUCCESS: '{result.get('text', 'NO_TEXT')}'")
                elif result.get("type") == "stt_error":
                    print(f"❌ STT ERROR: {result.get('message', 'NO_MESSAGE')}")
                else:
                    print(f"❓ UNEXPECTED TRANSCRIPTION: {result.get('type', 'NO_TYPE')}")
                    
        except Exception as e:
            print(f"❌ Error testing {url}: {e}")
            print(f"🔍 Exception type: {type(e).__name__}")

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