import nls
from WxMini.Utils.TokenUtil import *
from WxMini.Milvus.Config.MulvusConfig import *
from pydub import AudioSegment
import io

TTS_URL = "wss://nls-gateway-cn-shanghai.aliyuncs.com/ws/v1"

class TTS:
    def __init__(self, _file=None):
        self._file = _file
        self._f = None
        self._audio_data = bytearray()  # 用于存储生成的音频数据

    def start(self, text):
        """
        生成 TTS 并保存到文件（如果 _file 不为 None）
        """
        self._text = text
        if self._file:
            # 如果指定了文件路径，则保存到文件
            self._f = open(self._file, "wb")
        # https://nls-portal.console.aliyun.com/ttsProjectSetting/default?projectId=530183
        TOKEN = getToken()  # 参考https://help.aliyun.com/document_detail/450255.html获取token
        # 初始化 TTS
        tts = nls.NlsSpeechSynthesizer(
            url=TTS_URL,
            token=TOKEN,
            appkey=APPKEY,
            on_data=self.on_data,
            on_close=self.on_close
        )

        # 同步执行 TTS 生成
        tts.start(self._text, voice="xiaobei", aformat="mp3")

    def generate_audio(self, text):
        """
        生成 TTS 并返回音频数据和时长
        :param text: 要生成音频的文本
        :return: 音频数据（bytes）和时长（秒）
        """
        self._text = text
        self._audio_data = bytearray()  # 重置音频数据
        TOKEN = getToken()  # 参考https://help.aliyun.com/document_detail/450255.html获取token
        # 初始化 TTS
        tts = nls.NlsSpeechSynthesizer(
            url=TTS_URL,
            token=TOKEN,
            appkey=APPKEY,
            on_data=self.on_data,
            on_close=self.on_close
        )

        # 同步执行 TTS 生成
        tts.start(self._text, voice="xiaobei", aformat="mp3")
        audio_bytes = bytes(self._audio_data)  # 获取生成的音频数据

        # 计算音频时长
        audio = AudioSegment.from_file(io.BytesIO(audio_bytes), format="mp3")
        duration = len(audio) / 1000  # 转换为秒

        return audio_bytes, duration  # 返回音频数据和时长

    def on_close(self, *args):
        """
        TTS 生成完成时的回调
        """
        if self._f:
            self._f.close()
            print("TTS 生成完成，文件已关闭")
        else:
            print("TTS 生成完成，音频数据已返回")

    def on_data(self, data, *args):
        """
        接收 TTS 数据的回调
        """
        if self._f:
            self._f.write(data)  # 如果指定了文件路径，则写入文件
        self._audio_data.extend(data)  # 将数据存储到内存中