TermControlAgent/app/services/plnk_service.py

import asyncio
import time
import socket
import serial
from typing import Optional, Dict, Any
from app.core.device_manager import device_manager
from app.schemas.plnk import PLNKRequest, PLNKResponse, TCPConnectionInfo, PLNKConnectionInfo
from app.utils.log import get_logger

logger = get_logger(__name__)

class PLNKService:
    """PLNK协议服务类 - 实现自定义TCP/串口通信协议"""
    
    def __init__(self):
        self.connections = {}  # 存储连接
    
    async def _get_connection(self, device_id: str):
        """获取连接对象"""
        if device_id in self.connections:
            return self.connections[device_id]
        
        # 从设备管理器获取连接信息
        device = await device_manager.get_device(device_id)
        if not device:
            raise ValueError(f"设备 {device_id} 不存在")
        
        connection_info = device.connection_info
        connection_type = connection_info.get('connection_type', 'tcp')
        
        if connection_type == 'tcp':
            return await self._get_tcp_connection(device_id, connection_info)
        elif connection_type == 'serial':
            return await self._get_serial_connection(device_id, connection_info)
        else:
            raise ValueError(f"不支持的连接类型: {connection_type}")
    
    async def _get_tcp_connection(self, device_id: str, connection_info: Dict[str, Any]):
        """获取TCP连接"""
        try:
            tcp_info = connection_info.get('tcp_info', {})
            host = tcp_info.get('host')
            port = tcp_info.get('port')
            timeout = tcp_info.get('timeout', 30)
            
            # 创建TCP连接
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(timeout)
            sock.connect((host, port))
            
            # 存储连接
            self.connections[device_id] = {
                'type': 'tcp',
                'socket': sock,
                'host': host,
                'port': port
            }
            
            logger.info(f"TCP连接建立成功: {device_id} -> {host}:{port}")
            return self.connections[device_id]
            
        except Exception as e:
            logger.error(f"TCP连接失败: {device_id}, 错误: {e}")
            raise ValueError(f"TCP连接失败: {str(e)}")
    
    async def _get_serial_connection(self, device_id: str, connection_info: Dict[str, Any]):
        """获取串口连接"""
        try:
            serial_info = connection_info.get('serial_info', {})
            port = serial_info.get('port')
            baudrate = serial_info.get('baudrate', 115200)
            timeout = serial_info.get('timeout', 5)
            
            # 创建串口连接
            ser = serial.Serial(
                port=port,
                baudrate=baudrate,
                timeout=timeout
            )
            
            # 存储连接
            self.connections[device_id] = {
                'type': 'serial',
                'serial': ser,
                'port': port
            }
            
            logger.info(f"串口连接建立成功: {device_id} -> {port}")
            return self.connections[device_id]
            
        except Exception as e:
            logger.error(f"串口连接失败: {device_id}, 错误: {e}")
            raise ValueError(f"串口连接失败: {str(e)}")
    
    async def send_plnk_data(self, device_id: str, data: str, timeout: int = 30, 
                            wait_response: bool = True) -> PLNKResponse:
        """发送PLNK协议数据"""
        start_time = time.time()
        
        try:
            connection = await self._get_connection(device_id)
            connection_type = connection['type']
            
            # 将十六进制字符串转换为字节
            try:
                data_bytes = bytes.fromhex(data.replace(' ', ''))
            except ValueError as e:
                raise ValueError(f"无效的十六进制数据: {data}")
            
            sent_data = data
            received_data = ""
            
            if connection_type == 'tcp':
                sock = connection['socket']
                
                # 发送数据
                sock.send(data_bytes)
                
                if wait_response:
                    # 接收响应
                    sock.settimeout(timeout)
                    response_bytes = sock.recv(1024)
                    received_data = response_bytes.hex().upper()
                    
            elif connection_type == 'serial':
                ser = connection['serial']
                
                # 清空缓冲区
                ser.reset_input_buffer()
                ser.reset_output_buffer()
                
                # 发送数据
                ser.write(data_bytes)
                ser.flush()
                
                if wait_response:
                    # 接收响应
                    start_timeout = time.time()
                    response_bytes = b""
                    while time.time() - start_timeout < timeout:
                        if ser.in_waiting > 0:
                            response_bytes += ser.read(ser.in_waiting)
                            # 可以根据协议特点判断是否接收完整
                            if len(response_bytes) > 0:
                                break
                        await asyncio.sleep(0.1)
                    
                    received_data = response_bytes.hex().upper()
            
            execution_time = time.time() - start_time
            
            # 解析响应数据（这里可以根据具体协议实现）
            parsed_data = self._parse_plnk_response(received_data)
            
            success = len(received_data) > 0
            
            logger.info(f"PLNK数据发送完成: {device_id} -> {data}, 成功: {success}")
            
            return PLNKResponse(
                success=success,
                sent_data=sent_data,
                received_data=received_data,
                parsed_data=parsed_data,
                execution_time=execution_time
            )
            
        except Exception as e:
            execution_time = time.time() - start_time
            logger.error(f"PLNK数据发送失败: {device_id} -> {data}, 错误: {e}")
            
            return PLNKResponse(
                success=False,
                sent_data=data,
                received_data="",
                parsed_data=None,
                execution_time=execution_time
            )
    
    def _parse_plnk_response(self, hex_data: str) -> Optional[Dict[str, Any]]:
        """解析PLNK协议响应数据"""
        if not hex_data:
            return None
        
        try:
            # 这里可以根据具体的PLNK协议格式进行解析
            # 示例解析格式（需要根据实际协议调整）
            parsed = {
                "raw_data": hex_data,
                "length": len(hex_data) // 2,
                "timestamp": time.time()
            }
            
            # 如果数据长度足够，可以解析更多字段
            if len(hex_data) >= 4:
                parsed["header"] = hex_data[:4]
                parsed["payload"] = hex_data[4:-2] if len(hex_data) > 6 else hex_data[4:]
                if len(hex_data) >= 6:
                    parsed["checksum"] = hex_data[-2:]
            
            return parsed
            
        except Exception as e:
            logger.error(f"PLNK响应解析失败: {hex_data}, 错误: {e}")
            return None
    
    async def test_connection(self, device_id: str) -> bool:
        """测试连接"""
        try:
            connection = await self._get_connection(device_id)
            
            # 发送测试数据
            test_data = "0101"  # 简单的测试数据
            result = await self.send_plnk_data(device_id, test_data, timeout=5)
            
            return result.success
            
        except Exception as e:
            logger.error(f"PLNK连接测试失败: {device_id}, 错误: {e}")
            return False
    
    async def close_connection(self, device_id: str) -> bool:
        """关闭连接"""
        try:
            if device_id in self.connections:
                connection = self.connections[device_id]
                
                if connection['type'] == 'tcp':
                    connection['socket'].close()
                elif connection['type'] == 'serial':
                    connection['serial'].close()
                
                del self.connections[device_id]
                logger.info(f"PLNK连接已关闭: {device_id}")
                return True
            return False
        except Exception as e:
            logger.error(f"关闭PLNK连接失败: {device_id}, 错误: {e}")
            return False
    
    async def get_connection_status(self, device_id: str) -> Dict[str, Any]:
        """获取连接状态"""
        try:
            if device_id in self.connections:
                connection = self.connections[device_id]
                return {
                    "device_id": device_id,
                    "connected": True,
                    "type": connection['type'],
                    "status": "online"
                }
            else:
                return {
                    "device_id": device_id,
                    "connected": False,
                    "type": "unknown",
                    "status": "offline"
                }
        except Exception as e:
            logger.error(f"获取连接状态失败: {device_id}, 错误: {e}")
            return {
                "device_id": device_id,
                "connected": False,
                "type": "unknown",
                "status": "error"
            }