chat_agent/websocket/events.py

"""
WebSocket event registration and handlers.

This module registers all WebSocket events with Flask-SocketIO and provides
the main entry point for WebSocket functionality.
"""

import logging
from flask_socketio import SocketIO

from .chat_websocket import ChatWebSocketHandler


logger = logging.getLogger(__name__)


def register_websocket_events(socketio: SocketIO, websocket_handler: ChatWebSocketHandler) -> None:
    """
    Register all WebSocket events with Flask-SocketIO.
    
    Args:
        socketio: Flask-SocketIO instance
        websocket_handler: WebSocket handler instance
    """
    
    @socketio.on('connect')
    def handle_connect(auth=None):
        """Handle WebSocket connection."""
        logger.debug("WebSocket connect event received")
        return websocket_handler.handle_connect(auth)
    
    @socketio.on('disconnect')
    def handle_disconnect():
        """Handle WebSocket disconnection."""
        logger.debug("WebSocket disconnect event received")
        websocket_handler.handle_disconnect()
    
    @socketio.on('message')
    def handle_message(data):
        """Handle chat message."""
        logger.debug("WebSocket message event received")
        websocket_handler.handle_message(data)
    
    @socketio.on('language_switch')
    def handle_language_switch(data):
        """Handle language switch request."""
        logger.debug("WebSocket language_switch event received")
        websocket_handler.handle_language_switch(data)
    
    @socketio.on('typing_start')
    def handle_typing_start(data):
        """Handle typing indicator start."""
        logger.debug("WebSocket typing_start event received")
        websocket_handler.handle_typing_start(data)
    
    @socketio.on('typing_stop')
    def handle_typing_stop(data):
        """Handle typing indicator stop."""
        logger.debug("WebSocket typing_stop event received")
        websocket_handler.handle_typing_stop(data)
    
    @socketio.on('ping')
    def handle_ping(data):
        """Handle ping for connection health check."""
        logger.debug("WebSocket ping event received")
        websocket_handler.handle_ping(data)
    
    @socketio.on('get_session_info')
    def handle_get_session_info(data):
        """Handle session info request."""
        logger.debug("WebSocket get_session_info event received")
        websocket_handler.handle_get_session_info(data)
    
    # Error handlers
    @socketio.on_error_default
    def default_error_handler(e):
        """Handle WebSocket errors."""
        logger.error(f"WebSocket error: {e}")
        return {'error': 'Internal server error', 'code': 'WEBSOCKET_ERROR'}
    
    @socketio.on_error()
    def error_handler(e):
        """Handle general WebSocket errors."""
        logger.error(f"WebSocket general error: {e}")
        return {'error': 'Connection error', 'code': 'CONNECTION_ERROR'}
    
    logger.info("WebSocket events registered successfully")