"""
2.3 双向音频流对话处理器

处理音频双向流：K230→buffer→VAD→ASR→LLM→TTS→K230
"""

import asyncio
import base64
import json
import tempfile
import wave
from pathlib import Path

from src.Module.asr.asr import ASR
from src.Module.llm.llm import StreamingLLM
from src.Module.tts.tts import StreamingTTS
from src.Module.vad.vad import SileroVAD


class AudioStreamHandler:
    """音频流处理器 - 管理单个会话的音频双向流"""

    def __init__(self, websocket, init_data: dict):
        """
        初始化音频流处理器

        Args:
            websocket: WebSocket连接
            init_data: 初始化数据 {
                "type": "audio_stream_init",
                "session_id": "xxx",
                "audio_config": {
                    "sample_rate": 16000,
                    "bit_depth": 16,
                    "channels": 1,
                    "encoding": "pcm"
                }
            }
        """
        self.websocket = websocket
        self.session_id = init_data.get("session_id")
        self.audio_config = init_data.get("audio_config", {})

        # 音频配置参数
        self.sample_rate = self.audio_config.get("sample_rate", 16000)
        self.channels = self.audio_config.get("channels", 1)
        self.bit_depth = self.audio_config.get("bit_depth", 16)

        # 模块初始化
        self.vad = SileroVAD()
        self.asr = ASR()
        self.llm = StreamingLLM()
        self.tts = StreamingTTS(voice='Cherry')

        # 音频缓冲区
        self.audio_buffer = bytearray()
        self.is_speaking = False

    async def init_response(self):
        """发送初始化响应"""
        await self.websocket.send(json.dumps({
            "type": "audio_stream_init_response",
            "success": True,
            "message": "音频流连接已建立",
            "timestamp": ""
        }))
        print(f"[2.3] 会话 {self.session_id} 已初始化")

    async def handle_audio_upload(self, data: dict):
        """
        处理音频上传

        流程：
        1. 解码base64音频数据
        2. 累积到buffer
        3. VAD检测voice_end
        4. 如果检测到结束，触发ASR→LLM→TTS

        Args:
            data: {
                "type": "audio_stream_upload",
                "session_id": "xxx",
                "data": "base64-encoded-audio",
                "sequence": 1
            }
        """
        try:
            audio_b64 = data.get("data", "")
            audio_bytes = base64.b64decode(audio_b64)

            # 累积音频到buffer
            self.audio_buffer.extend(audio_bytes)

            # VAD检测（在线程池执行）
            vad_result = await asyncio.to_thread(self.vad.detect, audio_bytes)

            if vad_result.voice_start:
                self.is_speaking = True
                print(f"[VAD] 会话{self.session_id} 检测到语音开始")

            if vad_result.voice_end and self.is_speaking:
                print(f"[VAD] 会话{self.session_id} 检测到语音结束，开始处理...")
                self.is_speaking = False

                # 触发完整处理流程
                await self.process_user_speech()

                # 清空buffer，等待下一轮输入
                self.audio_buffer.clear()

        except Exception as e:
            print(f"[AudioUploadError] {e}")
            await self.websocket.send(json.dumps({
                "type": "error",
                "message": f"音频处理失败: {str(e)}"
            }))

    async def process_user_speech(self):
        """
        处理用户语音的完整流程

        流程：
        1. 保存音频到临时WAV文件
        2. ASR识别（线程池执行）
        3. LLM生成回复（线程池执行）
        4. TTS合成（线程池执行）
        5. 发送音频到K230
        """
        try:
            # 1. 保存到临时WAV文件
            with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_wav:
                temp_path = temp_wav.name

                # 写入WAV文件（标准格式）
                with wave.open(temp_path, 'wb') as wav:
                    wav.setnchannels(self.channels)
                    wav.setsampwidth(self.bit_depth // 8)
                    wav.setframerate(self.sample_rate)
                    wav.writeframes(self.audio_buffer)

            print(f"[临时文件] {temp_path}, 大小: {len(self.audio_buffer)} bytes")

            # 2. ASR识别（线程池执行）
            asr_result = await asyncio.to_thread(self.asr.recognize, temp_path)

            # 删除临时文件
            Path(temp_path).unlink(missing_ok=True)

            if not asr_result.success:
                print(f"[ASR Error] {asr_result.error}")
                await self.websocket.send(json.dumps({
                    "type": "error",
                    "message": f"ASR识别失败: {asr_result.error}"
                }))
                return

            user_text = asr_result.text
            print(f"[ASR] 识别结果: {user_text}")

            if not user_text.strip():
                print("[ASR] 识别为空")
                return

            # 3&4&5. LLM→TTS双向流→发送
            await self.generate_and_send_response(user_text)

        except Exception as e:
            print(f"[处理语音Error] {e}")
            await self.websocket.send(json.dumps({
                "type": "error",
                "message": f"处理语音失败: {str(e)}"
            }))

    async def generate_and_send_response(self, user_text: str):
        """
        生成并发送响应（LLM→TTS双向流）

        Args:
            user_text: 用户识别的文本
        """

        def text_and_audio_generator():
            """在线程池中执行的生成器：LLM→TTS双向流"""
            llm = StreamingLLM()
            tts = StreamingTTS(voice='Cherry')

            # LLM文本生成器
            def text_generator():
                for chunk in llm.chat(user_text):
                    if chunk.error:
                        raise Exception(f"LLM Error: {chunk.error}")
                    if chunk.content:
                        yield chunk.content

            # TTS音频生成：从LLM文本生成器获取文本
            audio_chunks = []
            for audio_chunk in tts.stream_from_generator(text_generator()):
                if audio_chunk.error:
                    raise Exception(f"TTS Error: {audio_chunk.error}")
                audio_chunks.append(audio_chunk)

            return audio_chunks

        # 在线程池执行
        audio_chunks = await asyncio.to_thread(text_and_audio_generator)

        # 发送音频到K230
        for i, audio_chunk in enumerate(audio_chunks):
            is_final = audio_chunk.is_final or (i == len(audio_chunks) - 1)
            if audio_chunk.data:
                await self.websocket.send(json.dumps({
                    "type": "audio_stream_download",
                    "session_id": self.session_id,
                    "data": base64.b64encode(audio_chunk.data).decode(),
                    "is_final": is_final,
                }))

        print(
            f"[TTS] 会话{self.session_id} 发送完成，"
            f"共 {len(audio_chunks)} 个音频块"
        )

    async def handle_control(self, data: dict):
        """
        处理控制消息

        Args:
            data: {
                "type": "audio_stream_control",
                "session_id": "xxx",
                "action": "pause|resume|end",
                "reason": "..."
            }
        """
        action = data.get("action")

        if action == "end":
            print(f"[控制] 会话{self.session_id} 结束请求")
            self.vad.reset()
            self.audio_buffer.clear()

        elif action == "pause":
            print(f"[控制] 会话{self.session_id} 暂停")

        elif action == "resume":
            print(f"[控制] 会话{self.session_id} 恢复")
