from fastapi import FastAPI, WebSocket, WebSocketDisconnect

from fastapi.middleware.cors import CORSMiddleware

import uvicorn

# 创建 FastAPI 实例

app = FastAPI()

# 配置 CORS (跨域资源共享)

# 允许所有来源访问，实际生产环境中应更严格地配置

app.add_middleware(

    CORSMiddleware,

    allow_origins=["*"],  # 允许所有来源

    allow_credentials=True,

    allow_methods=["*"],  # 允许所有 HTTP 方法

    allow_headers=["*"],  # 允许所有 HTTP 头

)

# 连接管理器类

# 用于管理活跃的 WebSocket 连接，实现消息广播

class ConnectionManager:

    def __init__(self):

        # 存储活跃的 WebSocket 连接列表

        self.active_connections: list[WebSocket] = []

    async def connect(self, websocket: WebSocket):

        """接受连接并将其添加到活跃连接列表"""

        await websocket.accept()

        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):

        """断开连接时从列表中移除"""

        self.active_connections.remove(websocket)

    async def broadcast(self, message: str):

        """向所有活跃连接发送消息"""

        for connection in self.active_connections:

            await connection.send_text(message)

# 实例化连接管理器

manager = ConnectionManager()

# WebSocket 路由 endpoint

# 路径参数 client_id 用于标识客户端

@app.websocket("/ws/{client_id}")

async def websocket_endpoint(websocket: WebSocket, client_id: int):

    # 1. 建立连接

    await manager.connect(websocket)

    try:

        # 2. 循环接收消息

        while True:

            # 等待客户端发送文本消息

            data = await websocket.receive_text()

            # 3. 广播消息给所有用户

            await manager.broadcast(f"Client #{client_id} says: {data}")

    except WebSocketDisconnect:

        # 4. 处理断开连接

        manager.disconnect(websocket)

        await manager.broadcast(f"Client #{client_id} left the chat")

if __name__ == "__main__":

    # 启动 uvicorn 服务器

    # host="0.0.0.0" 允许外部访问

    # port=8000 服务端口

    uvicorn.run(app, host="0.0.0.0", port=8000)