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

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
}


def pcm_to_wav(pcm_data: bytes, params: dict) -> bytes:
    """将裸 PCM 数据封装成 WAV 格式（带文件头）"""
    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()


async def broadcast_audio(audio_data: bytes):
    """广播音频数据给所有连接的客户端"""
    for connection in active_connections:
        try:
            await connection.send_bytes(audio_data)
        except Exception as e:
            print(f"广播失败：{e}")


@app.websocket("/ws/audio")
async def websocket_audio(websocket: WebSocket):
    await websocket.accept()
    active_connections.append(websocket)

    # 可选：后端本地播放测试（需安装 pyaudio）
    import pyaudio
    p = pyaudio.PyAudio()
    stream = p.open(format=p.get_format_from_width(AUDIO_PARAMS["sample_width"]),
                    channels=AUDIO_PARAMS["channels"],
                    rate=AUDIO_PARAMS["sample_rate"],
                    output=True)

    try:
        while True:
            # 1. 接收前端的裸 PCM 音频帧
            pcm_data = await websocket.receive_bytes()
            # print(f"收到 PCM 帧：{len(pcm_data)} bytes")

            # 2. 可选1：后端本地实时播放（测试用）
            stream.write(pcm_data)

            # 3. 可选2：封装成 WAV 并回传给前端（前端可直接播放）
            # wav_data = pcm_to_wav(pcm_data, AUDIO_PARAMS)
            # await websocket.send_bytes(wav_data)  # 仅返回给当前客户端
            # await broadcast_audio(wav_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:
        # 关闭本地播放流
        stream.stop_stream()
        stream.close()
        p.terminate()


if __name__ == "__main__":
    # 启动服务（确保安装了 websockets：pip install uvicorn[standard]）
    uvicorn.run(app, host="0.0.0.0", port=8000)