Create main.py

main.py

@@ -0,0 +1,87 @@
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request
+from fastapi.responses import HTMLResponse
+from fastapi.templating import Jinja2Templates
+import uuid
+import json
+
+app = FastAPI(title="FastAPI Chat")
+templates = Jinja2Templates(directory="templates")
+
+# Store active rooms: {room_id: [websocket_connections]}
+active_rooms = {}
+
+@app.get("/", response_class=HTMLResponse)
+async def home(request: Request):
+    return templates.TemplateResponse("index.html", {"request": request})
+
+@app.post("/create-room")
+async def create_room():
+    """Create a new chat room and return its ID"""
+    room_id = str(uuid.uuid4())[:8]  # Generate short, unique room ID
+    active_rooms[room_id] = []
+    return {"room_id": room_id, "success": True}
+
+@app.websocket("/ws/{room_id}")
+async def websocket_endpoint(websocket: WebSocket, room_id: str):
+    """Handle WebSocket connections for a specific room"""
+    if room_id not in active_rooms:
+        await websocket.close(code=1008, reason="Room not found")
+        return
+    
+    await websocket.accept()
+    active_rooms[room_id].append(websocket)
+    
+    # Notify all users about the new user
+    user_count = len(active_rooms[room_id])
+    system_message = json.dumps({
+        "type": "system",
+        "message": f"User joined. Total users: {user_count}",
+        "userCount": user_count
+    })
+    
+    for connection in active_rooms[room_id]:
+        try:
+            await connection.send_text(system_message)
+        except:
+            pass
+    
+    try:
+        while True:
+            # Receive message from client
+            data = await websocket.receive_text()
+            message_data = json.loads(data)
+            
+            # Broadcast to all users in the room
+            broadcast_message = json.dumps({
+                "type": "message",
+                "username": message_data.get("username", "Anonymous"),
+                "text": message_data.get("text", "")
+            })
+            
+            for connection in active_rooms[room_id]:
+                try:
+                    await connection.send_text(broadcast_message)
+                except:
+                    pass
+                    
+    except WebSocketDisconnect:
+        # Remove disconnected user
+        if websocket in active_rooms[room_id]:
+            active_rooms[room_id].remove(websocket)
+            
+            # Clean up empty rooms
+            if not active_rooms[room_id]:
+                del active_rooms[room_id]
+            else:
+                # Notify remaining users
+                user_count = len(active_rooms[room_id])
+                leave_message = json.dumps({
+                    "type": "system",
+                    "message": f"User left. Total users: {user_count}",
+                    "userCount": user_count
+                })
+                for connection in active_rooms[room_id]:
+                    try:
+                        await connection.send_text(leave_message)
+                    except:
+                        pass