Tolan/backend/websocket_handler.py

"""
WebSocket处理器模块
"""
import json
import asyncio
import time
from typing import Dict, Optional
from fastapi import WebSocket, WebSocketDisconnect
from loguru import logger
from asyncio import Queue

from config import settings
from connection_pool import connection_pool
from error_codes import ErrorCode, get_error_message
from performance import performance_monitor


class StreamingSession:
    """流式会话处理"""
    
    def __init__(self, websocket: WebSocket, session_id: str):
        self.websocket = websocket
        self.session_id = session_id
        self.doubao_client = None
        self.is_streaming = False
        self.is_recording = False
        
        # 音频数据队列
        self.to_doubao_queue: Queue = Queue()
        self.to_client_queue: Queue = Queue()
        
        # 性能指标
        self.start_time = 0
        self.audio_received_bytes = 0
        self.audio_sent_bytes = 0
        
        # 文本收集（用于日志）
        self.llm_text_chunks = []  # 收集LLM流式文本
        self.asr_text = ''  # 最后一次ASR识别结果
        
        # 对话轮次控制
        self.response_ended = False  # 标记当前轮次是否已结束（避免重复发送结束消息）
        
        # 心跳检测
        self.last_ping_time = time.time()
        self.heartbeat_timeout = 60  # 60秒无心跳则断开
    
    async def start(self):
        """启动会话"""
        try:
            # 从连接池获取豆包连接
            self.doubao_client = await connection_pool.acquire(timeout=5)
            
            if not self.doubao_client:
                await self.send_error(
                    ErrorCode.CONNECTION_POOL_EXHAUSTED,
                    "连接池已耗尽"
                )
                return False
            
            logger.info(f"会话 {self.session_id} 已启动")
            return True
            
        except Exception as e:
            logger.error(f"启动会话失败: {e}")
            await self.send_error(ErrorCode.INTERNAL_ERROR, f"启动会话失败: {str(e)}")
            return False
    
    async def process_stream(self):
        """处理流式数据"""
        try:
            self.is_streaming = True
            
            # 并发处理多个任务
            await asyncio.gather(
                self.receive_from_client(),
                self.send_to_doubao(),
                self.receive_from_doubao(),
                self.send_to_client(),
                self.check_heartbeat()
            )
            
        except Exception as e:
            logger.error(f"流式处理错误: {e}")
            await self.send_error(ErrorCode.INTERNAL_ERROR, f"流式处理错误: {str(e)}")
        finally:
            self.is_streaming = False
    
    async def receive_from_client(self):
        """接收客户端数据"""
        try:
            while self.is_streaming:
                message = await self.websocket.receive()
                
                if 'text' in message:
                    # 控制消息
                    await self.handle_control_message(message['text'])
                    
                elif 'bytes' in message:
                    # 音频数据
                    if self.is_recording:
                        audio_data = message['bytes']
                        chunk_size = len(audio_data)
                        self.audio_received_bytes += chunk_size
                        performance_monitor.record_audio_received(chunk_size)
                        # 只在DEBUG模式下记录每个音频块
                        logger.debug(f"收到音频块: {chunk_size} bytes, 累计: {self.audio_received_bytes} bytes")
                        await self.to_doubao_queue.put(audio_data)
                        
        except WebSocketDisconnect:
            logger.info(f"客户端 {self.session_id} 断开连接")
            self.is_streaming = False
        except Exception as e:
            logger.error(f"接收客户端数据错误: {e}")
            self.is_streaming = False
    
    async def send_to_doubao(self):
        """转发音频到豆包API"""
        try:
            while self.is_streaming:
                audio_chunk = await self.to_doubao_queue.get()
                
                if self.doubao_client:
                    # 转发给豆包API
                    chunk_size = len(audio_chunk)
                    await self.doubao_client.send_audio(audio_chunk)
                    # 使用实例变量中的累计值（在handle_start_recording中重置）
                    logger.debug(f"转发音频到豆包: {chunk_size} bytes")
                    
        except Exception as e:
            logger.error(f"转发豆包错误: {e}")
    
    async def receive_from_doubao(self):
        """接收豆包API响应"""
        try:
            # 持续接收豆包响应（支持多轮对话）
            while self.is_streaming:
                # 持续接收，不等待录音状态
                
                # 接收当前这一轮的音频
                if not self.doubao_client:
                    await asyncio.sleep(0.1)
                    continue
                
                try:
                    # 接收豆包音频数据，设置响应超时
                    audio_data = await asyncio.wait_for(
                        self.doubao_client.receive_audio(), 
                        timeout=settings.DOUBAO_RESPONSE_TIMEOUT
                    )
                    
                    # 检查返回值类型
                    if audio_data is False:
                        # False表示当前轮次音频结束
                        # 检查是否已经发送过结束消息（避免重复）
                        if self.response_ended:
                            logger.debug("已发送过结束消息，忽略重复的结束事件")
                            continue
                        
                        logger.info("当前轮次音频结束，等待下一轮")
                        
                        # 如果发送过音频，立即发送结束消息
                        if self.audio_sent_bytes > 0:
                            # 计算本轮对话的延迟
                            response_latency = int((time.time() - self.start_time) * 1000) if self.start_time > 0 else 0
                            
                            # 打印完整的LLM回复文本
                            full_llm_text = ''.join(self.llm_text_chunks)
                            logger.info(f"🔊 [AI回复完成] 音频: {self.audio_sent_bytes / 1024:.1f}KB, 延迟: {response_latency}ms")
                            if full_llm_text:
                                logger.info(f"💬 [完整回复] {full_llm_text}")
                            
                            await self.send_control_message({
                                'type': 'ai_response_end',
                                'payload': {
                                    'responseId': f'resp_{int(self.start_time)}',
                                    'audioSize': self.audio_sent_bytes,
                                    'message': 'AI回复完成'
                                }
                            })
                            
                            await self.send_control_message({
                                'type': 'state_change',
                                'payload': {
                                    'state': 'idle',
                                    'previousState': 'speaking',
                                    'reason': 'ai_response_completed'
                                }
                            })
                        else:
                            # 没有发送过音频，说明豆包没有回复
                            logger.warning("⚠️ 豆包没有返回音频，可能网络问题或API异常")
                            await self.send_control_message({
                                'type': 'ai_response_end',
                                'payload': {
                                    'responseId': f'resp_{int(self.start_time)}',
                                    'audioSize': 0,
                                    'message': 'AI无回复'
                                }
                            })
                            
                            await self.send_control_message({
                                'type': 'state_change',
                                'payload': {
                                    'state': 'idle',
                                    'previousState': 'processing',
                                    'reason': 'ai_no_response'
                                }
                            })
                        
                        # 标记已结束，避免重复发送
                        self.response_ended = True
                        # 重置音频计数器，为下一轮做准备
                        self.audio_sent_bytes = 0
                        # 不break，继续外层循环等待下一次录音
                        continue
                    elif isinstance(audio_data, tuple):
                        # 处理文本数据
                        if audio_data[0] == 'llm_text':
                            # LLM文本片段
                            self.llm_text_chunks.append(audio_data[1])
                        elif audio_data[0] == 'asr_text':
                            # ASR识别文本（保存最新的）
                            self.asr_text = audio_data[1]
                        continue
                    elif audio_data:
                        # bytes表示音频数据
                        # 发送AI响应开始消息（仅第一次）
                        if self.audio_sent_bytes == 0:
                            await self.send_control_message({
                                'type': 'ai_response_start',
                                'payload': {
                                    'responseId': f'resp_{int(time.time())}',
                                    'message': 'AI开始回复'
                                }
                            })
                            
                            # 发送状态变更为speaking
                            await self.send_control_message({
                                'type': 'state_change',
                                'payload': {
                                    'state': 'speaking',
                                    'previousState': 'processing',
                                    'reason': 'ai_started_response'
                                }
                            })
                        
                        # 将音频数据放入队列
                        await self.to_client_queue.put(audio_data)
                        self.audio_sent_bytes += len(audio_data)
                        performance_monitor.record_audio_sent(len(audio_data))
                        logger.debug(f"收到豆包音频: {len(audio_data)} bytes, 总计: {self.audio_sent_bytes} bytes")
                        
                        # 监控输出是否过长（超过100KB警告）
                        if self.audio_sent_bytes > 100 * 1024:
                            logger.warning(f"⚠️ AI回复过长: {self.audio_sent_bytes / 1024:.1f}KB，可能需要优化prompt")
                    else:
                        # None表示其他事件（文本、状态等），继续等待
                        logger.debug("收到豆包事件，继续等待音频")
                        await asyncio.sleep(0.05)
                    
                except asyncio.TimeoutError:
                    logger.warning(f"接收豆包音频超时（{settings.DOUBAO_RESPONSE_TIMEOUT}秒），发送超时结束消息")
                    
                    # 超时后发送结束消息
                    if self.audio_sent_bytes > 0:
                        logger.info(f"超时但有音频，已发送: {self.audio_sent_bytes} bytes")
                        await self.send_control_message({
                            'type': 'ai_response_end',
                            'payload': {
                                'responseId': f'resp_{int(self.start_time)}',
                                'audioSize': self.audio_sent_bytes,
                                'message': 'AI回复完成（超时）'
                            }
                        })
                    else:
                        logger.warning("超时且无音频，豆包无回复")
                        await self.send_control_message({
                            'type': 'ai_response_end',
                            'payload': {
                                'responseId': f'resp_{int(self.start_time)}',
                                'audioSize': 0,
                                'message': 'AI无回复（超时）'
                            }
                        })
                    
                    await self.send_control_message({
                        'type': 'state_change',
                        'payload': {
                            'state': 'idle',
                            'previousState': 'processing',
                            'reason': 'ai_timeout'
                        }
                    })
                    
                    # 重置音频计数器
                    self.audio_sent_bytes = 0
                    continue
                    
        except Exception as e:
            logger.error(f"接收豆包响应错误: {e}")
    
    async def send_to_client(self):
        """发送音频到客户端"""
        try:
            while self.is_streaming:
                audio_chunk = await self.to_client_queue.get()
                
                # 发送音频数据
                await self.websocket.send_bytes(audio_chunk)
                logger.debug(f"发送音频到客户端: {len(audio_chunk)} bytes")
                
        except Exception as e:
            logger.error(f"发送客户端数据错误: {e}")
    
    async def check_heartbeat(self):
        """检查心跳超时"""
        try:
            while self.is_streaming:
                await asyncio.sleep(10)  # 每10秒检查一次
                
                elapsed = time.time() - self.last_ping_time
                if elapsed > self.heartbeat_timeout:
                    logger.warning(f"会话 {self.session_id} 心跳超时 ({elapsed:.1f}s)")
                    await self.send_error(ErrorCode.SESSION_EXPIRED, "心跳超时，连接将关闭")
                    self.is_streaming = False
                    break
                    
        except Exception as e:
            logger.error(f"心跳检测错误: {e}")
    
    async def handle_control_message(self, message_text: str):
        """处理控制消息"""
        try:
            message = json.loads(message_text)
            msg_type = message.get('type')
            
            logger.info(f"收到控制消息: {msg_type}")
            
            if msg_type == 'init':
                await self.handle_init(message)
                
            elif msg_type == 'start_recording':
                await self.handle_start_recording(message)
                
            elif msg_type == 'stop_recording':
                await self.handle_stop_recording(message)
                
            elif msg_type == 'ping':
                # 更新心跳时间
                self.last_ping_time = time.time()
                await self.send_control_message({
                    'type': 'pong',
                    'timestamp': int(time.time() * 1000)
                })
                
            elif msg_type == 'close':
                self.is_streaming = False
                
            else:
                logger.warning(f"未知消息类型: {msg_type}")
                
        except json.JSONDecodeError as e:
            logger.error(f"JSON解析错误: {e}")
            await self.send_error(ErrorCode.INVALID_MESSAGE)
        except Exception as e:
            logger.error(f"处理控制消息错误: {e}")
            await self.send_error(ErrorCode.INTERNAL_ERROR, f"处理消息失败: {str(e)}")
    
    async def handle_init(self, message: dict):
        """处理初始化消息"""
        await self.send_control_message({
            'type': 'init_ack',
            'id': message.get('id'),
            'payload': {
                'status': 'success',
                'sessionId': self.session_id,
                'config': {
                    'maxAudioChunkSize': 8192,
                    'heartbeatInterval': 30000
                }
            }
        })
    
    async def handle_start_recording(self, message: dict):
        """处理开始录音"""
        self.is_recording = True
        self.start_time = time.time()
        
        # 重置音频计数器（新一轮对话）
        self.audio_received_bytes = 0
        self.audio_sent_bytes = 0
        self.llm_text_chunks = []  # 重置文本收集
        self.asr_text = ''  # 重置ASR文本
        self.response_ended = False  # 重置结束标志
        
        logger.info(f"🎙️ [开始录音] 新一轮对话")
        
        await self.send_control_message({
            'type': 'recording_started',
            'id': message.get('id'),
            'payload': {
                'status': 'ready',
                'message': '准备接收音频数据'
            }
        })
        
        await self.send_control_message({
            'type': 'state_change',
            'payload': {
                'state': 'listening',
                'previousState': 'idle',
                'reason': 'user_started_recording'
            }
        })
    
    async def handle_stop_recording(self, message: dict):
        """处理停止录音"""
        self.is_recording = False
        duration = int((time.time() - self.start_time) * 1000) if self.start_time > 0 else 0
        
        # 记录录音完成信息
        logger.info(f"🎤 [录音完成] 时长: {duration}ms, 音频: {self.audio_received_bytes / 1024:.1f}KB")
        
        # 等待一下让ASR完成识别
        await asyncio.sleep(0.3)
        
        # 打印ASR识别的完整文本
        if self.asr_text:
            logger.info(f"📝 [用户说] {self.asr_text}")
        
        # 等待队列中的音频数据全部发送完成
        await asyncio.sleep(0.15)
        
        # 发送静音帧让豆包VAD检测到结束
        # 豆包配置了 end_smooth_window_ms: 1200，需要1.2秒静音
        # 发送1300ms静音确保触发（16kHz 16bit = 32000 bytes/s）
        if self.doubao_client:
            try:
                silence_duration_ms = 1300
                silence_bytes = int(16000 * 2 * silence_duration_ms / 1000)
                silence_data = b'\x00' * silence_bytes
                
                await self.doubao_client.send_audio(silence_data)
                logger.info(f"发送{silence_duration_ms}ms静音数据，触发VAD结束检测")
            except Exception as e:
                logger.error(f"发送静音数据失败: {e}")
        
        await self.send_control_message({
            'type': 'recording_stopped',
            'id': message.get('id'),
            'payload': {
                'status': 'success',
                'duration': duration,
                'message': '录音已停止，等待AI回复'
            }
        })
        
        # 注意：AI响应开始消息会在receive_from_doubao中自动发送
        # 这里只发送状态变更
        await self.send_control_message({
            'type': 'state_change',
            'payload': {
                'state': 'processing',
                'previousState': 'listening',
                'reason': 'waiting_ai_response'
            }
        })
    
    async def send_control_message(self, message: dict):
        """发送控制消息"""
        try:
            if 'timestamp' not in message:
                message['timestamp'] = int(time.time() * 1000)
            
            await self.websocket.send_text(json.dumps(message))
            
        except Exception as e:
            logger.error(f"发送控制消息失败: {e}")
    
    async def send_error(self, error_code: int, custom_message: str = None):
        """发送错误消息"""
        message = custom_message or get_error_message(error_code)
        
        await self.send_control_message({
            'type': 'error',
            'payload': {
                'code': error_code,
                'message': message,
                'severity': 'error',
                'recoverable': True
            }
        })
    
    async def send_metrics(self):
        """发送性能指标"""
        latency = int((time.time() - self.start_time) * 1000) if self.start_time > 0 else 0
        
        await self.send_control_message({
            'type': 'metrics',
            'payload': {
                'latency': {
                    'total': latency,
                    'network': 0,  # TODO: 实际计算
                    'processing': 0,  # TODO: 实际计算
                    'ai': 0  # TODO: 实际计算
                },
                'audio': {
                    'received': self.audio_received_bytes,
                    'sent': self.audio_sent_bytes
                }
            }
        })
    
    async def cleanup(self):
        """清理会话"""
        try:
            # 记录会话
            performance_monitor.record_session()
            
            # 记录延迟
            if self.start_time > 0:
                latency = (time.time() - self.start_time) * 1000
                performance_monitor.record_latency(latency)
            
            # 注意：ai_response_end已在receive_from_doubao中发送，这里不重复发送
            
            # 归还连接到池中
            if self.doubao_client:
                await connection_pool.release(self.doubao_client)
                self.doubao_client = None
            
            logger.info(f"会话 {self.session_id} 已清理 - 收到: {self.audio_received_bytes}B, 发送: {self.audio_sent_bytes}B")
            
        except Exception as e:
            logger.error(f"清理会话失败: {e}")


class ConnectionManager:
    """连接管理器"""
    
    def __init__(self):
        self.active_connections: Dict[str, StreamingSession] = {}
    
    async def connect(self, websocket: WebSocket, session_id: str) -> StreamingSession:
        """建立连接"""
        await websocket.accept()
        
        session = StreamingSession(websocket, session_id)
        self.active_connections[session_id] = session
        
        logger.info(f"新连接建立: {session_id}, 当前活跃连接数: {len(self.active_connections)}")
        
        return session
    
    async def disconnect(self, session_id: str):
        """断开连接"""
        if session_id in self.active_connections:
            session = self.active_connections[session_id]
            await session.cleanup()
            del self.active_connections[session_id]
            
            logger.info(f"连接已断开: {session_id}, 当前活跃连接数: {len(self.active_connections)}")
    
    def get_session(self, session_id: str) -> Optional[StreamingSession]:
        """获取会话"""
        return self.active_connections.get(session_id)


# 创建全局连接管理器
connection_manager = ConnectionManager()