"""
实时语音合成模块 - 基于阿里云 DashScope 通义千问 TTS
核心特性：使用 yield 生成器实现流式音频返回
"""

import os
import base64
import threading
import queue
from typing import Generator, Optional
from pathlib import Path
from dataclasses import dataclass
from dotenv import load_dotenv

import dashscope
from dashscope.audio.qwen_tts_realtime import (
    QwenTtsRealtime,
    QwenTtsRealtimeCallback,
    AudioFormat
)

# 导入配置
from . import ttsconfig as config


@dataclass
class AudioChunk:
    """音频数据块"""
    data: bytes                     # 音频数据 (PCM)
    is_final: bool = False          # 是否为最后一块
    error: Optional[str] = None     # 错误信息


@dataclass
class AudioInfo:
    """音频参数信息"""
    sample_rate: int = config.SAMPLE_RATE
    channels: int = config.CHANNELS
    sample_width: int = config.SAMPLE_WIDTH


class _StreamingCallback(QwenTtsRealtimeCallback):
    """内部流式回调类"""

    def __init__(self):
        self._audio_queue: queue.Queue[AudioChunk] = queue.Queue()
        self._connected = threading.Event()
        self._session_updated = threading.Event()
        self.session_id: Optional[str] = None

    def on_open(self) -> None:
        self._connected.set()

    def on_close(self, close_status_code, close_msg) -> None:
        self._audio_queue.put(AudioChunk(data=b'', is_final=True))

    def on_event(self, response: dict) -> None:
        try:
            event_type = response.get('type', '')

            if event_type == 'session.created':
                self.session_id = response.get('session', {}).get('id')

            elif event_type == 'session.updated':
                self._session_updated.set()

            elif event_type == 'response.audio.delta':
                audio_b64 = response.get('delta', '')
                if audio_b64:
                    audio_bytes = base64.b64decode(audio_b64)
                    self._audio_queue.put(AudioChunk(data=audio_bytes))

            elif event_type == 'session.finished':
                self._audio_queue.put(AudioChunk(data=b'', is_final=True))

            elif event_type == 'error':
                error_msg = response.get('error', {}).get('message', '未知错误')
                self._audio_queue.put(
                    AudioChunk(
                        data=b'',
                        is_final=True,
                        error=error_msg))

        except Exception as e:
            self._audio_queue.put(AudioChunk(data=b'', is_final=True, error=str(e)))

    def get_chunk(self, timeout: float = config.CHUNK_TIMEOUT) -> Optional[AudioChunk]:
        try:
            return self._audio_queue.get(timeout=timeout)
        except queue.Empty:
            return AudioChunk(data=b'', is_final=True, error='超时')

    def wait_connected(self, timeout: float = config.CONNECT_TIMEOUT) -> bool:
        return self._connected.wait(timeout=timeout)


