Update main.py

main.py

@@ -1,14 +1,26 @@
-from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request
-from fastapi.responses import HTMLResponse
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request, UploadFile, File, HTTPException
+from fastapi.responses import HTMLResponse, FileResponse
+from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
 import uuid
 import json
+import os
+from pathlib import Path
 
 app = FastAPI(title="FastAPI Chat")
 templates = Jinja2Templates(directory="templates")
 
+# Create uploads directory
+UPLOAD_DIR = Path("uploads")
+UPLOAD_DIR.mkdir(exist_ok=True)
+
+# Mount static files
+app.mount("/uploads", StaticFiles(directory="uploads"), name="uploads")
+
 # Store active rooms: {room_id: [websocket_connections]}
 active_rooms = {}
+# Store uploaded files info: {room_id: [{filename, original_name, file_type}]}
+room_files = {}
 
 @app.get("/", response_class=HTMLResponse)
 async def home(request: Request):
@@ -17,10 +29,42 @@ async def home(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
+    room_id = str(uuid.uuid4())[:8]
     active_rooms[room_id] = []
+    room_files[room_id] = []
     return {"room_id": room_id, "success": True}
 
+@app.post("/upload-file/{room_id}")
+async def upload_file(room_id: str, file: UploadFile = File(...)):
+    """Handle file uploads for a room"""
+    if room_id not in active_rooms:
+        raise HTTPException(status_code=404, detail="Room not found")
+    
+    # Generate unique filename
+    file_ext = Path(file.filename).suffix
+    unique_name = f"{uuid.uuid4().hex}{file_ext}"
+    file_path = UPLOAD_DIR / unique_name
+    
+    # Save file
+    content = await file.read()
+    with open(file_path, "wb") as f:
+        f.write(content)
+    
+    # Store file info
+    file_info = {
+        "filename": unique_name,
+        "original_name": file.filename,
+        "file_type": file.content_type or "application/octet-stream",
+        "size": len(content)
+    }
+    room_files[room_id].append(file_info)
+    
+    return {
+        "success": True,
+        "file_url": f"/uploads/{unique_name}",
+        "file_info": file_info
+    }
+
 @app.websocket("/ws/{room_id}")
 async def websocket_endpoint(websocket: WebSocket, room_id: str):
     """Handle WebSocket connections for a specific room"""
@@ -47,15 +91,15 @@ async def websocket_endpoint(websocket: WebSocket, room_id: str):
     
     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",
+                "type": message_data.get("type", "message"),
                 "username": message_data.get("username", "Anonymous"),
-                "text": message_data.get("text", "")
+                "text": message_data.get("text", ""),
+                "file": message_data.get("file", None)
             })
             
             for connection in active_rooms[room_id]:
@@ -65,15 +109,20 @@ async def websocket_endpoint(websocket: WebSocket, room_id: str):
                     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]
+                # Clean up files for empty room
+                if room_id in room_files:
+                    for file_info in room_files[room_id]:
+                        try:
+                            (UPLOAD_DIR / file_info["filename"]).unlink()
+                        except:
+                            pass
+                    del room_files[room_id]
             else:
-                # Notify remaining users
                 user_count = len(active_rooms[room_id])
                 leave_message = json.dumps({
                     "type": "system",