TermControlAgent/app/core/websocket/send_controller.py

"""
WebSocket发送控制器
专门负责统一的消息发送逻辑
遵循单一职责原则
"""
import asyncio
from typing import Dict, Any, Optional
from datetime import datetime
from app.core.websocket.channel import ChannelMessage, WebSocketChannel
from app.core.websocket.client import WebSocketClient
from app.utils.structured_log import get_structured_logger, LogLevel

logger = get_structured_logger(__name__, LogLevel.INFO)


class WebSocketSendController:
    """WebSocket发送控制器
    
    单一职责：只负责统一的消息发送逻辑
    - 统一发送循环
    - 消息发送到WebSocket
    - 优先级控制
    - 发送任务管理
    """
    
    def __init__(self, client: WebSocketClient, heartbeat_channel: WebSocketChannel, send_channel: WebSocketChannel):
        self.client = client
        self.client_name = client.name
        self.heartbeat_channel = heartbeat_channel
        self.send_channel = send_channel
        self._send_task: Optional[asyncio.Task] = None
        self._running = False
        self._created_at = datetime.now()
        
        logger.info(f"WebSocket发送控制器初始化: {self.client_name}")
    
    async def start(self) -> bool:
        """启动发送控制器"""
        if self._running:
            logger.warning(f"WebSocket发送控制器 {self.client_name} 已在运行")
            return True
        
        try:
            self._running = True
            self._send_task = asyncio.create_task(self._unified_send_loop())
            logger.info(f"WebSocket发送控制器启动成功: {self.client_name}")
            return True
            
        except Exception as e:
            self._running = False
            logger.error(f"WebSocket发送控制器启动失败: {self.client_name} - {e}")
            return False
    
    async def stop(self) -> bool:
        """停止发送控制器"""
        if not self._running:
            logger.warning(f"WebSocket发送控制器 {self.client_name} 未在运行")
            return True
        
        try:
            self._running = False
            
            if self._send_task and not self._send_task.done():
                self._send_task.cancel()
                try:
                    await self._send_task
                except asyncio.CancelledError:
                    pass
            
            logger.info(f"WebSocket发送控制器停止成功: {self.client_name}")
            return True
            
        except Exception as e:
            logger.error(f"WebSocket发送控制器停止失败: {self.client_name} - {e}")
            return False
    
    async def _unified_send_loop(self):
        """统一发送循环
        
        单一职责：只负责从Channel获取消息并发送到WebSocket
        """
        logger.info(f"WebSocket发送控制器开始统一发送循环: {self.client_name}")
        
        try:
            while self._running and self.client.is_connected:
                try:
                    # 检查Channel是否存在且已连接
                    if not self.heartbeat_channel or not self.send_channel:
                        logger.warning(f"Channel不存在，等待重试: {self.client_name}")
                        await asyncio.sleep(1)
                        continue
                    
                    if not self.heartbeat_channel.is_connected or not self.send_channel.is_connected:
                        logger.warning(f"Channel未连接，等待重试: {self.client_name}")
                        await asyncio.sleep(1)
                        continue
                    
                    # 优先级发送：先发送业务数据，再发送心跳
                    message_sent = False
                    
                    # 1. 优先发送业务数据
                    if self.send_channel.queue_size > 0:
                        try:
                            message = await self.send_channel.receive_message()
                            if message:
                                success = await self._send_message(message)
                                if success:
                                    logger.debug(f"WebSocket发送控制器发送业务消息成功: {self.client_name} -> {message.type}")
                                    message_sent = True
                                else:
                                    logger.warning(f"WebSocket发送控制器发送业务消息失败: {self.client_name} -> {message.type}")
                        except Exception as e:
                            logger.error(f"WebSocket发送控制器发送业务消息异常: {self.client_name} - {e}")
                    
                    # 2. 如果没有业务数据，发送心跳
                    if not message_sent and self.heartbeat_channel.queue_size > 0:
                        try:
                            message = await self.heartbeat_channel.receive_message()
                            if message:
                                success = await self._send_message(message)
                                if success:
                                    logger.debug(f"WebSocket发送控制器发送心跳消息成功: {self.client_name} -> {message.type}")
                                    message_sent = True
                                else:
                                    logger.warning(f"WebSocket发送控制器发送心跳消息失败: {self.client_name} -> {message.type}")
                        except Exception as e:
                            logger.error(f"WebSocket发送控制器发送心跳消息异常: {self.client_name} - {e}")
                    
                    # 3. 如果没有消息需要发送，短暂等待
                    if not message_sent:
                        await asyncio.sleep(0.1)  # 100ms
                    
                except asyncio.CancelledError:
                    logger.info(f"WebSocket发送控制器发送循环被取消: {self.client_name}")
                    break
                except Exception as e:
                    logger.error(f"WebSocket发送控制器发送循环异常: {self.client_name} - {e}")
                    await asyncio.sleep(1)  # 异常时等待1秒再重试
            
            logger.info(f"WebSocket发送控制器统一发送循环结束: {self.client_name}")
            
        except Exception as e:
            logger.error(f"WebSocket发送控制器发送循环严重异常: {self.client_name} - {e}")
    
    async def _send_message(self, message: ChannelMessage) -> bool:
        """发送消息到WebSocket
        
        单一职责：只负责将消息发送到WebSocket客户端
        """
        try:
            if not self.client.is_connected:
                logger.warning(f"WebSocket客户端未连接，无法发送消息: {self.client_name}")
                return False
            
            # 在发送时添加Type字段
            send_data = {
                "Type": message.type,
                **message.data
            }
            
            # 使用序列化器序列化消息
            from app.core.websocket.serializer import websocket_serializer
            payload = websocket_serializer.serialize(send_data)
            
            # 发送数据到WebSocket
            success = await self.client.send_raw(payload)
            
            if success:
                logger.debug(f"WebSocket发送控制器消息发送成功: {self.client_name} -> {message.type}")
            else:
                logger.warning(f"WebSocket发送控制器消息发送失败: {self.client_name} -> {message.type}")
            
            return success
            
        except Exception as e:
            logger.error(f"WebSocket发送控制器发送消息异常: {self.client_name} -> {message.type} - {e}")
            return False
    
    def is_running(self) -> bool:
        """检查发送控制器是否在运行"""
        return self._running
    
    def get_stats(self) -> Dict[str, Any]:
        """获取发送控制器统计信息"""
        return {
            "client_name": self.client_name,
            "is_running": self._running,
            "is_connected": self.client.is_connected,
            "created_at": self._created_at.isoformat(),
            "heartbeat_channel": self.heartbeat_channel.name,
            "send_channel": self.send_channel.name
        }