class StreamingTTS:
    """
    流式语音合成类

    使用方式:
        tts = StreamingTTS(voice='Cherry')
        for chunk in tts.stream("你好，世界"):
            if chunk.data:
                audio_player.play(chunk.data)
    """

    # 音频参数信息 (只读)
    audio_info = AudioInfo()

    def __init__(
        self,
        voice: str = None,
        language: str = None,
        speech_rate: float = None,
        volume: int = None,
        pitch_rate: float = None,
        model: str = None,
    ):
        """
        初始化流式TTS

        所有参数可选，未传入时使用 ttsconfig.py 中的默认值

        Args:
            voice: 音色 ('Cherry', 'Serena', 'Ethan', 'Chelsie')
            language: 语言 ('Chinese', 'English', etc.)
            speech_rate: 语速 (0.5-2.0)
            volume: 音量 (0-100)
            pitch_rate: 语调 (0.5-2.0)
            model: 模型名称
        """
        self._load_api_key()

        # 使用传入值或配置文件默认值
        self.voice = voice or config.VOICE
        self.language = language or config.LANGUAGE
        self.speech_rate = self._clamp(speech_rate or config.SPEECH_RATE, 0.5, 2.0)
        self.volume = self._clamp(volume or config.VOLUME, 0, 100)
        self.pitch_rate = self._clamp(pitch_rate or config.PITCH_RATE, 0.5, 2.0)
        self.model = model or config.MODEL

        self._client: Optional[QwenTtsRealtime] = None
        self._callback: Optional[_StreamingCallback] = None

    @staticmethod
    def _clamp(value, min_val, max_val):
        return max(min_val, min(max_val, value))

    def _load_api_key(self) -> None:
        """从 .env 加载 API Key"""
        current_dir = Path(__file__).parent
        for _ in range(5):
            env_path = current_dir / '.env'
            if env_path.exists():
                load_dotenv(env_path)
                break
            current_dir = current_dir.parent

        api_key = os.environ.get('DASHSCOPE_API_KEY')
        if api_key:
            dashscope.api_key = api_key
        else:
            raise ValueError('未找到 DASHSCOPE_API_KEY')

    def stream(
        self,
        text: str,
        voice: str = None,
        language: str = None,
        speech_rate: float = None,
        volume: int = None,
        pitch_rate: float = None,
    ) -> Generator[AudioChunk, None, None]:
        """
        流式合成语音 - 核心接口

        Args:
            text: 要合成的文本
            voice: 临时覆盖音色
            language: 临时覆盖语言
            speech_rate: 临时覆盖语速
            volume: 临时覆盖音量
            pitch_rate: 临时覆盖语调

        Yields:
            AudioChunk: 音频数据块

        Example:
            >>> tts = StreamingTTS()
            >>> for chunk in tts.stream("你好", voice='Serena'):
            ...     process_audio(chunk.data)
        """
        # 合并参数：stream 传入的优先，否则用实例属性
        _voice = voice or self.voice
        _language = language or self.language
        _speech_rate = self._clamp(speech_rate or self.speech_rate, 0.5, 2.0)
        _volume = self._clamp(volume or self.volume, 0, 100)
        _pitch_rate = self._clamp(pitch_rate or self.pitch_rate, 0.5, 2.0)

        self._callback = _StreamingCallback()
        self._client = QwenTtsRealtime(
            model=self.model,
            callback=self._callback,
            url=config.WS_URL
        )

        # 启动发送线程
        send_thread = threading.Thread(
            target=self._send_text,
            args=(text, _voice, _language, _speech_rate, _volume, _pitch_rate),
            daemon=True
        )

        try:
            self._client.connect()
            send_thread.start()

            while True:
                chunk = self._callback.get_chunk()
                if chunk is None:
                    break
                yield chunk
                if chunk.is_final:
                    break

        finally:
            self._cleanup()

    def _send_text(
        self,
        text: str,
        voice: str,
        language: str,
        speech_rate: float,
        volume: int,
        pitch_rate: float
    ) -> None:
        """在后台发送文本"""
        try:
            if not self._callback.wait_connected():
                return

            self._client.update_session(
                voice=voice,
                response_format=AudioFormat.PCM_24000HZ_MONO_16BIT,
                mode=config.MODE,
                language_type=language,
                speech_rate=speech_rate,
                volume=volume,
                pitch_rate=pitch_rate
            )

            self._client.append_text(text)
            self._client.finish()

        except Exception as e:
            self._callback._audio_queue.put(
                AudioChunk(data=b'', is_final=True, error=str(e))
            )

    def _cleanup(self) -> None:
        if self._client:
            try:
                self._client.close()
            except BaseException:
                pass
            self._client = None

    def stream_from_generator(
        self,
        text_generator: Generator[str, None, None],
        voice: str = None,
        language: str = None,
        speech_rate: float = None,
        volume: int = None,
        pitch_rate: float = None,
    ) -> Generator[AudioChunk, None, None]:
        """
        双向流式合成 - 边发送文本边接收音频

        Args:
            text_generator: 文本生成器，逐块产生文本
            voice: 音色
            language: 语言
            speech_rate: 语速
            volume: 音量
            pitch_rate: 语调

        Yields:
            AudioChunk: 音频数据块

        Example:
            >>> def get_text():
            ...     yield "你好，"
            ...     yield "世界！"
            >>> tts = StreamingTTS()
            >>> for chunk in tts.stream_from_generator(get_text()):
            ...     play_audio(chunk.data)
        """
        _voice = voice or self.voice
        _language = language or self.language
        _speech_rate = self._clamp(speech_rate or self.speech_rate, 0.5, 2.0)
        _volume = self._clamp(volume or self.volume, 0, 100)
        _pitch_rate = self._clamp(pitch_rate or self.pitch_rate, 0.5, 2.0)

        self._callback = _StreamingCallback()
        self._client = QwenTtsRealtime(
            model=self.model,
            callback=self._callback,
            url=config.WS_URL
        )

        # 启动发送线程
        send_thread = threading.Thread(
            target=self._send_text_generator,
            args=(
                text_generator,
                _voice,
                _language,
                _speech_rate,
                _volume,
                _pitch_rate),
            daemon=True)

        try:
            self._client.connect()
            send_thread.start()

            while True:
                chunk = self._callback.get_chunk()
                if chunk is None:
                    break
                yield chunk
                if chunk.is_final:
                    break

        finally:
            self._cleanup()

    def _send_text_generator(
        self,
        text_generator: Generator[str, None, None],
        voice: str,
        language: str,
        speech_rate: float,
        volume: int,
        pitch_rate: float
    ) -> None:
        """后台逐块发送文本"""
        try:
            if not self._callback.wait_connected():
                return

            self._client.update_session(
                voice=voice,
                response_format=AudioFormat.PCM_24000HZ_MONO_16BIT,
                mode=config.MODE,
                language_type=language,
                speech_rate=speech_rate,
                volume=volume,
                pitch_rate=pitch_rate
            )

            # 逐块发送文本
            for text_chunk in text_generator:
                if text_chunk:
                    self._client.append_text(text_chunk)

            self._client.finish()

        except Exception as e:
            self._callback._audio_queue.put(
                AudioChunk(data=b'', is_final=True, error=str(e))
            )

    def stream_to_file(self, text: str, output_file: str, **kwargs) -> bool:
        """
        流式合成并保存到文件

        Args:
            text: 要合成的文本
            output_file: 输出文件路径 (.pcm 或 .wav)
            **kwargs: 传递给 stream() 的参数

        Returns:
            是否成功
        """
        audio_data = bytearray()
        error = None

        for chunk in self.stream(text, **kwargs):
            if chunk.error:
                error = chunk.error
                break
            if chunk.data:
                audio_data.extend(chunk.data)

        if error:
            print(f'[TTS Error] {error}')
            return False

        if output_file.endswith('.wav'):
            return self._save_as_wav(bytes(audio_data), output_file)
        else:
            with open(output_file, 'wb') as f:
                f.write(audio_data)
            return True

    def _save_as_wav(self, pcm_data: bytes, wav_file: str) -> bool:
        import wave
        try:
            with wave.open(wav_file, 'wb') as wav:
                wav.setnchannels(config.CHANNELS)
                wav.setsampwidth(config.SAMPLE_WIDTH)
                wav.setframerate(config.SAMPLE_RATE)
                wav.writeframes(pcm_data)
            return True
        except Exception as e:
            print(f'[TTS Error] WAV 保存失败: {e}')
            return False


# ============================================================
# 便捷函数
# ============================================================

def synthesize(text: str, **kwargs) -> Generator[AudioChunk, None, None]:
    """
    便捷的流式合成函数

    Args:
        text: 要合成的文本
        **kwargs: 传递给 StreamingTTS/stream() 的参数

    Yields:
        AudioChunk: 音频数据块
    """
    # 分离实例化参数和 stream 参数
    init_keys = {'model'}
    stream_keys = {'voice', 'language', 'speech_rate', 'volume', 'pitch_rate'}

    init_kwargs = {k: v for k, v in kwargs.items() if k in init_keys}
    stream_kwargs = {k: v for k, v in kwargs.items() if k in stream_keys}

    tts = StreamingTTS(**init_kwargs)
    yield from tts.stream(text, **stream_kwargs)


def synthesize_to_file(text: str, output_file: str, **kwargs) -> bool:
    """
    合成并保存到文件

    Args:
        text: 要合成的文本
        output_file: 输出文件路径
        **kwargs: 传递给 stream() 的参数

    Returns:
        是否成功
    """
    tts = StreamingTTS()
    return tts.stream_to_file(text, output_file, **kwargs)


def get_audio_info() -> AudioInfo:
    """获取音频参数信息"""
    return AudioInfo()
