relay/internal/simulation/websocket.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
模拟WebSocket通信模块，用于模拟与猪场主控的WebSocket通信
"""

import json
import time
import logging
from datetime import datetime
from .devices import DeviceManager, SimulatedDevice

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class SimulatedWebSocketServer:
    """模拟WebSocket服务器"""
    
    def __init__(self, config):
        """
        初始化模拟WebSocket服务器
        
        Args:
            config: 配置对象
        """
        self.config = config
        self.device_manager = DeviceManager()
        self._initialize_devices()
        logger.info("模拟WebSocket服务器初始化完成")
    
    def _initialize_devices(self):
        """初始化模拟设备"""
        logger.info("开始初始化模拟设备")
        # 根据配置创建模拟设备
        for device_config in self.config.simulation_devices:
            device = SimulatedDevice(
                device_id=device_config['id'],
                device_type=device_config['type'],
                status=device_config.get('status', 'stopped')
            )
            self.device_manager.add_device(device)
        logger.info(f"模拟设备初始化完成，共创建 {len(self.config.simulation_devices)} 个设备")
    
    def handle_message(self, message_str):
        """
        处理接收到的消息
        
        Args:
            message_str (str): JSON格式的消息字符串
            
        Returns:
            dict: 响应消息
        """
        logger.info(f"收到消息: {message_str}")
        
        try:
            message = json.loads(message_str)
            msg_type = message.get('type')
            logger.info(f"消息类型: {msg_type}")
            
            if msg_type == 'heartbeat':
                return self._handle_heartbeat(message)
            elif msg_type == 'command':
                command = message.get('command')
                logger.info(f"命令类型: {command}")
                if command == 'control_device':
                    return self._handle_control_device(message)
                elif command == 'query_device_status':
                    return self._handle_query_device_status(message)
                elif command == 'query_all_device_status':
                    return self._handle_query_all_device_status(message)
            
            # 未知消息类型
            result = {
                "type": "response",
                "status": "failed",
                "message": f"未知消息类型: {msg_type}",
                "timestamp": datetime.utcnow().isoformat() + 'Z'
            }
            logger.warning(f"未知消息类型: {msg_type}")
            return result
        except json.JSONDecodeError:
            result = {
                "type": "response",
                "status": "failed",
                "message": "无效的JSON格式",
                "timestamp": datetime.utcnow().isoformat() + 'Z'
            }
            logger.error("无效的JSON格式")
            return result
    
    def _handle_heartbeat(self, message):
        """
        处理心跳消息
        
        Args:
            message (dict): 心跳消息
            
        Returns:
            dict: 心跳响应
        """
        logger.info("处理心跳消息")
        result = {
            "type": "response",
            "command": "heartbeat",
            "status": "success",
            "message": "心跳响应",
            "timestamp": datetime.utcnow().isoformat() + 'Z'
        }
        logger.info("心跳响应已发送")
        return result
    
    def _handle_control_device(self, message):
        """
        处理设备控制命令
        
        Args:
            message (dict): 控制命令消息
            
        Returns:
            dict: 控制响应
        """
        data = message.get('data', {})
        device_id = data.get('device_id')
        action = data.get('action')
        
        logger.info(f"处理设备控制命令: 设备ID={device_id}, 动作={action}")
        result = self.device_manager.control_device(device_id, action)
        
        response = {
            "type": "response",
            "command": "control_device",
            "data": result,
            "timestamp": datetime.utcnow().isoformat() + 'Z'
        }
        logger.info(f"设备控制命令处理完成: {response}")
        return response
    
    def _handle_query_device_status(self, message):
        """
        处理查询设备状态命令
        
        Args:
            message (dict): 查询命令消息
            
        Returns:
            dict: 状态响应
        """
        data = message.get('data', {})
        device_id = data.get('device_id')
        
        logger.info(f"处理查询设备状态命令: 设备ID={device_id}")
        result = self.device_manager.query_device_status(device_id)
        
        response = {
            "type": "response",
            "command": "query_device_status",
            "data": result,
            "timestamp": datetime.utcnow().isoformat() + 'Z'
        }
        logger.info(f"设备状态查询完成: {response}")
        return response
    
    def _handle_query_all_device_status(self, message):
        """
        处理查询所有设备状态命令
        
        Args:
            message (dict): 查询命令消息
            
        Returns:
            dict: 状态响应
        """
        logger.info("处理查询所有设备状态命令")
        result = self.device_manager.query_all_device_status()
        
        response = {
            "type": "response",
            "command": "query_all_device_status",
            "data": result,
            "timestamp": datetime.utcnow().isoformat() + 'Z'
        }
        logger.info(f"所有设备状态查询完成，共 {len(result)} 个设备")
        return response