shelfmark/websocket_manager.py

"""WebSocket manager for real-time status updates."""

import logging
import threading
from typing import Optional, Dict, Any, Callable, List

from flask_socketio import SocketIO, emit

logger = logging.getLogger(__name__)

class WebSocketManager:
    """Manages WebSocket connections and broadcasts."""
    
    def __init__(self):
        self.socketio: Optional[SocketIO] = None
        self._enabled = False
        self._connection_count = 0
        self._connection_lock = threading.Lock()
        self._on_first_connect_callbacks: List[Callable[[], None]] = []
        self._on_all_disconnect_callbacks: List[Callable[[], None]] = []
        self._needs_rewarm = False  # Flag to trigger warmup callbacks on next connect
        
    def init_app(self, app, socketio: SocketIO):
        """Initialize the WebSocket manager with Flask-SocketIO instance."""
        self.socketio = socketio
        self._enabled = True
        logger.info("WebSocket manager initialized")
    
    def register_on_first_connect(self, callback: Callable[[], None]):
        """Register a callback to be called when the first client connects.
        
        This is useful for warming up resources (like the Cloudflare bypasser)
        when a user starts using the web UI.
        """
        self._on_first_connect_callbacks.append(callback)
        logger.debug(f"Registered on_first_connect callback: {callback.__name__}")
    
    def register_on_all_disconnect(self, callback: Callable[[], None]):
        """Register a callback to be called when all clients disconnect.

        This can be used to trigger cleanup or resource release.
        """
        self._on_all_disconnect_callbacks.append(callback)
        logger.debug(f"Registered on_all_disconnect callback: {callback.__name__}")

    def request_warmup_on_next_connect(self):
        """Request that warmup callbacks be triggered on the next client connect.

        This is used when resources (like the Cloudflare bypasser) shut down due to
        inactivity while clients are still connected. The next connect event should
        trigger warmup even though it's not technically the "first" connection.
        """
        with self._connection_lock:
            self._needs_rewarm = True
            logger.debug("Warmup requested for next client connect")
    
    def client_connected(self):
        """Track a new client connection. Call this from the connect event handler."""
        with self._connection_lock:
            was_zero = self._connection_count == 0
            needs_rewarm = self._needs_rewarm
            self._connection_count += 1
            current_count = self._connection_count
            # Clear rewarm flag if we're going to trigger warmup
            if was_zero or needs_rewarm:
                self._needs_rewarm = False

        logger.debug(f"Client connected. Active connections: {current_count}")

        # Trigger warmup callbacks if this is the first connection OR if rewarm was requested
        # (rewarm is requested when bypasser shuts down due to idle while clients are connected)
        if was_zero or needs_rewarm:
            reason = "First client connected" if was_zero else "Rewarm requested after idle shutdown"
            logger.info(f"{reason}, triggering warmup callbacks...")
            for callback in self._on_first_connect_callbacks:
                try:
                    # Run callbacks in a separate thread to not block the connection
                    thread = threading.Thread(target=callback, daemon=True)
                    thread.start()
                except Exception as e:
                    logger.error(f"Error in on_first_connect callback {callback.__name__}: {e}")
    
    def client_disconnected(self):
        """Track a client disconnection. Call this from the disconnect event handler."""
        with self._connection_lock:
            self._connection_count = max(0, self._connection_count - 1)
            current_count = self._connection_count
            is_now_zero = current_count == 0
        
        logger.debug(f"Client disconnected. Active connections: {current_count}")
        
        # If all clients have disconnected, trigger cleanup callbacks
        if is_now_zero:
            logger.info("All clients disconnected, triggering disconnect callbacks...")
            for callback in self._on_all_disconnect_callbacks:
                try:
                    callback()
                except Exception as e:
                    logger.error(f"Error in on_all_disconnect callback {callback.__name__}: {e}")
    
    def get_connection_count(self) -> int:
        """Get the current number of active WebSocket connections."""
        with self._connection_lock:
            return self._connection_count
    
    def has_active_connections(self) -> bool:
        """Check if there are any active WebSocket connections."""
        return self.get_connection_count() > 0
        
    def is_enabled(self) -> bool:
        """Check if WebSocket is enabled and ready."""
        return self._enabled and self.socketio is not None
    
    def broadcast_status_update(self, status_data: Dict[str, Any]):
        """Broadcast status update to all connected clients."""
        if not self.is_enabled():
            return
            
        try:
            # When calling socketio.emit() outside event handlers, it broadcasts by default
            self.socketio.emit('status_update', status_data)
            logger.debug(f"Broadcasted status update to all clients")
        except Exception as e:
            logger.error(f"Error broadcasting status update: {e}")
    
    def broadcast_download_progress(self, book_id: str, progress: float, status: str):
        """Broadcast download progress update for a specific book."""
        if not self.is_enabled():
            return
            
        try:
            data = {
                'book_id': book_id,
                'progress': progress,
                'status': status
            }
            # When calling socketio.emit() outside event handlers, it broadcasts by default
            self.socketio.emit('download_progress', data)
            logger.debug(f"Broadcasted progress for book {book_id}: {progress}%")
        except Exception as e:
            logger.error(f"Error broadcasting download progress: {e}")
    
    def broadcast_notification(self, message: str, notification_type: str = 'info'):
        """Broadcast a notification message to all clients."""
        if not self.is_enabled():
            return
            
        try:
            data = {
                'message': message,
                'type': notification_type
            }
            # When calling socketio.emit() outside event handlers, it broadcasts by default
            self.socketio.emit('notification', data)
            logger.debug(f"Broadcasted notification: {message}")
        except Exception as e:
            logger.error(f"Error broadcasting notification: {e}")

# Global WebSocket manager instance
ws_manager = WebSocketManager()