# AIfeng/2025-07-08 15:59:04
# 服务端录音WebSocket接口 - 集成StreamingRecorder实现服务端收音
# 核心功能：WebSocket通信、服务端录音控制、识别结果转发

import asyncio
import json
import weakref
from aiohttp import web, WSMsgType
from typing import Dict, Optional, Set
import sys
import os
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from streaming.streaming_recorder import StreamingRecorder
from logger import get_logger

logger = get_logger("ServerRecordingWebSocket")

class ServerRecordingWebSocket:
    """服务端录音WebSocket管理器
    
    功能：
    1. 管理WebSocket连接
    2. 控制服务端录音器
    3. 转发识别结果到前端
    4. 处理录音状态同步
    """
    
    def __init__(self, config_path: str = "streaming/streaming_config.json"):
        self.config = self._load_config(config_path)
        self.websocket_clients: Dict[str, Set[web.WebSocketResponse]] = {}
        self.recording_sessions: Dict[str, StreamingRecorder] = {}
        self.session_status: Dict[str, Dict] = {}
        
        logger.info("服务端录音WebSocket管理器初始化完成")
    
    def _load_config(self, config_path: str) -> dict:
        """加载配置文件"""
        try:
            import json
            with open(config_path, 'r', encoding='utf-8') as f:
                config = json.load(f)
            logger.info(f"配置文件加载成功: {config_path}")
            return config
        except Exception as e:
            logger.warning(f"配置文件加载失败: {e}，使用默认配置")
            return self._get_default_config()
    
    def _get_default_config(self) -> dict:
        """获取默认配置"""
        return {
            "streaming_recorder": {
                "audio": {
                    "chunk_size": 1024,
                    "sample_rate": 16000,
                    "channels": 1,
                    "format": "paInt16"
                },
                "audio_gain": 3.0,
                "enable_audio_gain": True
            },
            "streaming_vad": {
                "threshold": 0.03,
                "max_silence_duration": 1.5,
                "min_speech_duration": 0.5,
                "max_speech_duration": 30.0,
                "pre_buffer_duration": 0.5,
                "dynamic_threshold_factor": 0.8,
                "partial_result_interval": 2.0
            },
            "streaming_recognition": {
                "confidence_threshold": 0.6,
                "max_session_duration": 60.0,
                "result_merge_window": 1.0
            }
        }
    
    def register_routes(self, app: web.Application):
        """注册WebSocket路由"""
        app.router.add_get('/ws/server-recording', self.websocket_handler)
        logger.info("服务端录音WebSocket路由已注册: /ws/server-recording")
    
    async def websocket_handler(self, request):
        """WebSocket连接处理器"""
        ws = web.WebSocketResponse()
        await ws.prepare(request)
        
        session_id = None
        logger.info('新的服务端录音WebSocket连接建立')
        
        try:
            async for msg in ws:
                if msg.type == WSMsgType.TEXT:
                    try:
                        data = json.loads(msg.data)
                        await self._handle_websocket_message(ws, data)
                        
                    except json.JSONDecodeError:
                        await ws.send_str(json.dumps({
                            'type': 'error',
                            'message': '无效的JSON格式'
                        }))
                    except Exception as e:
                        logger.error(f'处理WebSocket消息错误: {e}')
                        await ws.send_str(json.dumps({
                            'type': 'error',
                            'message': f'消息处理错误: {str(e)}'
                        }))
                        
                elif msg.type == WSMsgType.ERROR:
                    logger.error(f'WebSocket错误: {ws.exception()}')
                    break
                elif msg.type == WSMsgType.CLOSE:
                    logger.info('服务端录音WebSocket连接正常关闭')
                    break
                    
        except ConnectionResetError:
            logger.warning('服务端录音WebSocket连接被远程主机重置')
        except ConnectionAbortedError:
            logger.warning('服务端录音WebSocket连接被中止')
        except Exception as e:
            logger.error(f'WebSocket连接错误: {e}')
        finally:
            if session_id:
                await self._cleanup_session(session_id, ws)
            logger.info('服务端录音WebSocket连接关闭')
        
        return ws
    
    async def _handle_websocket_message(self, ws: web.WebSocketResponse, data: Dict):
        """处理WebSocket消息"""
        message_type = data.get('type')
        session_id = data.get('session_id')
        
        if message_type == 'connect':
            await self._handle_connect(ws, session_id)
        elif message_type == 'start_recording':
            await self._handle_start_recording(ws, session_id, data.get('config', {}))
        elif message_type == 'stop_recording':
            await self._handle_stop_recording(ws, session_id)
        elif message_type == 'get_status':
            await self._handle_get_status(ws, session_id)
        elif message_type == 'ping':
            await ws.send_str(json.dumps({'type': 'pong'}))
        else:
            await ws.send_str(json.dumps({
                'type': 'error',
                'message': f'未知消息类型: {message_type}'
            }))
    
    async def _handle_connect(self, ws: web.WebSocketResponse, session_id: str):
        """处理连接请求"""
        if not session_id:
            session_id = f"session_{len(self.websocket_clients)}"
        
        # 初始化会话连接集合
        if session_id not in self.websocket_clients:
            self.websocket_clients[session_id] = weakref.WeakSet()
        
        # 添加WebSocket连接
        self.websocket_clients[session_id].add(ws)
        
        # 初始化会话状态
        self.session_status[session_id] = {
            'connected': True,
            'recording': False,
            'created_time': asyncio.get_event_loop().time()
        }
        
        logger.info(f'会话 {session_id} 连接成功，当前连接数: {len(self.websocket_clients[session_id])}')
        
        await ws.send_str(json.dumps({
            'type': 'connected',
            'session_id': session_id,
            'message': '服务端录音连接成功'
        }))
    
    async def _handle_start_recording(self, ws: web.WebSocketResponse, session_id: str, config: Dict):
        """处理开始录音请求"""
        try:
            if not session_id or session_id not in self.websocket_clients:
                await ws.send_str(json.dumps({
                    'type': 'error',
                    'message': '会话未连接，请先连接'
                }))
                return
            
            # 检查是否已在录音
            if session_id in self.recording_sessions:
                recorder = self.recording_sessions[session_id]
                if recorder.is_recording():
                    await ws.send_str(json.dumps({
                        'type': 'warning',
                        'message': '录音已在进行中'
                    }))
                    return
            
            # 合并配置
            merged_config = self.config.copy()
            if config:
                # 递归合并配置
                self._merge_config(merged_config, config)
            
            # 创建录音器
            recorder = self._create_recorder(session_id, merged_config)
            self.recording_sessions[session_id] = recorder
            
            # 开始录音
            device_index = config.get('device_index')
            success = recorder.start_recording(device_index)
            
            if success:
                self.session_status[session_id]['recording'] = True
                logger.info(f'会话 {session_id} 开始录音')
                
                await self._broadcast_to_session(session_id, {
                    'type': 'recording_started',
                    'session_id': session_id,
                    'message': '服务端录音已开始'
                })
            else:
                await ws.send_str(json.dumps({
                    'type': 'error',
                    'message': '启动录音失败'
                }))
                
        except Exception as e:
            logger.error(f'启动录音错误: {e}')
            await ws.send_str(json.dumps({
                'type': 'error',
                'message': f'启动录音失败: {str(e)}'
            }))
    
    async def _handle_stop_recording(self, ws: web.WebSocketResponse, session_id: str):
        """处理停止录音请求"""
        try:
            if session_id not in self.recording_sessions:
                await ws.send_str(json.dumps({
                    'type': 'warning',
                    'message': '没有正在进行的录音'
                }))
                return
            
            recorder = self.recording_sessions[session_id]
            recorder.stop_recording()
            
            self.session_status[session_id]['recording'] = False
            logger.info(f'会话 {session_id} 停止录音')
            
            await self._broadcast_to_session(session_id, {
                'type': 'recording_stopped',
                'session_id': session_id,
                'message': '服务端录音已停止'
            })
            
        except Exception as e:
            logger.error(f'停止录音错误: {e}')
            await ws.send_str(json.dumps({
                'type': 'error',
                'message': f'停止录音失败: {str(e)}'
            }))
    
    async def _handle_get_status(self, ws: web.WebSocketResponse, session_id: str):
        """处理状态查询请求"""
        status = {
            'session_id': session_id,
            'connected': session_id in self.websocket_clients,
            'recording': False
        }
        
        if session_id in self.recording_sessions:
            recorder = self.recording_sessions[session_id]
            status['recording'] = recorder.is_recording()
        
        if session_id in self.session_status:
            status.update(self.session_status[session_id])
        
        await ws.send_str(json.dumps({
            'type': 'status',
            'data': status
        }))
    
    def _create_recorder(self, session_id: str, config: Dict) -> StreamingRecorder:
        """创建录音器实例"""
        audio_config = config.get('streaming_recorder', {}).get('audio', {})
        vad_config = config.get('streaming_vad', {})
        recognition_config = config.get('streaming_recognition', {})
        
        recorder = StreamingRecorder(
            chunk=audio_config.get('chunk_size', 1024),
            rate=audio_config.get('sample_rate', 16000),
            channels=audio_config.get('channels', 1),
            volume_threshold=vad_config.get('threshold', 0.03),
            silence_duration=vad_config.get('max_silence_duration', 1.5),
            min_speech_duration=vad_config.get('min_speech_duration', 0.5),
            max_speech_duration=vad_config.get('max_speech_duration', 30.0),
            pre_buffer_duration=vad_config.get('pre_buffer_duration', 0.5),
            dynamic_threshold_factor=vad_config.get('dynamic_threshold_factor', 0.8),
            partial_result_interval=vad_config.get('partial_result_interval', 2.0),
            confidence_threshold=recognition_config.get('confidence_threshold', 0.6),
            max_session_duration=recognition_config.get('max_session_duration', 60.0),
            result_merge_window=recognition_config.get('result_merge_window', 1.0),
            username=f"server_recording_{session_id}",
            config=config
        )
        
        # 设置回调函数
        recorder.on_partial_result = lambda sid, text, confidence: asyncio.create_task(
            self._on_partial_result(session_id, sid, text, confidence)
        )
        recorder.on_final_result = lambda sid, text, confidence: asyncio.create_task(
            self._on_final_result(session_id, sid, text, confidence)
        )
        recorder.on_session_complete = lambda sid, results: asyncio.create_task(
            self._on_session_complete(session_id, sid, results)
        )
        recorder.on_status_update = lambda status: asyncio.create_task(
            self._on_status_update(session_id, status)
        )
        
        logger.info(f'为会话 {session_id} 创建录音器')
        return recorder
    
    async def _on_partial_result(self, session_id: str, recording_session_id: str, text: str, confidence: float):
        """处理部分识别结果"""
        await self._broadcast_to_session(session_id, {
            'type': 'asr_partial_result',
            'session_id': session_id,
            'recording_session_id': recording_session_id,
            'data': {
                'text': text,
                'confidence': confidence,
                'is_final': False
            }
        })
    
    async def _on_final_result(self, session_id: str, recording_session_id: str, text: str, confidence: float):
        """处理最终识别结果"""
        await self._broadcast_to_session(session_id, {
            'type': 'asr_final_result',
            'session_id': session_id,
            'recording_session_id': recording_session_id,
            'data': {
                'text': text,
                'confidence': confidence,
                'is_final': True
            }
        })
    
    async def _on_session_complete(self, session_id: str, recording_session_id: str, results: Dict):
        """处理会话完成"""
        await self._broadcast_to_session(session_id, {
            'type': 'asr_session_complete',
            'session_id': session_id,
            'recording_session_id': recording_session_id,
            'data': results
        })
    
    async def _on_status_update(self, session_id: str, status: Dict):
        """处理状态更新"""
        await self._broadcast_to_session(session_id, {
            'type': 'status_update',
            'session_id': session_id,
            'data': status
        })
    
    async def _broadcast_to_session(self, session_id: str, message: Dict):
        """向会话的所有WebSocket连接广播消息"""
        if session_id not in self.websocket_clients:
            return
        
        message_str = json.dumps(message)
        failed_connections = []
        
        for ws in list(self.websocket_clients[session_id]):
            try:
                await ws.send_str(message_str)
            except Exception as e:
                logger.error(f'向会话 {session_id} 发送消息失败: {e}')
                failed_connections.append(ws)
        
        # 清理失败的连接
        for ws in failed_connections:
            try:
                self.websocket_clients[session_id].discard(ws)
            except:
                pass
    
    async def _cleanup_session(self, session_id: str, ws: web.WebSocketResponse):
        """清理会话资源"""
        try:
            # 移除WebSocket连接
            if session_id in self.websocket_clients:
                self.websocket_clients[session_id].discard(ws)
                
                # 如果没有更多连接，清理会话
                if len(self.websocket_clients[session_id]) == 0:
                    # 停止录音
                    if session_id in self.recording_sessions:
                        recorder = self.recording_sessions[session_id]
                        if recorder.is_recording():
                            recorder.stop_recording()
                        del self.recording_sessions[session_id]
                    
                    # 清理状态
                    if session_id in self.session_status:
                        del self.session_status[session_id]
                    
                    # 清理连接集合
                    del self.websocket_clients[session_id]
                    
                    logger.info(f'会话 {session_id} 资源已清理')
                    
        except Exception as e:
            logger.error(f'清理会话 {session_id} 资源时出错: {e}')
    
    def _merge_config(self, base_config: Dict, override_config: Dict):
        """递归合并配置"""
        for key, value in override_config.items():
            if key in base_config and isinstance(base_config[key], dict) and isinstance(value, dict):
                self._merge_config(base_config[key], value)
            else:
                base_config[key] = value

# 全局实例
server_recording_ws = ServerRecordingWebSocket()

def register_server_recording_routes(app: web.Application):
    """注册服务端录音WebSocket路由"""
    server_recording_ws.register_routes(app)
    logger.info("服务端录音WebSocket路由注册完成")