import logging
from typing import Dict, Set
import socketio

logger = logging.getLogger(f'{__name__}.ConnectionManager')


class ConnectionManager:
    """Manages WebSocket connections and rooms"""

    def __init__(self, sio: socketio.AsyncServer):
        self.sio = sio
        self.user_sessions: Dict[str, str] = {}  # sid -> user_id
        self.game_rooms: Dict[str, Set[str]] = {}  # game_id -> set of sids

    async def connect(self, sid: str, user_id: str) -> None:
        """Register a new connection"""
        self.user_sessions[sid] = user_id
        logger.info(f"User {user_id} connected with session {sid}")

    async def disconnect(self, sid: str) -> None:
        """Handle disconnection"""
        user_id = self.user_sessions.pop(sid, None)
        if user_id:
            logger.info(f"User {user_id} disconnected (session {sid})")

            # Remove from all game rooms
            for game_id, sids in self.game_rooms.items():
                if sid in sids:
                    sids.remove(sid)
                    await self.broadcast_to_game(
                        game_id,
                        "user_disconnected",
                        {"user_id": user_id}
                    )

    async def join_game(self, sid: str, game_id: str, role: str) -> None:
        """Add user to game room"""
        await self.sio.enter_room(sid, game_id)

        if game_id not in self.game_rooms:
            self.game_rooms[game_id] = set()
        self.game_rooms[game_id].add(sid)

        user_id = self.user_sessions.get(sid)
        logger.info(f"User {user_id} joined game {game_id} as {role}")

        await self.broadcast_to_game(
            game_id,
            "user_connected",
            {"user_id": user_id, "role": role}
        )

    async def leave_game(self, sid: str, game_id: str) -> None:
        """Remove user from game room"""
        await self.sio.leave_room(sid, game_id)

        if game_id in self.game_rooms:
            self.game_rooms[game_id].discard(sid)

        user_id = self.user_sessions.get(sid)
        logger.info(f"User {user_id} left game {game_id}")

    async def broadcast_to_game(
        self,
        game_id: str,
        event: str,
        data: dict
    ) -> None:
        """Broadcast event to all users in game room"""
        await self.sio.emit(event, data, room=game_id)
        logger.debug(f"Broadcast {event} to game {game_id}")

    async def emit_to_user(self, sid: str, event: str, data: dict) -> None:
        """Emit event to specific user"""
        await self.sio.emit(event, data, room=sid)

    def get_game_participants(self, game_id: str) -> Set[str]:
        """Get all session IDs in game room"""
        return self.game_rooms.get(game_id, set())