#!/usr/bin/env python3
"""
Quick script to get JWT token and test WebSocket SOS connection
"""

import asyncio
import json
import sys

try:
    import httpx
    import websockets
except ImportError:
    print("Installing required packages...")
    import subprocess
    subprocess.check_call([sys.executable, "-m", "pip", "install", "httpx", "websockets"])
    import httpx
    import websockets

async def test_ws_connection():
    """Test WebSocket connection"""
    
    # 1. Get token
    print("\n🔐 Step 1: Getting JWT token...")
    try:
        async with httpx.AsyncClient() as client:
            response = await client.post(
                "http://localhost:8001/api/v1/auth/login",
                json={
                    "email": "wstester@test.com",
                    "password": "WsTest1234!"
                },
                timeout=10.0
            )
            
            print(f"Status: {response.status_code}")
            
            if response.status_code != 200:
                print(f"Response: {response.text}")
                print("❌ Login failed")
                return False
            
            data = response.json()
            token = data.get("access_token")
            
            # Decode token to get user_id
            import base64
            parts = token.split('.')
            if len(parts) == 3:
                payload = parts[1]
                # Add padding if needed
                padding = 4 - len(payload) % 4
                if padding != 4:
                    payload += '=' * padding
                decoded = base64.urlsafe_b64decode(payload)
                token_data = json.loads(decoded)
                user_id = token_data.get("sub")
            else:
                user_id = 51  # fallback
            
            print(f"✅ Token obtained: {token[:50]}...")
            print(f"✅ User ID: {user_id}")
            
    except Exception as e:
        print(f"❌ Error: {e}")
        return False
    
    # 2. Test WebSocket connection
    print("\n📡 Step 2: Testing WebSocket connection...")
    ws_url = f"ws://localhost:8002/api/v1/emergency/ws/{user_id}?token={token}"
    
    try:
        async with websockets.connect(ws_url) as websocket:
            print("✅ WebSocket connected!")
            
            # Receive connection message
            message = await asyncio.wait_for(websocket.recv(), timeout=5.0)
            print(f"📨 Received: {message}")
            
            # Send ping
            print("\n🏓 Step 3: Testing ping/pong...")
            await websocket.send(json.dumps({"type": "ping"}))
            print("📤 Ping sent")
            
            # Receive pong
            message = await asyncio.wait_for(websocket.recv(), timeout=5.0)
            print(f"📨 Received: {message}")
            
            # Keep connection open for a bit to receive any broadcasts
            print("\n👂 Step 4: Listening for broadcasts (10 seconds)...")
            print("(Leave this running and create an emergency alert in another terminal)")
            
            try:
                while True:
                    message = await asyncio.wait_for(websocket.recv(), timeout=10.0)
                    data = json.loads(message)
                    print(f"📨 Broadcast received: {json.dumps(data, indent=2)}")
                    
            except asyncio.TimeoutError:
                print("⏱️  No broadcasts received (timeout)")
            
            print("\n✅ WebSocket connection test completed successfully!")
            return True
            
    except Exception as e:
        print(f"❌ WebSocket error: {e}")
        return False

if __name__ == "__main__":
    try:
        result = asyncio.run(test_ws_connection())
        sys.exit(0 if result else 1)
    except KeyboardInterrupt:
        print("\n⚠️  Test interrupted by user")
        sys.exit(1)