"""
WebSocket Registry Module

This module provides efficient O(1) user lookup and management for WebSocket connections.
Replaces O(n) set iteration with dictionary-based lookups.
"""


class UserRegistry:
    """
    Efficient O(1) user lookup registry.

    Replaces O(n) set iteration with O(1) dictionary lookup.
    """

    def __init__(self):
        self._users = {}  # userID -> websocket

    def add(self, websocket) -> None:
        """Add a websocket connection"""
        if hasattr(websocket, 'userID'):
            self._users[websocket.userID] = websocket

    def remove(self, websocket) -> None:
        """Remove a websocket connection"""
        if hasattr(websocket, 'userID') and websocket.userID in self._users:
            del self._users[websocket.userID]

    def get(self, user_id: str):
        """Get websocket by user ID (O(1))"""
        return self._users.get(user_id)

    def count(self) -> int:
        """Get number of connected users"""
        return len(self._users)

    def __contains__(self, websocket) -> bool:
        """Check if websocket is registered"""
        if hasattr(websocket, 'userID'):
            return websocket.userID in self._users
        return False


def get_websocket_for_player(player_id):
    """
    Get websocket connection for a specific player ID.

    Args:
        player_id: The player's user ID

    Returns:
        WebSocket connection or None if not found
    """
    return USERS.get(player_id)


# Global user registry instance
USERS = UserRegistry()