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.perf_recording_start = 0  # 开始录音时间
        self.perf_recording_stop = 0   # 停止录音时间
        self.perf_asr_complete = 0     # ASR识别完成时间
        self.perf_llm_first_token = 0  # LLM首个token时间
        self.perf_tts_first_audio = 0  # TTS首个音频包时间
        self.perf_audio_complete = 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}")
            # 如果获取连接后出错，需要立即释放
            if self.doubao_client:
                await connection_pool.release(self.doubao_client)
                self.doubao_client = None
            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:
                    # 检查豆包连接是否健康
                    if not self.doubao_client.is_connected():
                        logger.warning(f"⚠️ 豆包连接已失效 (has_fatal_error={self.doubao_client.has_fatal_error})，停止接收循环")
                        
                        # 🔑 退出循环，等待下一轮录音时自动重连
                        # 发送idle状态，让前端准备好下次对话
                        await self.send_control_message({
                            'type': 'state_change',
                            'payload': {
                                'state': 'idle',
                                'previousState': 'speaking',
                                'reason': 'connection_needs_recovery'
                            }
                        })
                        
                        # 退出循环，避免无限循环占用资源
                        break
                    
                    # 接收豆包音频数据，设置响应超时
                    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:
                            # 记录音频完成时间
                            self.perf_audio_complete = time.time()
                            
                            # 计算本轮对话的延迟
                            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}")
                            
                            # 打印性能分析报告
                            self._log_performance_report()
                            
                            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:
                            # 记录TTS首包时间
                            self.perf_tts_first_audio = time.time()
                            
                            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):
        """处理开始录音"""
        # 🔑 关键：检查豆包连接是否健康，如果失效则重新获取
        if self.doubao_client and (not self.doubao_client.is_connected() or self.doubao_client.has_fatal_error):
            logger.warning(f"🔄 豆包连接已失效 (has_fatal_error={self.doubao_client.has_fatal_error})，重新获取连接...")
            
            # 释放旧连接
            try:
                await self.connection_pool.release(self.doubao_client)
            except Exception as e:
                logger.error(f"释放失效连接失败: {e}")
            
            # 重新获取新连接
            try:
                self.doubao_client = await self.connection_pool.acquire()
                logger.info(f"✅ 成功获取新的豆包连接: {self.doubao_client.session_id}")
            except Exception as e:
                logger.error(f"❌ 重新获取豆包连接失败: {e}")
                await self.send_error(
                    ErrorCode.INTERNAL_ERROR,
                    "语音服务连接失败，请稍后重试"
                )
                return
        
        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  # 重置结束标志
        
        # 记录性能时间点
        self.perf_recording_start = time.time()
        self.perf_recording_stop = 0
        self.perf_asr_complete = 0
        self.perf_llm_first_token = 0
        self.perf_tts_first_audio = 0
        self.perf_audio_complete = 0
        
        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
        self.perf_recording_stop = time.time()
        
        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完成识别和队列处理（优化：减少到100ms）
        await asyncio.sleep(0.1)
        
        # 打印ASR识别的完整文本
        if self.asr_text:
            logger.info(f"📝 [用户说] {self.asr_text}")
        
        # 发送静音帧让豆包VAD检测到结束
        # 豆包配置了 end_smooth_window_ms: 800，需要0.8秒静音
        # 发送900ms静音确保触发（16kHz 16bit = 32000 bytes/s）
        if self.doubao_client:
            try:
                silence_duration_ms = 900  # 优化：减少400ms等待时间
                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'
            }
        })
    
    def _log_performance_report(self):
        """打印性能分析报告"""
        if self.perf_recording_stop == 0:
            return
        
        # 计算各阶段耗时（毫秒）
        recording_duration = int((self.perf_recording_stop - self.perf_recording_start) * 1000)
        
        # ASR延迟（从停止录音到ASR完成）- 由于ASR是流式的，这里用TTS首包作为近似
        asr_latency = 0
        if self.perf_tts_first_audio > 0:
            asr_latency = int((self.perf_tts_first_audio - self.perf_recording_stop) * 1000) - 900  # 减去静音时间
        
        # TTS首包延迟
        tts_first_packet = 0
        if self.perf_tts_first_audio > 0:
            tts_first_packet = int((self.perf_tts_first_audio - self.perf_recording_stop) * 1000)
        
        # 音频传输时长
        audio_streaming = 0
        if self.perf_audio_complete > 0 and self.perf_tts_first_audio > 0:
            audio_streaming = int((self.perf_audio_complete - self.perf_tts_first_audio) * 1000)
        
        # 总延迟
        total_latency = 0
        if self.perf_audio_complete > 0:
            total_latency = int((self.perf_audio_complete - self.perf_recording_stop) * 1000)
        
        logger.info("=" * 60)
        logger.info("⚡ 性能分析报告")
        logger.info(f"📊 录音时长: {recording_duration}ms")
        logger.info(f"📊 ASR处理: ~{asr_latency}ms (估算)")
        logger.info(f"📊 TTS首包: {tts_first_packet}ms")
        logger.info(f"📊 音频流式: {audio_streaming}ms")
        logger.info(f"📊 总延迟: {total_latency}ms (停止录音→回复完成)")
        logger.info("=" * 60)
    
    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:
            logger.info(f"开始清理会话 {self.session_id}...")
            
            # 记录会话
            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:
                logger.debug(f"释放会话 {self.session_id} 的豆包连接到连接池")
                await connection_pool.release(self.doubao_client)
                self.doubao_client = None
            else:
                logger.warning(f"会话 {self.session_id} 没有豆包连接需要释放")
            
            logger.info(f"会话 {self.session_id} 已清理 - 收到: {self.audio_received_bytes}B, 发送: {self.audio_sent_bytes}B")
            
        except Exception as e:
            logger.error(f"清理会话失败: {e}")
            # 即使清理失败，也要尝试释放连接
            if self.doubao_client:
                try:
                    await connection_pool.release(self.doubao_client)
                    self.doubao_client = None
                    logger.info(f"在异常处理中成功释放连接")
                except Exception as release_error:
                    logger.error(f"释放连接失败: {release_error}")


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()