BodyBalanceEvaluation/backend/devices/utils/socket_manager.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Socket连接管理工具
提供设备间Socket.IO通信的统一管理
"""

import threading
import time
import logging
from typing import Dict, Any, Optional, Callable
from datetime import datetime


class SocketManager:
    """Socket连接管理器"""
    
    def __init__(self, socketio=None):
        """
        初始化Socket管理器
        
        Args:
            socketio: SocketIO实例
        """
        self.socketio = socketio
        self.logger = logging.getLogger(f"{__name__}.SocketManager")
        self._namespaces = {}
        self._event_handlers = {}
        self._connection_stats = {}
        self._lock = threading.RLock()
        
    def set_socketio(self, socketio):
        """
        设置SocketIO实例
        
        Args:
            socketio: SocketIO实例
        """
        self.socketio = socketio
        
    def register_namespace(self, namespace: str, device_name: str):
        """
        注册设备命名空间
        
        Args:
            namespace: 命名空间路径（如 '/camera'）
            device_name: 设备名称
        """
        with self._lock:
            self._namespaces[namespace] = {
                'device_name': device_name,
                'registered_at': datetime.now().isoformat(),
                'active': True
            }
            self._connection_stats[namespace] = {
                'messages_sent': 0,
                'messages_failed': 0,
                'last_message_time': None,
                'connected_clients': 0
            }
            
        self.logger.info(f"注册设备命名空间: {namespace} -> {device_name}")
        
    def unregister_namespace(self, namespace: str):
        """
        注销设备命名空间
        
        Args:
            namespace: 命名空间路径
        """
        with self._lock:
            if namespace in self._namespaces:
                self._namespaces[namespace]['active'] = False
                self.logger.info(f"注销设备命名空间: {namespace}")
                
    def emit_to_namespace(self, namespace: str, event: str, data: Any, 
                         callback: Optional[Callable] = None) -> bool:
        """
        向指定命名空间发送数据
        
        Args:
            namespace: 命名空间路径
            event: 事件名称
            data: 数据
            callback: 回调函数
            
        Returns:
            bool: 发送是否成功
        """
        if not self.socketio:
            self.logger.warning("SocketIO未初始化")
            return False
            
        try:
            with self._lock:
                if namespace not in self._namespaces:
                    self.logger.warning(f"命名空间未注册: {namespace}")
                    return False
                    
                if not self._namespaces[namespace]['active']:
                    self.logger.warning(f"命名空间已停用: {namespace}")
                    return False
            
            # 发送数据
            if callback:
                self.socketio.emit(event, data, namespace=namespace, callback=callback)
            else:
                self.socketio.emit(event, data, namespace=namespace)
                
            # 更新统计
            with self._lock:
                stats = self._connection_stats[namespace]
                stats['messages_sent'] += 1
                stats['last_message_time'] = time.time()
                
            return True
            
        except Exception as e:
            self.logger.error(f"发送数据到 {namespace} 失败: {e}")
            with self._lock:
                if namespace in self._connection_stats:
                    self._connection_stats[namespace]['messages_failed'] += 1
            return False
            
    def broadcast_to_all(self, event: str, data: Any, 
                        exclude_namespaces: Optional[list] = None) -> Dict[str, bool]:
        """
        向所有活跃命名空间广播数据
        
        Args:
            event: 事件名称
            data: 数据
            exclude_namespaces: 排除的命名空间列表
            
        Returns:
            Dict[str, bool]: 各命名空间的发送结果
        """
        exclude_namespaces = exclude_namespaces or []
        results = {}
        
        with self._lock:
            active_namespaces = [
                ns for ns, info in self._namespaces.items() 
                if info['active'] and ns not in exclude_namespaces
            ]
            
        for namespace in active_namespaces:
            results[namespace] = self.emit_to_namespace(namespace, event, data)
            
        return results
        
    def register_event_handler(self, namespace: str, event: str, handler: Callable):
        """
        注册事件处理器
        
        Args:
            namespace: 命名空间
            event: 事件名称
            handler: 处理函数
        """
        if not self.socketio:
            self.logger.warning("SocketIO未初始化，无法注册事件处理器")
            return
            
        handler_key = f"{namespace}:{event}"
        self._event_handlers[handler_key] = handler
        
        # 注册到SocketIO
        @self.socketio.on(event, namespace=namespace)
        def wrapper(*args, **kwargs):
            try:
                return handler(*args, **kwargs)
            except Exception as e:
                self.logger.error(f"事件处理器 {handler_key} 执行失败: {e}")
                
        self.logger.info(f"注册事件处理器: {handler_key}")
        
    def get_namespace_stats(self, namespace: str) -> Optional[Dict[str, Any]]:
        """
        获取命名空间统计信息
        
        Args:
            namespace: 命名空间
            
        Returns:
            Optional[Dict[str, Any]]: 统计信息
        """
        with self._lock:
            if namespace in self._connection_stats:
                stats = self._connection_stats[namespace].copy()
                if stats['last_message_time']:
                    stats['last_message_ago'] = time.time() - stats['last_message_time']
                return stats
        return None
        
    def get_all_stats(self) -> Dict[str, Any]:
        """
        获取所有统计信息
        
        Returns:
            Dict[str, Any]: 所有统计信息
        """
        with self._lock:
            return {
                'namespaces': self._namespaces.copy(),
                'stats': {ns: self.get_namespace_stats(ns) for ns in self._namespaces},
                'total_namespaces': len(self._namespaces),
                'active_namespaces': len([ns for ns, info in self._namespaces.items() if info['active']])
            }
            
    def cleanup_namespace(self, namespace: str):
        """
        清理命名空间资源
        
        Args:
            namespace: 命名空间
        """
        with self._lock:
            if namespace in self._namespaces:
                del self._namespaces[namespace]
            if namespace in self._connection_stats:
                del self._connection_stats[namespace]
                
        # 清理事件处理器
        handlers_to_remove = [
            key for key in self._event_handlers.keys() 
            if key.startswith(f"{namespace}:")
        ]
        for key in handlers_to_remove:
            del self._event_handlers[key]
            
        self.logger.info(f"清理命名空间: {namespace}")
        
    def cleanup_all(self):
        """
        清理所有资源
        """
        with self._lock:
            namespaces = list(self._namespaces.keys())
            
        for namespace in namespaces:
            self.cleanup_namespace(namespace)
            
        self.logger.info("清理所有Socket资源")