from fastapi import FastAPI, WebSocket, WebSocketDisconnect
import uvicorn
from starlette.middleware.cors import CORSMiddleware
import wave
import io
import asyncio
import websockets
import json

app = FastAPI()

# 跨域配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 存储活跃连接
active_connections = []

# 音频参数（必须和前端一致！）
AUDIO_PARAMS = {
    "sample_rate": 16000,  # 前端设置的 16000
    "channels": 1,         # 单声道
    "sample_width": 2,     # 16位深度（2字节/样本）
    "frame_size": 1024     # 前端设置的 1024
}

# 语音识别服务配置（替换为你的识别服务地址）
ASR_CONFIG = {
    # "host": "192.168.109.129",
    "host": "10.10.10.202",
    "port": 10096,
    "mode": "2pass",        # 识别模式：online/offline/2pass
    "chunk_size": [5, 10, 5],
    "chunk_interval": 10,
    "use_itn": 1,
    "hotwords": ""          # 热词（可选）
}

# 封装 PCM 为 WAV（如果识别服务需要 WAV 格式）
def pcm_to_wav(pcm_data: bytes, params: dict) -> bytes:
    output = io.BytesIO()
    with wave.open(output, 'wb') as wf:
        wf.setnchannels(params["channels"])
        wf.setsampwidth(params["sample_width"])
        wf.setframerate(params["sample_rate"])
        wf.writeframes(pcm_data)
    output.seek(0)
    return output.read()

# 语音识别客户端协程：连接 ASR 服务，发送音频，接收结果
async def asr_client(websocket: WebSocket, audio_queue: asyncio.Queue):
    """
    连接语音识别服务，处理音频转发和结果回传
    :param websocket: 前端的 WebSocket 连接
    :param audio_queue: 音频数据队列（接收前端的 PCM 数据）
    """
    # 1. 连接 ASR 服务
    asr_uri = f"ws://{ASR_CONFIG['host']}:{ASR_CONFIG['port']}"
    try:
        async with websockets.connect(
            asr_uri,
            subprotocols=["binary"],
            ping_interval=None
        ) as asr_ws:
            # 2. 发送 ASR 初始化配置
            init_msg = json.dumps({
                "mode": ASR_CONFIG["mode"],
                "chunk_size": ASR_CONFIG["chunk_size"],
                "chunk_interval": ASR_CONFIG["chunk_interval"],
                "wav_name": "frontend_stream",
                "is_speaking": True,
                "hotwords": ASR_CONFIG["hotwords"],
                "itn": bool(ASR_CONFIG["use_itn"]),
                "audio_fs": AUDIO_PARAMS["sample_rate"]
            })
            await asr_ws.send(init_msg)
            print("已发送 ASR 初始化配置", init_msg)

            # 3. 协程1：从队列读取前端音频，发送给 ASR 服务
            async def send_audio():
                while True:
                    try:
                        # 从队列获取前端的 PCM 数据
                        pcm_data = await audio_queue.get()
                        # 如果识别服务需要 WAV 格式，取消下面注释
                        # wav_data = pcm_to_wav(pcm_data, AUDIO_PARAMS)
                        # await asr_ws.send(wav_data)
                        # 否则直接发送 PCM 裸数据
                        await asr_ws.send(pcm_data)
                        await asyncio.sleep(0.005)  # 控制发送速率
                    except Exception as e:
                        print(f"发送音频到 ASR 失败：{e}")
                        break

            # 4. 协程2：接收 ASR 识别结果，回传给前端
            async def recv_asr_result():
                while True:
                    try:
                        asr_result = await asr_ws.recv()
                        asr_result_json = json.loads(asr_result)
                        # 仅保留核心识别结果（按需调整）
                        result = {
                            "text": asr_result_json.get("text", ""),
                            "mode": asr_result_json.get("mode", ""),
                            "timestamp": asr_result_json.get("timestamp", ""),
                            "is_final": asr_result_json.get("is_final", False)
                        }
                        # 回传给当前前端客户端2
                        print('result', result)
                        ORIGINAL_STR = "Hello 世界！123@#"
                        binary_data = ORIGINAL_STR.encode("utf-8")
                        print(f"原始字符串：{ORIGINAL_STR}")
                        print(f"转二进制后：{binary_data}（长度：{len(binary_data)} 字节）")

                        # 2. 发送二进制数据
                        await websocket.send_bytes(binary_data)
                    except Exception as e:
                        print(f"接收 ASR 结果失败：{e}")
                        break

            # 5. 并发运行两个协程
            await asyncio.gather(send_audio(), recv_asr_result())

    except Exception as e:
        print(f"ASR 客户端连接失败：{e}")
        # 给前端返回错误
        await websocket.send_json({"error": str(e), "text": ""})

@app.websocket("/ws/audio")
async def websocket_audio(websocket: WebSocket):
    await websocket.accept()
    active_connections.append(websocket)
    print(f"客户端连接成功，当前连接数：{len(active_connections)}")

    # 创建音频队列：前端音频 -> 队列 -> ASR 客户端
    audio_queue = asyncio.Queue(maxsize=100)  # 队列大小防止积压

    # 启动 ASR 客户端协程
    asr_task = asyncio.create_task(asr_client(websocket, audio_queue))

    try:
        while True:
            # 1. 接收前端的裸 PCM 音频帧
            pcm_data = await websocket.receive_bytes()
            # 2. 将音频放入队列，交给 ASR 客户端处理
            if not audio_queue.full():
                audio_queue.put_nowait(pcm_data)

    except WebSocketDisconnect:
        active_connections.remove(websocket)
        print(f"客户端断开，当前连接数：{len(active_connections)}")
    except Exception as e:
        if websocket in active_connections:
            active_connections.remove(websocket)
        print(f"处理音频失败：{e}")
    finally:
        # 停止 ASR 客户端协程
        asr_task.cancel()
        await asr_task

if __name__ == "__main__":
    # 启动 FastAPI 服务
    uvicorn.run(app, host="0.0.0.0", port=8000)