import asyncio
import websockets
import soundfile as sf
import uuid

# --- 配置 ---
HOST = "localhost"
PORT = 11096
SESSION_ID = str(uuid.uuid4())
SENDER_URI = f"ws://{HOST}:{PORT}/ws/asr/{SESSION_ID}?mode=sender"
RECEIVER_URI = f"ws://{HOST}:{PORT}/ws/asr/{SESSION_ID}?mode=receiver"

AUDIO_FILE_PATH = "tests/XT_ZZY_denoise.wav" # 确保此测试文件存在且为 16kHz, 16-bit, 单声道
CHUNK_DURATION_MS = 100 # 每次发送100ms的音频数据
CHUNK_SIZE = int(16000 * 2 * CHUNK_DURATION_MS / 1000) # 3200 bytes

async def run_receiver():
    """作为接收者连接，并打印收到的所有消息。"""
    print(f"▶️ [Receiver] 尝试连接到: {RECEIVER_URI}")
    try:
        async with websockets.connect(RECEIVER_URI) as websocket:
            print("✅ [Receiver] 连接成功，等待消息...")
            try:
                while True:
                    message = await websocket.recv()
                    print(f"🎧 [Receiver] 收到结果: {message}")
            except websockets.exceptions.ConnectionClosed as e:
                print(f"✅ [Receiver] 连接已由服务器正常关闭: {e.reason}")
    except Exception as e:
        print(f"❌ [Receiver] 连接失败: {e}")

async def run_sender():
    """
    作为发送者连接，同时负责发送音频和接收自己会话的广播结果。
    """
    await asyncio.sleep(1) # 等待receiver有机会先连接
    print(f"▶️ [Sender] 尝试连接到: {SENDER_URI}")
    try:
        async with websockets.connect(SENDER_URI) as websocket:
            print("✅ [Sender] 连接成功。")
            
            # --- 并行任务：接收消息 ---
            async def receive_task():
                print("▶️ [Sender-Receiver] 开始监听广播消息...")
                try:
                    while True:
                        message = await websocket.recv()
                        print(f"🎧 [Sender-Receiver] 收到结果: {message}")
                except websockets.exceptions.ConnectionClosed:
                    print("✅ [Sender-Receiver] 连接已关闭，停止监听。")

            receiver_sub_task = asyncio.create_task(receive_task())

            # --- 主任务：发送音频 ---
            try:
                print("▶️ [Sender] 准备发送音频...")
                audio_data, sample_rate = sf.read(AUDIO_FILE_PATH, dtype='int16')
                if sample_rate != 16000:
                    print(f"❌ [Sender] 错误：音频文件采样率必须是 16kHz。")
                    receiver_sub_task.cancel()
                    return

                total_samples = len(audio_data)
                chunk_samples = CHUNK_SIZE // 2
                samples_sent = 0
                print(f"音频加载成功，总长度: {total_samples} samples。开始分块发送...")
                
                for i in range(0, total_samples, chunk_samples):
                    chunk = audio_data[i:i + chunk_samples]
                    if len(chunk) == 0:
                        break
                    await websocket.send(chunk.tobytes())
                    samples_sent += len(chunk)
                    print(f"🎧 [Sender] 正在发送: {samples_sent}/{total_samples} samples", end="\r")
                    await asyncio.sleep(CHUNK_DURATION_MS / 1000)
                
                print()
                print("🏁 [Sender] 音频流发送完毕，发送 'close' 信号。")
                await websocket.send("close")

            except FileNotFoundError:
                print(f"❌ [Sender] 错误：找不到音频文件 {AUDIO_FILE_PATH}")
            except Exception as e:
                print(f"❌ [Sender] 发送过程中发生错误: {e}")
            
            # 等待接收任务自然结束（当连接关闭时）
            await receiver_sub_task

    except Exception as e:
        print(f"❌ [Sender] 连接失败: {e}")

async def main():
    """同时运行 sender 和 receiver 任务。"""
    print("--- 开始 WebSocket ASR 服务端到端测试 ---")
    print(f"会话 ID: {SESSION_ID}")
    
    # 创建 receiver 和 sender 任务
    sender_task = asyncio.create_task(run_sender())
    await asyncio.sleep(7)
    receiver_task = asyncio.create_task(run_receiver())
    
    # 等待两个任务完成
    await asyncio.gather(receiver_task, sender_task)
    
    print("--- 测试结束 ---")

if __name__ == "__main__":
    # 在运行此脚本前，请确保 FastAPI 服务器正在运行。
    # python main.py
    asyncio.run(main())