Sebastiankay/eft_group_map_websocket

1from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect2from fastapi.middleware.cors import CORSMiddleware3from fastapi.responses import HTMLResponse4from fastapi.staticfiles import StaticFiles5from fastapi.templating import Jinja2Templates6import json7import uvicorn8 9app = FastAPI()10 11# Configure CORS12app.add_middleware(13    CORSMiddleware,14    allow_origins=["*"],15    allow_credentials=True,16    allow_methods=["*"],17    allow_headers=["*"],18)19 20# Serve static files21app.mount("/static", StaticFiles(directory="static"), name="static")22 23# Set up Jinja2 templates24templates = Jinja2Templates(directory="templates")25 26with open("data.json", "r", encoding="utf8") as file:27    maps_data = json.load(file)28 29 30@app.get("/", response_class=HTMLResponse)31async def read_index(request: Request):32    return templates.TemplateResponse("index.html", {"request": request, "maps": maps_data})33 34 35@app.get("/map/{normalized_name}", response_class=HTMLResponse)36async def get_map_by_normalized_name(request: Request, normalized_name: str):37    map_entry = maps_data.get(normalized_name)38    if map_entry:39        filtered_maps = [map_entry]40    else:41        filtered_maps = []42    return templates.TemplateResponse("map.html", {"request": request, "maps": filtered_maps})43 44 45class ConnectionManager:46    def __init__(self):47        self.groups: dict[str, set[WebSocket]] = {}48        self.connections: dict[WebSocket, str] = {}49 50    async def connect(self, websocket: WebSocket):51        await websocket.accept()52 53    def disconnect(self, websocket: WebSocket):54        if websocket in self.connections:55            group_name = self.connections[websocket]56            self.remove_from_group(websocket, group_name)57 58    async def join_group(self, websocket: WebSocket, group_name: str):59        # Remove from previous group if exists60        if websocket in self.connections:61            old_group = self.connections[websocket]62            self.remove_from_group(websocket, old_group)63        64        # Add to new group65        if group_name not in self.groups:66            self.groups[group_name] = set()67        self.groups[group_name].add(websocket)68        self.connections[websocket] = group_name69        print(f"Client joined group '{group_name}'")70 71    def remove_from_group(self, websocket: WebSocket, group_name: str):72        if group_name in self.groups and websocket in self.groups[group_name]:73            self.groups[group_name].remove(websocket)74            if len(self.groups[group_name]) == 0:75                del self.groups[group_name]76                print(f"Group '{group_name}' deleted (no members)")77        if websocket in self.connections:78            del self.connections[websocket]79        print(f"Client removed from group '{group_name}'")80 81    async def broadcast_to_group(self, message: str, group_name: str):82        if group_name in self.groups:83            for connection in self.groups[group_name]:84                await connection.send_text(message)85 86 87manager = ConnectionManager()88 89 90@app.websocket("/ws")91async def websocket_endpoint(websocket: WebSocket):92    await manager.connect(websocket)93    try:94        while True:95            data = await websocket.receive_text()96            try:97                message = json.loads(data)98                # Handle join messages99                if message.get('type') == 'join':100                    group_name = message.get('group')101                    if group_name:102                        await manager.join_group(websocket, group_name)103                    else:104                        error = json.dumps({105                            "type": "error",106                            "message": "Missing group name in join request"107                        })108                        await websocket.send_text(error)109                110                # Broadcast all other messages to the group111                else:112                    if websocket in manager.connections:113                        group_name = manager.connections[websocket]114                        await manager.broadcast_to_group(data, group_name)115                    else:116                        error = json.dumps({117                            "type": "error",118                            "message": "Join a group before sending messages"119                        })120                        await websocket.send_text(error)121            122            except json.JSONDecodeError:123                # Handle non-JSON messages by broadcasting directly124                if websocket in manager.connections:125                    group_name = manager.connections[websocket]126                    await manager.broadcast_to_group(data, group_name)127    128    except WebSocketDisconnect:129        manager.disconnect(websocket)130 131 132if __name__ == "__main__":133    uvicorn.run(app, host="0.0.0.0", port=8000)