ai-voicebot/server/websocket/connection.py

"""
WebSocket connection management.

This module handles WebSocket connections and integrates with the message router.
"""

from typing import Dict, Any, Optional, TYPE_CHECKING
from fastapi import WebSocket, WebSocketDisconnect

from logger import logger
from .message_handlers import MessageRouter

if TYPE_CHECKING:
    # Use absolute imports to avoid relative import issues
    try:
        from core.session_manager import Session, SessionManager
        from core.lobby_manager import Lobby, LobbyManager
        from core.auth_manager import AuthManager
    except ImportError:
        # Fallback for when running from different directory structure
        from ..core.session_manager import Session, SessionManager
        from ..core.lobby_manager import Lobby, LobbyManager
        from ..core.auth_manager import AuthManager


class WebSocketConnectionManager:
    """Manages WebSocket connections and message processing"""
    
    def __init__(
        self,
        session_manager: "SessionManager",
        lobby_manager: "LobbyManager", 
        auth_manager: "AuthManager"
    ):
        self.session_manager = session_manager
        self.lobby_manager = lobby_manager
        self.auth_manager = auth_manager
        self.message_router = MessageRouter()
        
        # Managers dict for injection into handlers
        self.managers = {
            "session_manager": session_manager,
            "lobby_manager": lobby_manager,
            "auth_manager": auth_manager,
        }

    async def handle_connection(
        self, 
        websocket: WebSocket, 
        lobby_id: str, 
        session_id: str
    ):
        """Handle a WebSocket connection for a session in a lobby"""
        await websocket.accept()
        
        # Validate inputs
        if not lobby_id:
            await websocket.send_json({
                "type": "error", 
                "data": {"error": "Invalid or missing lobby"}
            })
            await websocket.close()
            return
            
        if not session_id:
            await websocket.send_json({
                "type": "error", 
                "data": {"error": "Invalid or missing session"}
            })
            await websocket.close()
            return
        
        # Get session
        session = self.session_manager.get_session(session_id)
        if not session:
            await websocket.send_json({
                "type": "error", 
                "data": {"error": f"Invalid session ID {session_id}"}
            })
            await websocket.close()
            return
        
        # Get lobby
        lobby = self.lobby_manager.get_lobby(lobby_id)
        if not lobby:
            await websocket.send_json({
                "type": "error", 
                "data": {"error": f"Lobby not found: {lobby_id}"}
            })
            await websocket.close()
            return
        
        logger.info(f"{session.getName()} <- lobby_joined({lobby.getName()})")
        
        # Set up connection
        session.ws = websocket
        session.update_last_used()
        
        # Clean up stale session in lobby if needed
        if session.id in lobby.sessions:
            logger.info(f"{session.getName()} - Stale session in lobby {lobby.getName()}. Re-joining.")
            try:
                await session.leave_lobby(lobby)
            except Exception as e:
                logger.warning(f"Error cleaning up stale session: {e}")
        
        # Notify existing peers about new user
        await self._notify_peers_of_join(session, lobby)
        
        try:
            # Message processing loop
            while True:
                packet = await websocket.receive_json()
                session.update_last_used()
                
                message_type = packet.get("type", None)
                data: Optional[Dict[str, Any]] = packet.get("data", None)
                
                if not message_type:
                    logger.error(f"{session.getName()} - Invalid request: {packet}")
                    await websocket.send_json({
                        "type": "error", 
                        "data": {"error": "Invalid request"}
                    })
                    continue
                
                # Route message to appropriate handler
                await self.message_router.route(
                    message_type, session, lobby, data or {}, websocket, self.managers
                )
                
        except WebSocketDisconnect:
            logger.info(f"{session.getName()} <- WebSocket disconnected")
        except Exception as e:
            logger.error(f"Error in WebSocket connection for {session.getName()}: {e}")
        finally:
            # Clean up connection
            await self._cleanup_connection(session, lobby)

    async def _notify_peers_of_join(self, session: "Session", lobby: "Lobby"):
        """Notify existing peers about a new user joining"""
        failed_peers = []
        
        with lobby.lock:
            peer_sessions = list(lobby.sessions.values())
        
        for peer_session in peer_sessions:
            if not peer_session.ws:
                logger.warning(
                    f"{session.getName()} - Live peer session {peer_session.id} not found in lobby {lobby.getName()}. Marking for removal."
                )
                failed_peers.append(peer_session.id)
                continue

            logger.info(f"{session.getName()} -> user_joined({peer_session.getName()})")
            try:
                await peer_session.ws.send_json({
                    "type": "user_joined",
                    "data": {
                        "session_id": session.id,
                        "name": session.name,
                    },
                })
            except Exception as e:
                logger.warning(f"Failed to notify {peer_session.getName()} of user join: {e}")
                failed_peers.append(peer_session.id)

        # Clean up failed peers
        with lobby.lock:
            for failed_peer_id in failed_peers:
                if failed_peer_id in lobby.sessions:
                    del lobby.sessions[failed_peer_id]

    async def _cleanup_connection(self, session: "Session", lobby: "Lobby"):
        """Clean up when connection is closed"""
        try:
            # Clear WebSocket reference
            session.ws = None
            
            # Remove from lobby if present
            if session.id in lobby.sessions:
                await session.leave_lobby(lobby)
                logger.info(f"Removed {session.getName()} from lobby {lobby.getName()} on disconnect")
            
        except Exception as e:
            logger.error(f"Error during connection cleanup for {session.getName()}: {e}")

    def add_message_handler(self, message_type: str, handler):
        """Add a custom message handler"""
        self.message_router.register(message_type, handler)

    def get_supported_message_types(self) -> list[str]:
        """Get list of supported message types"""
        return self.message_router.get_supported_types()