from fastapi import WebSocket
from typing import Dict, Set
import uuid

class WebSocketManager:
    """Manages WebSocket connections for multiplayer"""
    
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}
        self.player_nicknames: Dict[str, str] = {}
        self.nickname_to_id: Dict[str, str] = {}
    
    async def connect(self, websocket: WebSocket, nickname: str):
        """Connect a new player"""
        await websocket.accept()
        
        # Generate or reuse player ID
        if nickname in self.nickname_to_id:
            player_id = self.nickname_to_id[nickname]
        else:
            player_id = str(uuid.uuid4())
            self.nickname_to_id[nickname] = player_id
        
        self.active_connections[player_id] = websocket
        self.player_nicknames[player_id] = nickname
        
        # Broadcast player joined
        await self.broadcast({
            "type": "player_joined",
            "player_id": player_id,
            "nickname": nickname
        })
    
    async def disconnect(self, websocket: WebSocket):
        """Disconnect a player"""
        player_id = None
        for pid, ws in self.active_connections.items():
            if ws == websocket:
                player_id = pid
                break
        
        if player_id:
            del self.active_connections[player_id]
            nickname = self.player_nicknames.pop(player_id, None)
            
            # Broadcast player left
            await self.broadcast({
                "type": "player_left",
                "player_id": player_id,
                "nickname": nickname
            })
    
    async def broadcast(self, message: dict, exclude: WebSocket = None):
        """Broadcast message to all connected players"""
        disconnected = []
        
        for player_id, websocket in self.active_connections.items():
            if websocket == exclude:
                continue
            
            try:
                await websocket.send_json(message)
            except Exception:
                disconnected.append(player_id)
        
        # Clean up disconnected websockets
        for player_id in disconnected:
            if player_id in self.active_connections:
                del self.active_connections[player_id]
            if player_id in self.player_nicknames:
                del self.player_nicknames[player_id]
    
    async def broadcast_game_state(self, state: dict):
        """Broadcast full game state"""
        await self.broadcast({
            "type": "game_state_update",
            "state": state
        })
    
    async def get_player_id(self, websocket: WebSocket) -> str:
        """Get player ID from websocket"""
        for player_id, ws in self.active_connections.items():
            if ws == websocket:
                return player_id
        return None
    
    async def get_nickname(self, websocket: WebSocket) -> str:
        """Get nickname from websocket"""
        player_id = await self.get_player_id(websocket)
        return self.player_nicknames.get(player_id, "Unknown")