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
from core.error_handling import (
    with_webrtc_error_handling,
    WebRTCError,
    ErrorSeverity,
    error_handler
)

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
    @with_webrtc_error_handling
    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
    @with_webrtc_error_handling
    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}")

    @staticmethod
    @with_webrtc_error_handling
    async def handle_add_peer(
        session: "Session",
        peer_session: "Session", 
        lobby: "Lobby"
    ) -> None:
        """
        Handle adding WebRTC peer connections between two sessions in a lobby.
        
        Args:
            session: The session joining the lobby
            peer_session: The existing peer session in the lobby
            lobby: The lobby context
        """
        # Only establish WebRTC connections if at least one has media
        if session.has_media or peer_session.has_media:
            # Add peer_session to session's peer list
            with session.session_lock:
                if lobby.id not in session.lobby_peers:
                    session.lobby_peers[lobby.id] = []
                session.lobby_peers[lobby.id].append(peer_session.id)

            # Add session to peer_session's peer list
            with peer_session.session_lock:
                if lobby.id not in peer_session.lobby_peers:
                    peer_session.lobby_peers[lobby.id] = []
                peer_session.lobby_peers[lobby.id].append(session.id)

            # Notify existing peer about new peer (they should not create offer)
            logger.info(
                f"{session.getName()} -> {peer_session.getName()}:addPeer("
                f"{session.getName()}, {lobby.getName()}, should_create_offer=False, "
                f"has_media={session.has_media})"
            )
            try:
                if peer_session.ws:
                    await peer_session.ws.send_json(
                        {
                            "type": "addPeer",
                            "data": {
                                "peer_id": session.id,
                                "peer_name": session.name,
                                "has_media": session.has_media,
                                "should_create_offer": False,
                            },
                        }
                    )
            except Exception as e:
                logger.warning(
                    f"Failed to send addPeer to {peer_session.getName()}: {e}"
                )

            # Notify new session about existing peer (they should create offer)
            logger.info(
                f"{session.getName()} -> {session.getName()}:addPeer("
                f"{peer_session.getName()}, {lobby.getName()}, should_create_offer=True, "
                f"has_media={peer_session.has_media})"
            )
            try:
                if session.ws:
                    await session.ws.send_json(
                        {
                            "type": "addPeer",
                            "data": {
                                "peer_id": peer_session.id,
                                "peer_name": peer_session.name,
                                "has_media": peer_session.has_media,
                                "should_create_offer": True,
                            },
                        }
                    )
            except Exception as e:
                logger.warning(f"Failed to send addPeer to {session.getName()}: {e}")
        else:
            logger.info(
                f"{session.getName()} - Skipping WebRTC connection with "
                f"{peer_session.getName()} (neither has media: "
                f"self={session.has_media}, peer={peer_session.has_media})"
            )

    @staticmethod
    @with_webrtc_error_handling
    async def handle_remove_peer(
        session: "Session",
        peer_session: "Session",
        lobby: "Lobby"
    ) -> None:
        """
        Handle removing WebRTC peer connections between two sessions.
        
        Args:
            session: The session leaving the lobby
            peer_session: The peer session to disconnect from
            lobby: The lobby context
        """
        # Notify peer about session removal
        if peer_session.ws:
            logger.info(
                f"{peer_session.getName()} <- remove_peer({session.getName()})"
            )
            try:
                await peer_session.ws.send_json({
                    "type": "removePeer",
                    "data": {"peer_name": session.name, "peer_id": session.id},
                })
            except Exception as e:
                logger.warning(
                    f"Failed to send removePeer to {peer_session.getName()}: {e}"
                )
        else:
            logger.warning(
                f"{session.getName()} <- part({lobby.getName()}) - "
                f"No WebSocket connection for {peer_session.getName()}. Skipping."
            )

        # Remove from peer's lobby_peers
        with peer_session.session_lock:
            if (lobby.id in peer_session.lobby_peers and 
                session.id in peer_session.lobby_peers[lobby.id]):
                peer_session.lobby_peers[lobby.id].remove(session.id)

        # Notify session about peer removal
        if session.ws:
            logger.info(
                f"{session.getName()} <- remove_peer({peer_session.getName()})"
            )
            try:
                await session.ws.send_json({
                    "type": "removePeer",
                    "data": {
                        "peer_name": peer_session.name,
                        "peer_id": peer_session.id,
                    },
                })
            except Exception as e:
                logger.warning(
                    f"Failed to send removePeer to {session.getName()}: {e}"
                )
        else:
            logger.error(
                f"{session.getName()} <- part({lobby.getName()}) - No WebSocket connection."
            )