ai-voicebot/server/websocket/webrtc_signaling.py

"""
WebRTC Signaling Handlers

This module contains WebRTC signaling message handlers for peer-to-peer communication.
Handles ICE candidate relay and session description exchange between peers.
"""

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

from logger import logger

if TYPE_CHECKING:
    from core.session_manager import Session
    from core.lobby_manager import Lobby


class WebRTCSignalingHandlers:
    """WebRTC signaling message handlers for peer-to-peer communication."""

    @staticmethod
    async def handle_relay_ice_candidate(
        websocket: WebSocket,
        session: "Session",
        lobby: "Lobby",
        data: Dict[str, Any]
    ) -> None:
        """
        Handle ICE candidate relay between peers.
        
        Args:
            websocket: The WebSocket connection
            session: The sender session
            lobby: The lobby context
            data: Message data containing peer_id and candidate
        """
        logger.info(f"{session.getName()} <- relayICECandidate")
        
        if not data:
            logger.error(f"{session.getName()} - relayICECandidate missing data")
            await websocket.send_json({
                "type": "error",
                "data": {"error": "relayICECandidate missing data"}
            })
            return

        # Check if session is properly joined to lobby with RTC peers
        with session.session_lock:
            if (lobby.id not in session.lobby_peers or 
                session.id not in lobby.sessions):
                logger.error(
                    f"{session.short}:{session.name} <- relayICECandidate - "
                    f"Not an RTC peer ({session.id})"
                )
                await websocket.send_json({
                    "type": "error",
                    "data": {"error": "Not joined to lobby"}
                })
                return
            
            session_peers = session.lobby_peers[lobby.id]

        # Validate peer_id
        peer_id = data.get("peer_id")
        if peer_id not in session_peers:
            logger.error(
                f"{session.getName()} <- relayICECandidate - "
                f"Not an RTC peer({peer_id}) in {session_peers}"
            )
            await websocket.send_json({
                "type": "error",
                "data": {"error": f"Target peer {peer_id} not found"}
            })
            return

        # Get candidate data
        candidate = data.get("candidate")

        # Prepare message for target peer
        message: Dict[str, Any] = {
            "type": "iceCandidate",
            "data": {
                "peer_id": session.id,
                "peer_name": session.name,
                "candidate": candidate,
            },
        }

        # Find target peer session and relay the message
        peer_session = lobby.getSession(peer_id)
        if not peer_session or not peer_session.ws:
            logger.warning(
                f"{session.getName()} - Live peer session {peer_id} "
                f"not found in lobby {lobby.getName()}."
            )
            return

        logger.info(
            f"{session.getName()} -> iceCandidate({peer_session.getName()})"
        )
        
        try:
            await peer_session.ws.send_json(message)
        except Exception as e:
            logger.warning(f"Failed to relay ICE candidate: {e}")

    @staticmethod
    async def handle_relay_session_description(
        websocket: WebSocket,
        session: "Session",
        lobby: "Lobby",
        data: Dict[str, Any]
    ) -> None:
        """
        Handle session description relay between peers.
        
        Args:
            websocket: The WebSocket connection
            session: The sender session
            lobby: The lobby context
            data: Message data containing peer_id and session_description
        """
        logger.info(f"{session.getName()} <- relaySessionDescription")
        
        if not data:
            logger.error(f"{session.getName()} - relaySessionDescription missing data")
            await websocket.send_json({
                "type": "error",
                "data": {"error": "relaySessionDescription missing data"}
            })
            return

        # Check if session is properly joined to lobby with RTC peers
        with session.session_lock:
            if (lobby.id not in session.lobby_peers or 
                session.id not in lobby.sessions):
                logger.error(
                    f"{session.short}:{session.name} <- relaySessionDescription - "
                    f"Not an RTC peer ({session.id})"
                )
                await websocket.send_json({
                    "type": "error",
                    "data": {"error": "Not joined to lobby"}
                })
                return

            lobby_peers = session.lobby_peers[lobby.id]

        # Validate peer_id
        peer_id = data.get("peer_id")
        if not peer_id:
            logger.error(f"{session.getName()} - relaySessionDescription missing peer_id")
            await websocket.send_json({
                "type": "error",
                "data": {"error": "relaySessionDescription missing peer_id"}
            })
            return

        if peer_id not in lobby_peers:
            logger.error(
                f"{session.getName()} <- relaySessionDescription - "
                f"Not an RTC peer({peer_id}) in {lobby_peers}"
            )
            await websocket.send_json({
                "type": "error",
                "data": {"error": f"Target peer {peer_id} not found"}
            })
            return

        # Find target peer session
        peer_session = lobby.getSession(peer_id)
        if not peer_session or not peer_session.ws:
            logger.warning(
                f"{session.getName()} - Live peer session {peer_id} "
                f"not found in lobby {lobby.getName()}."
            )
            return

        # Get session description data
        session_description = data.get("session_description")
        
        # Prepare message for target peer
        message: Dict[str, Any] = {
            "type": "sessionDescription",
            "data": {
                "peer_id": session.id,
                "peer_name": session.name,
                "session_description": session_description,
            },
        }

        logger.info(
            f"{session.getName()} -> sessionDescription({peer_session.getName()})"
        )
        
        try:
            await peer_session.ws.send_json(message)
        except Exception as e:
            logger.warning(f"Failed to relay session description: {e}")