# AIfeng/2025-07-11 13:36:00
"""
豆包ASR客户端核心模块
提供完整的语音识别服务接口，支持流式和非流式识别
"""

import asyncio
import json
import logging
import time
import uuid
from pathlib import Path
from typing import Dict, Any, Optional, Callable, AsyncGenerator

import aiofiles
import websockets
from websockets.exceptions import ConnectionClosedError, WebSocketException

from .protocol import DoubaoProtocol, MessageType
from .audio_utils import AudioProcessor


class DoubaoASRClient:
    """豆包ASR客户端"""
    
    def __init__(self, config: Dict[str, Any]):
        """
        初始化ASR客户端
        
        Args:
            config: 配置字典
        """
        self.config = config
        self.asr_config = config.get('asr_config', {})
        self.auth_config = config.get('auth_config', {})
        self.audio_config = config.get('audio_config', {})
        self.connection_config = config.get('connection_config', {})
        self.logging_config = config.get('logging_config', {})
        
        # 设置日志
        self.logger = self._setup_logger()
        
        # 协议处理器
        self.protocol = DoubaoProtocol()
        
        # 音频处理器
        self.audio_processor = AudioProcessor()
        
        # 连接状态
        self.is_connected = False
        self.current_session_id = None
    
    def _setup_logger(self) -> logging.Logger:
        """设置日志记录器"""
        logger = logging.getLogger('doubao_asr')
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            )
            handler.setFormatter(formatter)
            logger.addHandler(handler)
        
        if self.logging_config.get('enable_debug', False):
            logger.setLevel(logging.DEBUG)
        else:
            logger.setLevel(logging.INFO)
        
        return logger
    
    def _get_ws_url(self, streaming: bool = True) -> str:
        """获取WebSocket URL"""
        if streaming:
            return self.asr_config.get('ws_url', 'wss://openspeech.bytedance.com/api/v3/sauc/bigmodel')
        else:
            return self.asr_config.get('ws_url_nostream', 'wss://openspeech.bytedance.com/api/v3/sauc/bigmodel_nostream')
    
    def _build_auth_headers(self, request_id: str) -> Dict[str, str]:
        """构建认证头部"""
        headers = {
            'X-Api-Resource-Id': self.asr_config.get('resource_id', 'volc.bigasr.sauc.duration'),
            'X-Api-Access-Key': self.auth_config.get('access_key', ''),
            'X-Api-App-Key': self.auth_config.get('app_key', ''),
            'X-Api-Request-Id': request_id
        }
        return headers
    
    def _build_request_params(
        self,
        request_id: str,
        audio_format: str = 'wav',
        sample_rate: int = 16000,
        bits: int = 16,
        channels: int = 1,
        uid: str = 'default_user'
    ) -> Dict[str, Any]:
        """构建请求参数"""
        return {
            'user': {
                'uid': uid
            },
            'audio': {
                'format': audio_format,
                'sample_rate': sample_rate,
                'bits': bits,
                'channel': channels,
                'codec': self.audio_config.get('default_codec', 'raw')
            },
            'request': {
                'model_name': self.asr_config.get('model_name', 'bigmodel'),
                'enable_punc': self.asr_config.get('enable_punc', True)
            }
        }
    
    async def recognize_file(
        self,
        audio_path: str,
        streaming: bool = True,
        result_callback: Optional[Callable[[Dict[str, Any]], None]] = None,
        **kwargs
    ) -> Dict[str, Any]:
        """
        识别音频文件
        
        Args:
            audio_path: 音频文件路径
            streaming: 是否使用流式识别
            result_callback: 结果回调函数
            **kwargs: 其他参数
            
        Returns:
            Dict: 识别结果
        """
        try:
            # 读取音频文件
            async with aiofiles.open(audio_path, mode='rb') as f:
                audio_data = await f.read()
            
            self.logger.info(f"开始识别音频文件: {audio_path}, 大小: {len(audio_data)} 字节")
            
            # 识别音频数据
            return await self.recognize_audio_data(
                audio_data,
                streaming=streaming,
                result_callback=result_callback,
                **kwargs
            )
        
        except Exception as e:
            self.logger.error(f"识别音频文件失败: {e}")
            return {
                'success': False,
                'error': str(e),
                'audio_path': audio_path
            }
    
    async def recognize_audio_data(
        self,
        audio_data: bytes,
        streaming: bool = True,
        result_callback: Optional[Callable[[Dict[str, Any]], None]] = None,
        **kwargs
    ) -> Dict[str, Any]:
        """
        识别音频数据
        
        Args:
            audio_data: 音频数据
            streaming: 是否使用流式识别
            result_callback: 结果回调函数
            **kwargs: 其他参数
            
        Returns:
            Dict: 识别结果
        """
        request_id = str(uuid.uuid4())
        self.current_session_id = request_id
        
        try:
            # 准备音频数据
            audio_format, segment_size, metadata = self.audio_processor.prepare_audio_for_recognition(
                audio_data,
                segment_duration_ms=self.asr_config.get('seg_duration', 200)
            )
            
            self.logger.info(f"音频格式: {audio_format}, 分片大小: {segment_size}, 元数据: {metadata}")
            
            # 构建请求参数
            request_params = self._build_request_params(
                request_id,
                audio_format=audio_format,
                sample_rate=metadata.get('sample_rate', 16000),
                bits=metadata.get('sample_width', 2) * 8,
                channels=metadata.get('channels', 1),
                uid=kwargs.get('uid', 'default_user')
            )
            
            # 执行识别
            if streaming:
                return await self._streaming_recognize(
                    audio_data,
                    request_params,
                    segment_size,
                    request_id,
                    result_callback
                )
            else:
                return await self._non_streaming_recognize(
                    audio_data,
                    request_params,
                    request_id
                )
        
        except Exception as e:
            self.logger.error(f"识别音频数据失败: {e}")
            return {
                'success': False,
                'error': str(e),
                'request_id': request_id
            }
    
    async def _streaming_recognize(
        self,
        audio_data: bytes,
        request_params: Dict[str, Any],
        segment_size: int,
        request_id: str,
        result_callback: Optional[Callable[[Dict[str, Any]], None]] = None
    ) -> Dict[str, Any]:
        """流式识别处理"""
        ws_url = self._get_ws_url(streaming=True)
        headers = self._build_auth_headers(request_id)
        
        results = []
        final_result = None
        
        try:
            # 兼容不同版本的websockets库
            connect_kwargs = {
                'uri': ws_url,
                'max_size': self.connection_config.get('max_size', 1000000000)
            }
            
            # 尝试使用新版本的additional_headers参数
            try:
                async with websockets.connect(
                    **connect_kwargs,
                    additional_headers=headers
                ) as ws:
                    await self._handle_streaming_connection(ws, audio_data, request_params, segment_size, request_id, result_callback, results, final_result)
            except TypeError:
                # 回退到旧版本的extra_headers参数
                async with websockets.connect(
                    **connect_kwargs,
                    extra_headers=headers
                ) as ws:
                    await self._handle_streaming_connection(ws, audio_data, request_params, segment_size, request_id, result_callback, results, final_result)
            
            return {
                'success': True,
                'request_id': request_id,
                'results': results,
                'final_result': final_result,
                'total_results': len(results)
            }
        
        except ConnectionClosedError as e:
            self.logger.error(f"WebSocket连接关闭: {e.code} - {e.reason}")
            return {
                'success': False,
                'error': f"连接关闭: {e.reason}",
                'error_code': e.code,
                'request_id': request_id
            }
        
        except WebSocketException as e:
            self.logger.error(f"WebSocket异常: {e}")
            return {
                'success': False,
                'error': str(e),
                'request_id': request_id
            }
        
        except Exception as e:
            self.logger.error(f"流式识别异常: {e}")
            return {
                'success': False,
                'error': str(e),
                'request_id': request_id
            }
        
        finally:
            self.is_connected = False
    
    async def _handle_streaming_connection(
        self,
        ws,
        audio_data: bytes,
        request_params: Dict[str, Any],
        segment_size: int,
        request_id: str,
        result_callback: Optional[Callable[[Dict[str, Any]], None]],
        results: list,
        final_result: Any
    ):
        """处理流式连接的核心逻辑"""
        self.is_connected = True
        self.logger.info(f"WebSocket连接建立成功")
        
        # 发送初始请求
        seq = 1
        full_request = self.protocol.build_full_request(request_params, seq)
        await ws.send(full_request)
        
        # 接收初始响应
        response = await ws.recv()
        result = self.protocol.parse_response(response)
        
        if self.logging_config.get('log_responses', True):
            self.logger.debug(f"初始响应: {result}")
        
        # 分片发送音频数据
        for chunk, is_last in self.audio_processor.slice_audio_data(audio_data, segment_size):
            seq += 1
            if is_last:
                seq = -seq
            
            start_time = time.time()
            
            # 构建音频请求
            audio_request = self.protocol.build_audio_request(
                chunk, seq, is_last
            )
            
            # 发送音频数据
            await ws.send(audio_request)
            
            # 接收响应
            response = await ws.recv()
            result = self.protocol.parse_response(response)
            
            # 处理结果
            if result.get('payload_msg'):
                results.append(result)
                
                # 调用回调函数
                if result_callback:
                    try:
                        result_callback(result)
                    except Exception as e:
                        self.logger.warning(f"回调函数执行失败: {e}")
            
            if result.get('is_last_package'):
                final_result = result
                break
            
            # 流式识别延时控制
            if self.asr_config.get('streaming_mode', True):
                elapsed = time.time() - start_time
                sleep_time = max(0, (self.asr_config.get('seg_duration', 200) / 1000.0) - elapsed)
                if sleep_time > 0:
                    await asyncio.sleep(sleep_time)
    
    async def _non_streaming_recognize(
        self,
        audio_data: bytes,
        request_params: Dict[str, Any],
        request_id: str
    ) -> Dict[str, Any]:
        """非流式识别处理"""
        ws_url = self._get_ws_url(streaming=False)
        headers = self._build_auth_headers(request_id)
        
        try:
            # 兼容不同版本的websockets库
            connect_kwargs = {
                'uri': ws_url,
                'max_size': self.connection_config.get('max_size', 1000000000)
            }
            
            # 尝试使用新版本的additional_headers参数
            try:
                async with websockets.connect(
                    **connect_kwargs,
                    additional_headers=headers
                ) as ws:
                    return await self._handle_non_streaming_connection(ws, audio_data, request_params, request_id)
            except TypeError:
                # 回退到旧版本的extra_headers参数
                async with websockets.connect(
                    **connect_kwargs,
                    extra_headers=headers
                ) as ws:
                    return await self._handle_non_streaming_connection(ws, audio_data, request_params, request_id)
        
        except Exception as e:
            self.logger.error(f"非流式识别异常: {e}")
            return {
                'success': False,
                'error': str(e),
                'request_id': request_id
            }
        
        finally:
            self.is_connected = False
    
    async def _handle_non_streaming_connection(
        self,
        ws,
        audio_data: bytes,
        request_params: Dict[str, Any],
        request_id: str
    ) -> Dict[str, Any]:
        """处理非流式连接的核心逻辑"""
        self.is_connected = True
        self.logger.info(f"WebSocket连接建立成功")
        
        # 发送完整请求（包含音频数据）
        full_request = self.protocol.build_full_request(request_params, 1)
        await ws.send(full_request)
        
        # 发送音频数据
        audio_request = self.protocol.build_audio_request(
            audio_data, -1, is_last=True
        )
        await ws.send(audio_request)
        
        # 接收最终结果
        response = await ws.recv()
        result = self.protocol.parse_response(response)
        
        self.is_connected = False
        
        return {
            'success': True,
            'request_id': request_id,
            'result': result
        }
    
    async def close(self):
        """关闭客户端"""
        self.is_connected = False
        self.current_session_id = None
        self.logger.info("ASR客户端已关闭")
    
    def get_status(self) -> Dict[str, Any]:
        """获取客户端状态"""
        return {
            'is_connected': self.is_connected,
            'current_session_id': self.current_session_id,
            'config': {
                'ws_url': self._get_ws_url(),
                'model_name': self.asr_config.get('model_name'),
                'streaming_mode': self.asr_config.get('streaming_mode')
            }
        }