from fastapi import WebSocket, WebSocketDisconnect, Depends
from typing import Dict, Set
import json
import asyncio
from uuid import UUID


class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[UUID, Set[WebSocket]] = {}
        self.global_connections: Set[WebSocket] = set()

    async def connect(self, websocket: WebSocket, stream_id: UUID = None):
        await websocket.accept()
        if stream_id:
            if stream_id not in self.active_connections:
                self.active_connections[stream_id] = set()
            self.active_connections[stream_id].add(websocket)
        else:
            self.global_connections.add(websocket)

    def disconnect(self, websocket: WebSocket, stream_id: UUID = None):
        if stream_id and stream_id in self.active_connections:
            self.active_connections[stream_id].discard(websocket)
            if not self.active_connections[stream_id]:
                del self.active_connections[stream_id]
        self.global_connections.discard(websocket)

    async def send_personal_message(self, message: str, websocket: WebSocket):
        await websocket.send_text(message)

    async def broadcast_to_stream(self, message: dict, stream_id: UUID):
        if stream_id in self.active_connections:
            dead_connections = set()
            for connection in self.active_connections[stream_id]:
                try:
                    await connection.send_text(json.dumps(message))
                except Exception:
                    dead_connections.add(connection)
            
            # Remove dead connections
            for connection in dead_connections:
                self.active_connections[stream_id].discard(connection)
            
            # Clean up empty stream sets
            if not self.active_connections[stream_id]:
                del self.active_connections[stream_id]

    async def broadcast_global(self, message: dict):
        dead_connections = set()
        for connection in self.global_connections:
            try:
                await connection.send_text(json.dumps(message))
            except Exception:
                dead_connections.add(connection)
        
        # Remove dead connections
        for connection in dead_connections:
            self.global_connections.discard(connection)


manager = ConnectionManager()