Update main.py

main.py

@@ -1,13 +1,15 @@
-import os
-import uuid
-from fastapi import FastAPI, WebSocket, WebSocketDisconnect
-from fastapi.staticfiles import StaticFiles
+from fastapi import FastAPI, Request, HTTPException
+from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
-from typing import Dict, List
+from pydantic import BaseModel
+from typing import List, Optional
+from datetime import datetime
+import uuid
+import json
+import os
 
 app = FastAPI()
 
-# Allow CORS for development
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -16,80 +18,200 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Store active connections and messages
-class ConnectionManager:
-    def __init__(self):
-        self.active_connections: Dict[str, WebSocket] = {}
-        self.messages: List[dict] = []
-        self.users: Dict[str, dict] = {}
-
-    async def connect(self, websocket: WebSocket, user_id: str, username: str, email: str):
-        await websocket.accept()
-        self.active_connections[user_id] = websocket
-        self.users[user_id] = {"username": username, "email": email}
-        await self.send_user_list()
-
-    def disconnect(self, user_id: str):
-        if user_id in self.active_connections:
-            del self.active_connections[user_id]
-            del self.users[user_id]
-
-    async def send_personal_message(self, message: str, websocket: WebSocket):
-        await websocket.send_text(message)
-
-    async def broadcast(self, message: dict):
-        for connection in self.active_connections.values():
-            await connection.send_json(message)
-
-    async def send_user_list(self):
-        user_list = {
-            "type": "user_list",
-            "data": list(self.users.values())
-        }
-        await self.broadcast(user_list)
-
-manager = ConnectionManager()
-
-@app.websocket("/ws/{user_id}/{username}/{email}")
-async def websocket_endpoint(websocket: WebSocket, user_id: str, username: str, email: str):
-    await manager.connect(websocket, user_id, username, email)
-    try:
-        while True:
-            data = await websocket.receive_json()
-            if data["type"] == "message":
-                message = {
-                    "id": str(uuid.uuid4()),
-                    "sender": {
-                        "id": user_id,
-                        "username": username,
-                        "email": email
-                    },
-                    "content": data["content"],
-                    "timestamp": data["timestamp"],
-                    "type": data["message_type"],
-                    "fileName": data.get("fileName"),
-                    "fileSize": data.get("fileSize")
-                }
-                manager.messages.append(message)
-                await manager.broadcast({
-                    "type": "new_message",
-                    "data": message
-                })
-    except WebSocketDisconnect:
-        manager.disconnect(user_id)
-        await manager.send_user_list()
-
-@app.get("/messages")
-async def get_messages():
-    return manager.messages
-
-@app.get("/users")
-async def get_users():
-    return list(manager.users.values())
-
-# Serve frontend files
-app.mount("/", StaticFiles(directory="../frontend", html=True), name="frontend")
-
-if __name__ == "__main__":
-    import uvicorn
-    
+# In-memory database (for demo purposes)
+users_db = {}
+chats_db = {}
+messages_db = {}
+
+# Helper functions
+def save_data():
+    with open('users.json', 'w') as f:
+        json.dump(users_db, f)
+    with open('chats.json', 'w') as f:
+        json.dump(chats_db, f)
+    with open('messages.json', 'w') as f:
+        json.dump(messages_db, f)
+
+def load_data():
+    global users_db, chats_db, messages_db
+    if os.path.exists('users.json'):
+        with open('users.json', 'r') as f:
+            users_db = json.load(f)
+    if os.path.exists('chats.json'):
+        with open('chats.json', 'r') as f:
+            chats_db = json.load(f)
+    if os.path.exists('messages.json'):
+        with open('messages.json', 'r') as f:
+            messages_db = json.load(f)
+
+load_data()
+
+# Request Models
+class CreateUserRequest(BaseModel):
+    username: str
+    email: str
+
+class SendMessageRequest(BaseModel):
+    sender_id: str
+    receiver_id: Optional[str] = None
+    chat_id: Optional[str] = None
+    content: str
+
+class CreateGroupRequest(BaseModel):
+    creator_id: str
+    group_name: str
+    participants: List[str]
+
+class UpdatePresenceRequest(BaseModel):
+    user_id: str
+    online: bool
+
+# API Endpoints
+@app.post("/api/create_user")
+async def create_user(data: CreateUserRequest):
+    user_id = str(uuid.uuid4())
+    users_db[user_id] = {
+        'username': data.username,
+        'email': data.email,
+        'online': False,
+        'last_seen': None,
+        'created_at': datetime.now().isoformat()
+    }
+    save_data()
+    return {"user_id": user_id, "status": "success"}
+
+@app.get("/api/get_user/{user_id}")
+async def get_user(user_id: str):
+    if user_id in users_db:
+        return {"user": users_db[user_id], "status": "success"}
+    raise HTTPException(status_code=404, detail="User not found")
+
+@app.post("/api/send_message")
+async def send_message(data: SendMessageRequest):
+    message_id = str(uuid.uuid4())
+    timestamp = datetime.now().isoformat()
+    chat_id = data.chat_id
+
+    if not chat_id:
+        participants = sorted([data.sender_id, data.receiver_id])
+        chat_id = 'private_' + '_'.join(participants)
+        if chat_id not in chats_db:
+            chats_db[chat_id] = {
+                'type': 'private',
+                'participants': participants,
+                'created_at': timestamp
+            }
+    elif chat_id not in chats_db:
+        raise HTTPException(status_code=404, detail="Chat not found")
+
+    messages_db[message_id] = {
+        'chat_id': chat_id,
+        'sender_id': data.sender_id,
+        'content': data.content,
+        'timestamp': timestamp,
+        'status': 'sent'
+    }
+    save_data()
+    return {"message_id": message_id, "status": "success"}
+
+@app.get("/api/get_messages")
+async def get_messages(chat_id: str, user_id: str):
+    if chat_id not in chats_db:
+        raise HTTPException(status_code=404, detail="Chat not found")
+    if user_id not in chats_db[chat_id]['participants']:
+        raise HTTPException(status_code=403, detail="Unauthorized")
+
+    chat_messages = [msg for msg in messages_db.values() if msg['chat_id'] == chat_id]
+    chat_messages.sort(key=lambda x: x['timestamp'])
+
+    for msg in chat_messages:
+        if msg['sender_id'] != user_id and msg['status'] == 'sent':
+            msg['status'] = 'delivered'
+
+    save_data()
+    return {"messages": chat_messages}
+
+@app.get("/api/get_user_chats")
+async def get_user_chats(user_id: str):
+    if user_id not in users_db:
+        raise HTTPException(status_code=404, detail="User not found")
+
+    user_chats = []
+    for chat_id, chat in chats_db.items():
+        if user_id in chat['participants']:
+            chat_messages = [msg for msg in messages_db.values() if msg['chat_id'] == chat_id]
+            chat_messages.sort(key=lambda x: x['timestamp'])
+            last_message = chat_messages[-1] if chat_messages else None
+
+            if chat['type'] == 'private':
+                other_user_id = [uid for uid in chat['participants'] if uid != user_id][0]
+                chat_name = users_db.get(other_user_id, {}).get('username', f'User {other_user_id[:4]}')
+            else:
+                chat_name = chat.get('group_name', 'Group Chat')
+
+            user_chats.append({
+                'chat_id': chat_id,
+                'type': chat['type'],
+                'name': chat_name,
+                'participants': chat['participants'],
+                'last_message': last_message,
+                'created_at': chat['created_at']
+            })
+
+    user_chats.sort(key=lambda x: x['last_message']['timestamp'] if x['last_message'] else x['created_at'], reverse=True)
+    return {"chats": user_chats}
+
+@app.post("/api/create_group")
+async def create_group(data: CreateGroupRequest):
+    user_id = data.creator_id
+    group_name = data.group_name
+    participants = data.participants
+
+    if user_id not in participants:
+        participants.append(user_id)
+
+    if len(participants) < 2:
+        raise HTTPException(status_code=400, detail="Group must have at least 2 participants")
+
+    chat_id = 'group_' + str(uuid.uuid4())
+    timestamp = datetime.now().isoformat()
+
+    chats_db[chat_id] = {
+        'type': 'group',
+        'group_name': group_name,
+        'creator_id': user_id,
+        'participants': participants,
+        'created_at': timestamp
+    }
+
+    save_data()
+    return {"chat_id": chat_id, "status": "success"}
+
+@app.post("/api/update_presence")
+async def update_presence(data: UpdatePresenceRequest):
+    user_id = data.user_id
+    online = data.online
+
+    if user_id in users_db:
+        users_db[user_id]['online'] = online
+        users_db[user_id]['last_seen'] = None if online else datetime.now().isoformat()
+        save_data()
+        return {"status": "success"}
+
+    raise HTTPException(status_code=404, detail="User not found")
+
+@app.get("/api/search_users")
+async def search_users(query: str = ""):
+    query = query.lower()
+    matching_users = []
+    for user_id, user in users_db.items():
+        if (query in user['username'].lower() or 
+            query in user['email'].lower() or 
+            query in user_id.lower()):
+            matching_users.append({
+                'user_id': user_id,
+                'username': user['username'],
+                'email': user['email']
+            })
+
+    return {"users": matching_users}