#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
语音识别模块
基于 SAUC API 为录音文件提供语音识别功能
"""

import os
import json
import time
import logging
import asyncio
import aiohttp
import struct
import gzip
import uuid
from typing import Optional, List, Dict, Any, AsyncGenerator
from dataclasses import dataclass

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 常量定义
DEFAULT_SAMPLE_RATE = 16000

class ProtocolVersion:
    V1 = 0b0001

class MessageType:
    CLIENT_FULL_REQUEST = 0b0001
    CLIENT_AUDIO_ONLY_REQUEST = 0b0010
    SERVER_FULL_RESPONSE = 0b1001
    SERVER_ERROR_RESPONSE = 0b1111

class MessageTypeSpecificFlags:
    NO_SEQUENCE = 0b0000
    POS_SEQUENCE = 0b0001
    NEG_SEQUENCE = 0b0010
    NEG_WITH_SEQUENCE = 0b0011

class SerializationType:
    NO_SERIALIZATION = 0b0000
    JSON = 0b0001

class CompressionType:
    GZIP = 0b0001

@dataclass
class RecognitionResult:
    """语音识别结果"""
    text: str
    confidence: float
    is_final: bool
    start_time: Optional[float] = None
    end_time: Optional[float] = None

class AudioUtils:
    """音频处理工具类"""
    
    @staticmethod
    def gzip_compress(data: bytes) -> bytes:
        """GZIP压缩"""
        return gzip.compress(data)

    @staticmethod
    def gzip_decompress(data: bytes) -> bytes:
        """GZIP解压缩"""
        return gzip.decompress(data)

    @staticmethod
    def is_wav_file(data: bytes) -> bool:
        """检查是否为WAV文件"""
        if len(data) < 44:
            return False
        return data[:4] == b'RIFF' and data[8:12] == b'WAVE'

    @staticmethod
    def read_wav_info(data: bytes) -> tuple:
        """读取WAV文件信息"""
        if len(data) < 44:
            raise ValueError("Invalid WAV file: too short")
            
        # 解析WAV头
        chunk_id = data[:4]
        if chunk_id != b'RIFF':
            raise ValueError("Invalid WAV file: not RIFF format")
            
        format_ = data[8:12]
        if format_ != b'WAVE':
            raise ValueError("Invalid WAV file: not WAVE format")
            
        # 解析fmt子块
        audio_format = struct.unpack('<H', data[20:22])[0]
        num_channels = struct.unpack('<H', data[22:24])[0]
        sample_rate = struct.unpack('<I', data[24:28])[0]
        bits_per_sample = struct.unpack('<H', data[34:36])[0]
        
        # 查找data子块
        pos = 36
        while pos < len(data) - 8:
            subchunk_id = data[pos:pos+4]
            subchunk_size = struct.unpack('<I', data[pos+4:pos+8])[0]
            if subchunk_id == b'data':
                wave_data = data[pos+8:pos+8+subchunk_size]
                return (
                    num_channels,
                    bits_per_sample // 8,
                    sample_rate,
                    subchunk_size // (num_channels * (bits_per_sample // 8)),
                    wave_data
                )
            pos += 8 + subchunk_size
            
        raise ValueError("Invalid WAV file: no data subchunk found")

class AsrConfig:
    """ASR配置"""
    
    def __init__(self, app_key: str = None, access_key: str = None):
        self.auth = {
            "app_key": app_key or os.getenv("SAUC_APP_KEY", "your_app_key"),
            "access_key": access_key or os.getenv("SAUC_ACCESS_KEY", "your_access_key")
        }
    
    @property
    def app_key(self) -> str:
        return self.auth["app_key"]
    
    @property
    def access_key(self) -> str:
        return self.auth["access_key"]

class AsrRequestHeader:
    """ASR请求头"""
    
    def __init__(self):
        self.message_type = MessageType.CLIENT_FULL_REQUEST
        self.message_type_specific_flags = MessageTypeSpecificFlags.POS_SEQUENCE
        self.serialization_type = SerializationType.JSON
        self.compression_type = CompressionType.GZIP
        self.reserved_data = bytes([0x00])

    def with_message_type(self, message_type: int) -> 'AsrRequestHeader':
        self.message_type = message_type
        return self

    def with_message_type_specific_flags(self, flags: int) -> 'AsrRequestHeader':
        self.message_type_specific_flags = flags
        return self

    def with_serialization_type(self, serialization_type: int) -> 'AsrRequestHeader':
        self.serialization_type = serialization_type
        return self

    def with_compression_type(self, compression_type: int) -> 'AsrRequestHeader':
        self.compression_type = compression_type
        return self

    def with_reserved_data(self, reserved_data: bytes) -> 'AsrRequestHeader':
        self.reserved_data = reserved_data
        return self

    def to_bytes(self) -> bytes:
        header = bytearray()
        header.append((ProtocolVersion.V1 << 4) | 1)
        header.append((self.message_type << 4) | self.message_type_specific_flags)
        header.append((self.serialization_type << 4) | self.compression_type)
        header.extend(self.reserved_data)
        return bytes(header)

    @staticmethod
    def default_header() -> 'AsrRequestHeader':
        return AsrRequestHeader()

class RequestBuilder:
    """请求构建器"""
    
    @staticmethod
    def new_auth_headers(config: AsrConfig) -> Dict[str, str]:
        """创建认证头"""
        reqid = str(uuid.uuid4())
        return {
            "X-Api-Resource-Id": "volc.bigasr.sauc.duration",
            "X-Api-Request-Id": reqid,
            "X-Api-Access-Key": config.access_key,
            "X-Api-App-Key": config.app_key
        }

    @staticmethod
    def new_full_client_request(seq: int) -> bytes:
        """创建完整客户端请求"""
        header = AsrRequestHeader.default_header() \
            .with_message_type_specific_flags(MessageTypeSpecificFlags.POS_SEQUENCE)
        
        payload = {
            "user": {
                "uid": "local_voice_user"
            },
            "audio": {
                "format": "wav",
                "codec": "raw",
                "rate": 16000,
                "bits": 16,
                "channel": 1
            },
            "request": {
                "model_name": "bigmodel",
                "enable_itn": True,
                "enable_punc": True,
                "enable_ddc": True,
                "show_utterances": True,
                "enable_nonstream": False
            }
        }
        
        payload_bytes = json.dumps(payload).encode('utf-8')
        compressed_payload = AudioUtils.gzip_compress(payload_bytes)
        payload_size = len(compressed_payload)
        
        request = bytearray()
        request.extend(header.to_bytes())
        request.extend(struct.pack('>i', seq))
        request.extend(struct.pack('>U', payload_size))
        request.extend(compressed_payload)
        
        return bytes(request)

    @staticmethod
    def new_audio_only_request(seq: int, segment: bytes, is_last: bool = False) -> bytes:
        """创建纯音频请求"""
        header = AsrRequestHeader.default_header()
        if is_last:
            header.with_message_type_specific_flags(MessageTypeSpecificFlags.NEG_WITH_SEQUENCE)
            seq = -seq
        else:
            header.with_message_type_specific_flags(MessageTypeSpecificFlags.POS_SEQUENCE)
        header.with_message_type(MessageType.CLIENT_AUDIO_ONLY_REQUEST)
        
        request = bytearray()
        request.extend(header.to_bytes())
        request.extend(struct.pack('>i', seq))
        
        compressed_segment = AudioUtils.gzip_compress(segment)
        request.extend(struct.pack('>U', len(compressed_segment)))
        request.extend(compressed_segment)
        
        return bytes(request)

class AsrResponse:
    """ASR响应"""
    
    def __init__(self):
        self.code = 0
        self.event = 0
        self.is_last_package = False
        self.payload_sequence = 0
        self.payload_size = 0
        self.payload_msg = None

    def to_dict(self) -> Dict[str, Any]:
        return {
            "code": self.code,
            "event": self.event,
            "is_last_package": self.is_last_package,
            "payload_sequence": self.payload_sequence,
            "payload_size": self.payload_size,
            "payload_msg": self.payload_msg
        }

class ResponseParser:
    """响应解析器"""
    
    @staticmethod
    def parse_response(msg: bytes) -> AsrResponse:
        """解析响应"""
        response = AsrResponse()
        
        header_size = msg[0] & 0x0f
        message_type = msg[1] >> 4
        message_type_specific_flags = msg[1] & 0x0f
        serialization_method = msg[2] >> 4
        message_compression = msg[2] & 0x0f
        
        payload = msg[header_size*4:]
        
        # 解析message_type_specific_flags
        if message_type_specific_flags & 0x01:
            response.payload_sequence = struct.unpack('>i', payload[:4])[0]
            payload = payload[4:]
        if message_type_specific_flags & 0x02:
            response.is_last_package = True
        if message_type_specific_flags & 0x04:
            response.event = struct.unpack('>i', payload[:4])[0]
            payload = payload[4:]
            
        # 解析message_type
        if message_type == MessageType.SERVER_FULL_RESPONSE:
            response.payload_size = struct.unpack('>U', payload[:4])[0]
            payload = payload[4:]
        elif message_type == MessageType.SERVER_ERROR_RESPONSE:
            response.code = struct.unpack('>i', payload[:4])[0]
            response.payload_size = struct.unpack('>U', payload[4:8])[0]
            payload = payload[8:]
            
        if not payload:
            return response
            
        # 解压缩
        if message_compression == CompressionType.GZIP:
            try:
                payload = AudioUtils.gzip_decompress(payload)
            except Exception as e:
                logger.error(f"Failed to decompress payload: {e}")
                return response
                
        # 解析payload
        try:
            if serialization_method == SerializationType.JSON:
                response.payload_msg = json.loads(payload.decode('utf-8'))
        except Exception as e:
            logger.error(f"Failed to parse payload: {e}")
            
        return response

class SpeechRecognizer:
    """语音识别器"""
    
    def __init__(self, app_key: str = None, access_key: str = None, 
                 url: str = "wss://openspeech.bytedance.com/api/v3/sauc/bigmodel_nostream"):
        self.config = AsrConfig(app_key, access_key)
        self.url = url
        self.seq = 1
        
    async def recognize_file(self, file_path: str) -> List[RecognitionResult]:
        """识别音频文件"""
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"Audio file not found: {file_path}")
        
        results = []
        
        try:
            async with aiohttp.ClientSession() as session:
                # 读取音频文件
                with open(file_path, 'rb') as f:
                    content = f.read()
                
                if not AudioUtils.is_wav_file(content):
                    raise ValueError("Audio file must be in WAV format")
                
                # 获取音频信息
                try:
                    _, _, sample_rate, _, audio_data = AudioUtils.read_wav_info(content)
                    if sample_rate != DEFAULT_SAMPLE_RATE:
                        logger.warning(f"Sample rate {sample_rate} != {DEFAULT_SAMPLE_RATE}, may affect recognition accuracy")
                except Exception as e:
                    logger.error(f"Failed to read audio info: {e}")
                    raise
                
                # 计算分段大小 (200ms per segment)
                segment_size = 1 * 2 * DEFAULT_SAMPLE_RATE * 200 // 1000  # channel * bytes_per_sample * sample_rate * duration_ms / 1000
                
                # 创建WebSocket连接
                headers = RequestBuilder.new_auth_headers(self.config)
                async with session.ws_connect(self.url, headers=headers) as ws:
                    
                    # 发送完整客户端请求
                    request = RequestBuilder.new_full_client_request(self.seq)
                    self.seq += 1
                    await ws.send_bytes(request)
                    
                    # 接收初始响应
                    msg = await ws.receive()
                    if msg.type == aiohttp.WSMsgType.BINARY:
                        response = ResponseParser.parse_response(msg.data)
                        logger.info(f"Initial response: {response.to_dict()}")
                    
                    # 分段发送音频数据
                    audio_segments = self._split_audio(audio_data, segment_size)
                    total_segments = len(audio_segments)
                    
                    for i, segment in enumerate(audio_segments):
                        is_last = (i == total_segments - 1)
                        request = RequestBuilder.new_audio_only_request(
                            self.seq, 
                            segment,
                            is_last=is_last
                        )
                        await ws.send_bytes(request)
                        logger.info(f"Sent audio segment {i+1}/{total_segments}")
                        
                        if not is_last:
                            self.seq += 1
                        
                        # 短暂延迟模拟实时流
                        await asyncio.sleep(0.1)
                    
                    # 接收识别结果
                    final_text = ""
                    while True:
                        msg = await ws.receive()
                        if msg.type == aiohttp.WSMsgType.BINARY:
                            response = ResponseParser.parse_response(msg.data)
                            
                            if response.payload_msg and 'text' in response.payload_msg:
                                text = response.payload_msg['text']
                                if text:
                                    final_text += text
                                    
                                    result = RecognitionResult(
                                        text=text,
                                        confidence=0.9,  # 默认置信度
                                        is_final=response.is_last_package
                                    )
                                    results.append(result)
                                    
                                    logger.info(f"Recognized: {text}")
                            
                            if response.is_last_package or response.code != 0:
                                break
                                
                        elif msg.type == aiohttp.WSMsgType.ERROR:
                            logger.error(f"WebSocket error: {msg.data}")
                            break
                        elif msg.type == aiohttp.WSMsgType.CLOSED:
                            logger.info("WebSocket connection closed")
                            break
                    
                    # 如果没有获得最终结果，创建一个包含所有文本的结果
                    if final_text and not any(r.is_final for r in results):
                        final_result = RecognitionResult(
                            text=final_text,
                            confidence=0.9,
                            is_final=True
                        )
                        results.append(final_result)
                    
                    return results
                    
        except Exception as e:
            logger.error(f"Speech recognition failed: {e}")
            raise
    
    def _split_audio(self, data: bytes, segment_size: int) -> List[bytes]:
        """分割音频数据"""
        if segment_size <= 0:
            return []
            
        segments = []
        for i in range(0, len(data), segment_size):
            end = i + segment_size
            if end > len(data):
                end = len(data)
            segments.append(data[i:end])
        return segments
    
    async def recognize_latest_recording(self, directory: str = ".") -> Optional[RecognitionResult]:
        """识别最新的录音文件"""
        # 查找最新的录音文件
        recording_files = [f for f in os.listdir(directory) if f.startswith('recording_') and f.endswith('.wav')]
        
        if not recording_files:
            logger.warning("No recording files found")
            return None
        
        # 按文件名排序（包含时间戳）
        recording_files.sort(reverse=True)
        latest_file = recording_files[0]
        latest_path = os.path.join(directory, latest_file)
        
        logger.info(f"Recognizing latest recording: {latest_file}")
        
        try:
            results = await self.recognize_file(latest_path)
            if results:
                # 返回最终的识别结果
                final_results = [r for r in results if r.is_final]
                if final_results:
                    return final_results[-1]
                else:
                    # 如果没有标记为final的结果，返回最后一个
                    return results[-1]
        except Exception as e:
            logger.error(f"Failed to recognize latest recording: {e}")
        
        return None

async def main():
    """测试函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description="语音识别测试")
    parser.add_argument("--file", type=str, help="音频文件路径")
    parser.add_argument("--latest", action="store_true", help="识别最新的录音文件")
    parser.add_argument("--app-key", type=str, help="SAUC App Key")
    parser.add_argument("--access-key", type=str, help="SAUC Access Key")
    
    args = parser.parse_args()
    
    recognizer = SpeechRecognizer(
        app_key=args.app_key,
        access_key=args.access_key
    )
    
    try:
        if args.latest:
            result = await recognizer.recognize_latest_recording()
            if result:
                print(f"识别结果: {result.text}")
                print(f"置信度: {result.confidence}")
                print(f"最终结果: {result.is_final}")
            else:
                print("未能识别到语音内容")
        elif args.file:
            results = await recognizer.recognize_file(args.file)
            for i, result in enumerate(results):
                print(f"结果 {i+1}: {result.text}")
                print(f"置信度: {result.confidence}")
                print(f"最终结果: {result.is_final}")
                print("-" * 40)
        else:
            print("请指定 --file 或 --latest 参数")
            
    except Exception as e:
        print(f"识别失败: {e}")

if __name__ == "__main__":
    asyncio.run(main())