"""
WebSocket support for real-time updates
"""

from typing import Dict, List, Set
from fastapi import WebSocket, WebSocketDisconnect
import json
import asyncio


class ConnectionManager:
    """Manage WebSocket connections"""
    
    def __init__(self):
        self.active_connections: Dict[str, List[WebSocket]] = {}
        self.user_connections: Dict[str, WebSocket] = {}
    
    async def connect(self, websocket: WebSocket, client_id: str, room: str = "default"):
        """Connect a client"""
        await websocket.accept()
        
        if room not in self.active_connections:
            self.active_connections[room] = []
        
        self.active_connections[room].append(websocket)
        self.user_connections[client_id] = websocket
    
    def disconnect(self, websocket: WebSocket, client_id: str, room: str = "default"):
        """Disconnect a client"""
        if room in self.active_connections:
            if websocket in self.active_connections[room]:
                self.active_connections[room].remove(websocket)
        
        if client_id in self.user_connections:
            del self.user_connections[client_id]
    
    async def broadcast(self, message: dict, room: str = "default"):
        """Broadcast message to all clients in a room"""
        if room not in self.active_connections:
            return
        
        disconnected = []
        for connection in self.active_connections[room]:
            try:
                await connection.send_json(message)
            except:
                disconnected.append(connection)
        
        # Clean up disconnected clients
        for conn in disconnected:
            if conn in self.active_connections[room]:
                self.active_connections[room].remove(conn)
    
    async def send_to_client(self, client_id: str, message: dict):
        """Send message to specific client"""
        if client_id in self.user_connections:
            await self.user_connections[client_id].send_json(message)
    
    async def broadcast_observation(self, observation: dict):
        """Broadcast new observation"""
        await self.broadcast({
            "type": "new_observation",
            "data": observation
        }, room="observations")
    
    async def broadcast_signal(self, signal: dict):
        """Broadcast new signal"""
        await self.broadcast({
            "type": "new_signal",
            "data": signal
        }, room="signals")
    
    async def broadcast_decision(self, decision: dict):
        """Broadcast new decision"""
        await self.broadcast({
            "type": "new_decision",
            "data": decision
        }, room="decisions")
    
    async def broadcast_outcome(self, outcome: dict):
        """Broadcast new outcome"""
        await self.broadcast({
            "type": "new_outcome",
            "data": outcome
        }, room="outcomes")
    
    def get_stats(self) -> dict:
        """Get connection statistics"""
        return {
            "total_connections": sum(len(conns) for conns in self.active_connections.values()),
            "rooms": {
                room: len(conns)
                for room, conns in self.active_connections.items()
            },
            "unique_clients": len(self.user_connections)
        }


# Global connection manager
manager = ConnectionManager()


async def websocket_endpoint(websocket: WebSocket, client_id: str, room: str = "default"):
    """WebSocket endpoint"""
    await manager.connect(websocket, client_id, room)
    
    try:
        while True:
            # Receive message from client
            data = await websocket.receive_text()
            message = json.loads(data)
            
            # Handle different message types
            if message.get("type") == "subscribe":
                # Client wants to subscribe to a room
                new_room = message.get("room", "default")
                # In production, validate subscription permissions
                pass
            
            elif message.get("type") == "ping":
                await websocket.send_json({"type": "pong"})
            
    except WebSocketDisconnect:
        manager.disconnect(websocket, client_id, room)
    except Exception as e:
        manager.disconnect(websocket, client_id, room)
        print(f"WebSocket error: {e}")