import asyncio
import websockets
import json
import numpy as np
import sounddevice as sd  # 用于播放音频（pip install sounddevice）

# 配置
WS_URL = "ws://10.10.10.202:50000/ws/tts"
sample_rate = None  # 从服务端获取采样率
audio_buffer = []   # 存储流式音频块

# 播放音频的回调函数（流式播放）
def play_audio_callback(outdata, frames, time, status):
    global audio_buffer
    if status:
        print(status, file=sys.stderr)
    # 从缓冲区取数据
    chunk = np.zeros(frames, dtype=np.float32)
    if len(audio_buffer) >= frames:
        chunk = np.array(audio_buffer[:frames], dtype=np.float32)
        audio_buffer = audio_buffer[frames:]
    outdata[:] = chunk.reshape(-1, 1)

async def tts_client():
    global sample_rate, audio_buffer
    async with websockets.connect(WS_URL) as websocket:
        # 1. 构造推理参数（示例：预训练音色+流式推理）
        tts_request = {
            "tts_text": "我是通义实验室语音团队全新推出的生成式语音大模型，提供舒适自然的语音合成能力。",
            "mode": "3s极速复刻",
            "sft_spk": "",  #
            "stream": True,
            "seed": 123456,
            "speed": 1.0
        }
        # 发送推理请求
        await websocket.send(json.dumps(tts_request))
        print("已发送推理请求")

        # 2. 初始化音频播放流（先等待服务端返回采样率）
        play_stream = None
        try:
            async for message in websocket:
                data = json.loads(message)
                if data["status"] == "start":
                    # 开始接收流式数据，初始化播放
                    sample_rate = data["sample_rate"]
                    print(f"开始合成，采样率：{sample_rate}")
                    # 启动音频播放流（阻塞式，需在后台运行）
                    play_stream = sd.OutputStream(
                        samplerate=sample_rate,
                        channels=1,
                        callback=play_audio_callback,
                        blocksize=1024  # 每次播放的块大小，可调整
                    )
                    play_stream.start()
                elif data["status"] == "stream":
                    # 接收音频块，加入缓冲区
                    audio_chunk = np.array(data["audio_chunk"], dtype=np.float32)
                    audio_buffer.extend(audio_chunk)
                elif data["status"] == "end":
                    # 合成结束，停止播放
                    print("合成完成：", data["msg"])
                    if play_stream:
                        # 等待缓冲区播放完毕
                        while len(audio_buffer) > 0:
                            await asyncio.sleep(0.1)
                        play_stream.stop()
                        play_stream.close()
                    break
                elif data["status"] == "error":
                    print("合成失败：", data["msg"])
                    if play_stream:
                        play_stream.stop()
                        play_stream.close()
                    break
        except Exception as e:
            print("客户端异常：", e)
            if play_stream:
                play_stream.stop()
                play_stream.close()

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