Update main.py

main.py

@@ -10,29 +10,22 @@ from fastapi.middleware.cors import CORSMiddleware
 
 app = FastAPI(title="Antaram Chat")
 
-# --- Fix: Add CORS to allow connections from localhost/cloud ---
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Allows all origins
+    allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# Setup Directories
 BASE_DIR = Path(__file__).resolve().parent
 UPLOAD_DIR = BASE_DIR / "uploads"
 UPLOAD_DIR.mkdir(exist_ok=True)
 
-# Templates
 templates = Jinja2Templates(directory="templates")
-
-# Mount static files (for uploaded images)
 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)
@@ -41,119 +34,90 @@ 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].upper() # Using uppercase for easier reading
+    room_id = str(uuid.uuid4())[:8].upper()
     active_rooms[room_id] = []
     room_files[room_id] = []
-    print(f"Room Created: {room_id}") # Debug log
     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"""
-    # Check if room exists in active_rooms OR room_files (in case WS hasn't connected yet)
     if room_id not in room_files:
-        raise HTTPException(status_code=404, detail="Room not found")
+        # If the room exists in active_rooms but no files yet, initialize list
+        if room_id in active_rooms:
+            room_files[room_id] = []
+        else:
+            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),
-        "file_url": f"/uploads/{unique_name}" # Pre-calculate URL
+        "file_url": f"/uploads/{unique_name}"
     }
     room_files[room_id].append(file_info)
     
-    return {
-        "success": True,
-        "file_info": file_info
-    }
+    return {"success": True, "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"""
-    
-    # 1. Validate Room Existence
     if room_id not in active_rooms:
-        print(f"Connection rejected: Room {room_id} not found")
-        # 1000 = Normal Closure, but usually we imply an error. 
-        # Note: Sending text before accept is risky, so we just close.
-        await websocket.close(code=1008) # 1008 = Policy Violation (Room doesn't exist)
+        await websocket.close(code=1008)
         return
     
-    # 2. Accept Connection
     await websocket.accept()
     active_rooms[room_id].append(websocket)
     
-    # 3. Notify Join
+    # Notify Join
     user_count = len(active_rooms[room_id])
-    system_message = json.dumps({
+    await broadcast_to_room(room_id, json.dumps({
         "type": "system",
-        "message": "A user has joined the channel.",
+        "message": "A user joined.",
         "userCount": user_count
-    })
-    
-    await broadcast_to_room(room_id, system_message)
+    }))
     
     try:
         while True:
-            # 4. Listen for messages
             data = await websocket.receive_text()
             message_data = json.loads(data)
             
-            # Construct broadcast message
-            broadcast_payload = json.dumps({
+            payload = json.dumps({
                 "type": message_data.get("type", "message"),
                 "username": message_data.get("username", "Anonymous"),
                 "text": message_data.get("text", ""),
                 "file": message_data.get("file", None)
             })
-            
-            # 5. Broadcast to ALL users (including sender, for sync)
-            await broadcast_to_room(room_id, broadcast_payload)
+            await broadcast_to_room(room_id, payload)
                     
     except WebSocketDisconnect:
         if room_id in active_rooms and websocket in active_rooms[room_id]:
             active_rooms[room_id].remove(websocket)
-            
             if not active_rooms[room_id]:
-                # If room is empty, clean up
                 del active_rooms[room_id]
-                if room_id in room_files:
-                    # Optional: Delete actual files here if you want temporary storage
-                    del room_files[room_id]
-                print(f"Room {room_id} deleted (empty)")
+                if room_id in room_files: del room_files[room_id]
             else:
-                # Notify remaining users
                 user_count = len(active_rooms[room_id])
-                leave_message = json.dumps({
+                await broadcast_to_room(room_id, json.dumps({
                     "type": "system",
-                    "message": "A user has disconnected.",
+                    "message": "A user left.",
                     "userCount": user_count
-                })
-                await broadcast_to_room(room_id, leave_message)
+                }))
 
 async def broadcast_to_room(room_id: str, message: str):
-    """Helper to send message to all sockets in a room"""
     if room_id in active_rooms:
         for connection in active_rooms[room_id]:
             try:
                 await connection.send_text(message)
-            except Exception as e:
-                print(f"Error broadcasting: {e}")
+            except: pass
 
 if __name__ == "__main__":
     import uvicorn
-    # Run with: python main.py
     uvicorn.run(app, host="0.0.0.0", port=7860)