evolution-client-python/evolutionapi/services/websocket.py

import logging
import time
from typing import Callable

import socketio

from ..models.websocket import WebSocketConfig, WebSocketInfo


class WebSocketService:
    def __init__(self, client):
        self.client = client

    def set_websocket(self, instance_id: str, config: WebSocketConfig, instance_token: str):
        """
        Configure WebSocket settings for an instance

        Args:
            instance_id (str): The instance ID
            config (WebSocketConfig): The WebSocket configuration
            instance_token (str): The instance token

        Returns:
            dict: The response from the API
        """
        return self.client.post(
            f"websocket/set/{instance_id}", data=config.__dict__, instance_token=instance_token
        )

    def find_websocket(self, instance_id: str, instance_token: str) -> WebSocketInfo:
        """
        Get WebSocket settings for an instance

        Args:
            instance_id (str): The instance ID
            instance_token (str): The instance token

        Returns:
            WebSocketInfo: The WebSocket information
        """
        response = self.client.get(f"websocket/find/{instance_id}", instance_token=instance_token)
        return WebSocketInfo(**response)


class WebSocketManager:
    def __init__(
        self,
        base_url: str,
        instance_id: str,
        api_token: str,
        max_retries: int = 5,
        retry_delay: float = 1.0,
    ):
        """
        Initialize the WebSocket manager

        Args:
            base_url (str): Base URL of the API
            instance_id (str): Instance ID
            api_token (str): API authentication token
            max_retries (int): Maximum number of reconnection attempts
            retry_delay (float): Initial delay between attempts in seconds
        """
        self.base_url = base_url.rstrip("/")
        self.instance_id = instance_id
        self.api_token = api_token
        self.max_retries = max_retries
        self.retry_delay = retry_delay
        self.retry_count = 0
        self.should_reconnect = True

        # Socket.IO configuration
        self.sio = socketio.Client(
            ssl_verify=False,  # For local development
            logger=False,
            engineio_logger=False,
            request_timeout=30,
        )

        # Configure class logger to INFO
        self.logger = logging.getLogger(__name__)
        self.logger.setLevel(logging.INFO)

        # Dictionary to store registered handlers
        self._handlers = {}

        # Configure event handlers
        self.sio.on("connect", self._on_connect)
        self.sio.on("disconnect", self._on_disconnect)
        self.sio.on("error", self._on_error)

        # Register global handler in instance-specific namespace
        self.sio.on("*", self._handle_event, namespace=f"/{self.instance_id}")

    def _on_connect(self):
        """Handler for connection event"""
        self.logger.info("Socket.IO connected")
        self.retry_count = 0  # Reset retry counter after successful connection

    def _on_disconnect(self):
        """Handler for disconnection event"""
        self.logger.warning(
            f"Socket.IO disconnected. Attempt {self.retry_count + 1}/{self.max_retries}"
        )
        if self.should_reconnect and self.retry_count < self.max_retries:
            self._attempt_reconnect()
        else:
            self.logger.error("Maximum number of reconnection attempts reached")

    def _on_error(self, error):
        """Handler for error events"""
        self.logger.error(f"Socket.IO error: {str(error)}", exc_info=True)

    def _attempt_reconnect(self):
        """Attempt to reconnect with exponential backoff"""
        try:
            delay = self.retry_delay * (2**self.retry_count)  # Exponential backoff
            self.logger.info(f"Attempting to reconnect in {delay:.2f} seconds...")
            time.sleep(delay)
            self.connect()
            self.retry_count += 1
        except Exception as e:
            self.logger.error(f"Error during reconnection attempt: {str(e)}", exc_info=True)
            if self.retry_count < self.max_retries:
                self._attempt_reconnect()
            else:
                self.logger.error("All reconnection attempts failed")

    def _handle_event(self, event, *args):
        """Global handler for all events"""
        # Only process registered events
        if event in self._handlers:
            self.logger.debug(f"Event received in namespace /{self.instance_id}: {event}")
            self.logger.debug(f"Event data: {args}")

            try:
                # Extract event data
                raw_data = args[0] if args else {}

                # Ensure we're passing the correct object to the callback
                if isinstance(raw_data, dict):
                    self.logger.debug(f"Calling handler for {event} with data: {raw_data}")
                    self._handlers[event](raw_data)
                else:
                    self.logger.error(f"Invalid data received for event {event}: {raw_data}")
            except Exception as e:
                self.logger.error(f"Error processing event {event}: {str(e)}", exc_info=True)

    def connect(self):
        """Connect to Socket.IO server"""
        try:
            # Connect only to instance namespace with authentication header
            self.sio.connect(
                f"{self.base_url}?apikey={self.api_token}",
                transports=["websocket"],
                namespaces=[f"/{self.instance_id}"],
                wait_timeout=30,
            )

            # Join instance-specific room
            self.sio.emit(
                "subscribe", {"instance": self.instance_id}, namespace=f"/{self.instance_id}"
            )

        except Exception as e:
            self.logger.error(f"Error connecting to Socket.IO: {str(e)}", exc_info=True)
            raise

    def disconnect(self):
        """Disconnect from Socket.IO server"""
        self.should_reconnect = False  # Prevent reconnection attempts
        if self.sio.connected:
            self.sio.disconnect()

    def on(self, event: str, callback: Callable):
        """
        Register a callback for a specific event

        Args:
            event (str): Event name
            callback (Callable): Function to be called when the event occurs
        """
        self._handlers[event] = callback