Update app.py

app.py

@@ -1,49 +1,57 @@
-# app.py
 import os
 import json
 import time
 import uuid
 import base64
-from typing import Optional, Dict, Set
+from typing import Optional, Dict, Set, List, Any
 import asyncio
 import aiosqlite
-import tempfile
-import io
 import logging
 from contextlib import asynccontextmanager
+from datetime import datetime
+from urllib.parse import unquote
 
 from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect, Header, Query, Request, UploadFile, File
 from fastapi.staticfiles import StaticFiles
 from fastapi.responses import Response, JSONResponse, FileResponse
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.middleware.trustedhost import TrustedHostMiddleware
 
 from cryptography.hazmat.primitives.ciphers.aead import AESGCM
 from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
 from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.backends import default_backend
 
-from storage_handler import store_file, retrieve_file, delete_file, store_file_stream
+from storage_handler import store_file, retrieve_file, delete_file, get_bucket_info, file_exists
 
-# --------------------------- Configuration ---------------------------
+# ------------------------------------------------------------------------
+# Configuration
+# ------------------------------------------------------------------------
 DATABASE_URL = os.environ.get("DATABASE_URL", "/data/infinitychat.db")
 MESSAGE_KEY_B64 = os.environ.get("SECRET_KEY", None)
 if MESSAGE_KEY_B64 is None:
-    # Auto-generate (set permanently in production!)
     MESSAGE_KEY_B64 = base64.urlsafe_b64encode(os.urandom(32)).decode()
-    print("WARNING: SECRET_KEY not set – generated random key.")
+    print("WARNING: SECRET_KEY not set – generated random key. Set it permanently in production!")
 MESSAGE_KEY = base64.urlsafe_b64decode(MESSAGE_KEY_B64)
-assert len(MESSAGE_KEY) == 32
+assert len(MESSAGE_KEY) == 32, "SECRET_KEY must decode to exactly 32 bytes"
 
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger("InfinityChat")
 
-# --------------------------- Database initialization ---------------------------
+# ------------------------------------------------------------------------
+# Database initialization
+# ------------------------------------------------------------------------
 async def init_database():
-    """Initialize the SQLite database with proper async context."""
+    """Initialize SQLite database with all required tables and indexes."""
     async with aiosqlite.connect(DATABASE_URL) as db:
         db.row_factory = aiosqlite.Row
         await db.execute("PRAGMA journal_mode=WAL;")
         await db.execute("PRAGMA foreign_keys=ON;")
+        await db.execute("PRAGMA synchronous=NORMAL;")
+        await db.execute("PRAGMA cache_size=-20000;")  # 20MB cache
+        await db.execute("PRAGMA busy_timeout=5000;")
         
+        # Users table
         await db.execute("""
             CREATE TABLE IF NOT EXISTS users (
                 id INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -53,9 +61,13 @@ async def init_database():
                 salt TEXT NOT NULL,
                 avatar_path TEXT,
                 token TEXT UNIQUE,
-                created_at INTEGER DEFAULT (strftime('%s','now'))
+                created_at INTEGER DEFAULT (strftime('%s','now')),
+                last_seen INTEGER DEFAULT (strftime('%s','now')),
+                status TEXT DEFAULT 'offline'
             )
         """)
+        
+        # Messages table with optimized indexes
         await db.execute("""
             CREATE TABLE IF NOT EXISTS messages (
                 id INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -68,72 +80,140 @@ async def init_database():
                 file_path TEXT,
                 file_type TEXT,
                 file_name TEXT,
+                file_size INTEGER DEFAULT 0,
                 FOREIGN KEY(sender_id) REFERENCES users(id)
             )
         """)
+        
+        # Read receipts
         await db.execute("""
             CREATE TABLE IF NOT EXISTS read_receipts (
                 user_id INTEGER NOT NULL,
                 message_id INTEGER NOT NULL,
-                PRIMARY KEY (user_id, message_id)
+                read_at INTEGER DEFAULT (strftime('%s','now')),
+                PRIMARY KEY (user_id, message_id),
+                FOREIGN KEY(user_id) REFERENCES users(id),
+                FOREIGN KEY(message_id) REFERENCES messages(id)
             )
         """)
-        await db.execute("CREATE INDEX IF NOT EXISTS idx_msg_time ON messages(timestamp_ms)")
+        
+        # Conversation index for faster queries
+        await db.execute("""
+            CREATE INDEX IF NOT EXISTS idx_messages_sender_time 
+            ON messages(sender_id, timestamp_ms)
+        """)
+        await db.execute("""
+            CREATE INDEX IF NOT EXISTS idx_messages_time_id 
+            ON messages(timestamp_ms, id)
+        """)
+        await db.execute("""
+            CREATE INDEX IF NOT EXISTS idx_messages_not_deleted 
+            ON messages(is_deleted, id)
+        """)
+        await db.execute("""
+            CREATE INDEX IF NOT EXISTS idx_receipts_user_msg 
+            ON read_receipts(user_id, message_id)
+        """)
+        
         await db.commit()
     
     logger.info("✅ Database initialized successfully")
 
-# --------------------------- Lifespan handler (proper async startup) ---------------------------
+# ------------------------------------------------------------------------
+# FastAPI app with lifespan
+# ------------------------------------------------------------------------
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    """Handle startup and shutdown events properly."""
-    # Startup
+    """Handle startup and shutdown events."""
     await init_database()
     yield
-    # Shutdown
-    # Cleanup if needed
 
-# --------------------------- FastAPI app with lifespan ---------------------------
-app = FastAPI(title="InfinityChat", lifespan=lifespan)
+app = FastAPI(
+    title="InfinityChat",
+    version="1.0.0",
+    lifespan=lifespan,
+    docs_url=None,
+    redoc_url=None
+)
+
+# Middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Static files
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
-# --------------------------- Database helpers ---------------------------
-async def get_db():
-    """Get a database connection with proper async context."""
-    db = await aiosqlite.connect(DATABASE_URL)
-    db.row_factory = aiosqlite.Row
-    await db.execute("PRAGMA journal_mode=WAL")
-    await db.execute("PRAGMA foreign_keys=ON")
-    return db
-
-# --------------------------- Encryption helpers ---------------------------
-def encrypt_message(plain: str) -> str:
+# ------------------------------------------------------------------------
+# Database connection pool
+# ------------------------------------------------------------------------
+class DatabasePool:
+    """Simple database connection management."""
+    
+    def __init__(self):
+        self._conn = None
+    
+    async def get_connection(self) -> aiosqlite.Connection:
+        if self._conn is None:
+            self._conn = await aiosqlite.connect(DATABASE_URL)
+            self._conn.row_factory = aiosqlite.Row
+            await self._conn.execute("PRAGMA journal_mode=WAL")
+            await self._conn.execute("PRAGMA foreign_keys=ON")
+            await self._conn.execute("PRAGMA synchronous=NORMAL")
+        return self._conn
+    
+    async def close(self):
+        if self._conn:
+            await self._conn.close()
+            self._conn = None
+
+db_pool = DatabasePool()
+
+async def get_db() -> aiosqlite.Connection:
+    """Get or create a database connection."""
+    return await db_pool.get_connection()
+
+# ------------------------------------------------------------------------
+# Encryption (AES-256-GCM)
+# ------------------------------------------------------------------------
+def encrypt_message(plaintext: str) -> str:
     """Encrypt message content with AES-256-GCM."""
     aesgcm = AESGCM(MESSAGE_KEY)
     nonce = os.urandom(12)
-    ct = aesgcm.encrypt(nonce, plain.encode(), None)
-    return base64.urlsafe_b64encode(nonce + ct).decode()
+    ciphertext = aesgcm.encrypt(nonce, plaintext.encode('utf-8'), None)
+    return base64.urlsafe_b64encode(nonce + ciphertext).decode('utf-8')
 
 def decrypt_message(encrypted_b64: str) -> str:
     """Decrypt message content."""
-    raw = base64.urlsafe_b64decode(encrypted_b64)
-    nonce, ct = raw[:12], raw[12:]
-    aesgcm = AESGCM(MESSAGE_KEY)
-    return aesgcm.decrypt(nonce, ct, None).decode()
+    try:
+        raw = base64.urlsafe_b64decode(encrypted_b64)
+        nonce, ciphertext = raw[:12], raw[12:]
+        aesgcm = AESGCM(MESSAGE_KEY)
+        return aesgcm.decrypt(nonce, ciphertext, None).decode('utf-8')
+    except Exception as e:
+        logger.error(f"Decryption failed: {e}")
+        return "[Decryption Error]"
 
-# --------------------------- Password handling ---------------------------
+# ------------------------------------------------------------------------
+# Password hashing (PBKDF2)
+# ------------------------------------------------------------------------
 def hash_password(password: str, salt: Optional[str] = None) -> tuple[str, str]:
-    """Hash password with PBKDF2."""
+    """Hash password with PBKDF2-SHA256."""
     if salt is None:
-        salt = os.urandom(16).hex()
+        salt = os.urandom(32).hex()
     kdf = PBKDF2HMAC(
         algorithm=hashes.SHA256(),
         length=32,
-        salt=salt.encode(),
+        salt=salt.encode('utf-8'),
         iterations=600000,
         backend=default_backend()
     )
-    return base64.urlsafe_b64encode(kdf.derive(password.encode())).decode(), salt
+    key = base64.urlsafe_b64encode(kdf.derive(password.encode('utf-8'))).decode('utf-8')
+    return key, salt
 
 def verify_password(password: str, salt: str, stored_hash: str) -> bool:
     """Verify password against stored hash."""
@@ -142,447 +222,874 @@ def verify_password(password: str, salt: str, stored_hash: str) -> bool:
 
 def generate_token() -> str:
     """Generate a secure random token."""
-    return base64.urlsafe_b64encode(os.urandom(32)).decode()
+    return base64.urlsafe_b64encode(os.urandom(32)).decode('utf-8')
 
-# --------------------------- User lookup ---------------------------
-async def get_user_by_token(token: str):
-    """Get user by token with proper async context."""
+# ------------------------------------------------------------------------
+# Authentication helpers
+# ------------------------------------------------------------------------
+async def authenticate_user(token: str) -> Optional[dict]:
+    """Validate token and return user data."""
+    if not token or token in ('null', 'undefined'):
+        return None
+    
     db = await get_db()
     try:
-        async with db.execute("SELECT * FROM users WHERE token = ?", (token,)) as cur:
-            return await cur.fetchone()
-    finally:
-        await db.close()
+        async with db.execute(
+            "SELECT id, username, display_name, avatar_path, status FROM users WHERE token = ?",
+            (token,)
+        ) as cursor:
+            user = await cursor.fetchone()
+            if user:
+                # Update last_seen
+                await db.execute(
+                    "UPDATE users SET last_seen = ? WHERE id = ?",
+                    (int(time.time()), user['id'])
+                )
+                await db.commit()
+                return dict(user)
+            return None
+    except Exception as e:
+        logger.error(f"Authentication error: {e}")
+        return None
 
-async def get_user_by_username(username: str):
-    """Get user by username."""
-    db = await get_db()
-    try:
-        async with db.execute("SELECT * FROM users WHERE username = ?", (username,)) as cur:
-            return await cur.fetchone()
-    finally:
-        await db.close()
+# ------------------------------------------------------------------------
+# HTTP Endpoints
+# ------------------------------------------------------------------------
+@app.get("/api/health")
+async def health_check():
+    """Health check endpoint."""
+    return {
+        "status": "healthy",
+        "timestamp": int(time.time()),
+        "version": "1.0.0"
+    }
 
-# --------------------------- HTTP endpoints ---------------------------
-@app.post("/signup")
-async def signup(username: str, password: str, display_name: str = ""):
+@app.post("/api/auth/signup")
+async def signup(
+    username: str = Query(..., min_length=3, max_length=30),
+    password: str = Query(..., min_length=6),
+    display_name: Optional[str] = Query(None, max_length=50)
+):
     """Create a new user account."""
+    # Validate username format
+    if not username.isalnum() and '_' not in username:
+        raise HTTPException(400, "Username can only contain letters, numbers, and underscores")
+    
     db = await get_db()
     try:
-        existing = await get_user_by_username(username)
-        if existing:
-            raise HTTPException(400, "Username already taken")
+        # Check if username exists
+        async with db.execute(
+            "SELECT id FROM users WHERE username = ?", (username.lower(),)
+        ) as cursor:
+            if await cursor.fetchone():
+                raise HTTPException(409, "Username is already taken")
         
+        # Create user
         pwd_hash, salt = hash_password(password)
         token = generate_token()
+        display = display_name or username
         
-        await db.execute(
-            "INSERT INTO users (username, display_name, password_hash, salt, token) VALUES (?, ?, ?, ?, ?)",
-            (username, display_name, pwd_hash, salt, token)
+        cursor = await db.execute(
+            """INSERT INTO users (username, display_name, password_hash, salt, token, status)
+               VALUES (?, ?, ?, ?, ?, 'online')""",
+            (username.lower(), display, pwd_hash, salt, token)
         )
         await db.commit()
+        user_id = cursor.lastrowid
+        
+        logger.info(f"New user registered: {username} (ID: {user_id})")
         
-        return {"token": token, "username": username, "display_name": display_name}
+        return {
+            "token": token,
+            "user": {
+                "id": user_id,
+                "username": username,
+                "display_name": display,
+                "avatar_path": None,
+                "status": "online"
+            }
+        }
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Signup error: {e}")
+        raise HTTPException(500, "Internal server error")
     finally:
         await db.close()
 
-@app.post("/login")
-async def login(username: str, password: str):
+@app.post("/api/auth/login")
+async def login(
+    username: str = Query(...),
+    password: str = Query(...)
+):
     """Authenticate user and return token."""
     db = await get_db()
     try:
-        user = await get_user_by_username(username)
-        if not user or not verify_password(password, user["salt"], user["password_hash"]):
-            raise HTTPException(400, "Invalid credentials")
+        async with db.execute(
+            "SELECT * FROM users WHERE username = ?", (username.lower(),)
+        ) as cursor:
+            user = await cursor.fetchone()
+        
+        if not user:
+            # Use constant-time comparison to prevent timing attacks
+            verify_password("dummy", os.urandom(32).hex(), 
+                          base64.urlsafe_b64encode(os.urandom(32)).decode())
+            raise HTTPException(401, "Invalid username or password")
         
+        if not verify_password(password, user['salt'], user['password_hash']):
+            raise HTTPException(401, "Invalid username or password")
+        
+        # Generate new token
         token = generate_token()
-        await db.execute("UPDATE users SET token = ? WHERE id = ?", (token, user["id"]))
+        await db.execute(
+            "UPDATE users SET token = ?, last_seen = ?, status = 'online' WHERE id = ?",
+            (token, int(time.time()), user['id'])
+        )
         await db.commit()
         
+        logger.info(f"User logged in: {user['username']}")
+        
         return {
             "token": token,
-            "username": user["username"],
-            "display_name": user["display_name"],
-            "avatar_path": user["avatar_path"]
+            "user": {
+                "id": user['id'],
+                "username": user['username'],
+                "display_name": user['display_name'],
+                "avatar_path": user['avatar_path'],
+                "status": "online"
+            }
         }
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Login error: {e}")
+        raise HTTPException(500, "Internal server error")
     finally:
         await db.close()
 
-@app.get("/me")
-async def me(token: str = Header(...)):
-    """Get current user info."""
-    user = await get_user_by_token(token)
+@app.get("/api/auth/verify")
+async def verify_token(token: str = Header(..., alias="X-Auth-Token")):
+    """Verify authentication token."""
+    user = await authenticate_user(token)
     if not user:
-        raise HTTPException(401, "Invalid token")
-    return {
-        "id": user["id"],
-        "username": user["username"],
-        "display_name": user["display_name"],
-        "avatar_path": user["avatar_path"]
-    }
+        raise HTTPException(401, "Invalid or expired token")
+    return {"user": user}
+
+@app.post("/api/auth/logout")
+async def logout(token: str = Header(..., alias="X-Auth-Token")):
+    """Logout user by invalidating token."""
+    db = await get_db()
+    try:
+        await db.execute(
+            "UPDATE users SET token = NULL, status = 'offline', last_seen = ? WHERE token = ?",
+            (int(time.time()), token)
+        )
+        await db.commit()
+        return {"status": "logged_out"}
+    finally:
+        await db.close()
 
-@app.post("/update_profile")
-async def update_profile(display_name: str = None, token: str = Header(...)):
+# ------------------------------------------------------------------------
+# Profile Endpoints
+# ------------------------------------------------------------------------
+@app.get("/api/profile")
+async def get_profile(token: str = Header(..., alias="X-Auth-Token")):
+    """Get current user profile."""
+    user = await authenticate_user(token)
+    if not user:
+        raise HTTPException(401)
+    return {"user": user}
+
+@app.patch("/api/profile")
+async def update_profile(
+    display_name: Optional[str] = Query(None, max_length=50),
+    token: str = Header(..., alias="X-Auth-Token")
+):
     """Update user profile."""
-    user = await get_user_by_token(token)
+    user = await authenticate_user(token)
     if not user:
         raise HTTPException(401)
     
-    if display_name is not None:
-        db = await get_db()
-        try:
-            await db.execute("UPDATE users SET display_name = ? WHERE id = ?", (display_name, user["id"]))
+    db = await get_db()
+    try:
+        if display_name:
+            await db.execute(
+                "UPDATE users SET display_name = ? WHERE id = ?",
+                (display_name, user['id'])
+            )
             await db.commit()
-        finally:
-            await db.close()
-    
-    return {"status": "ok"}
+            user['display_name'] = display_name
+        
+        return {"user": user}
+    finally:
+        await db.close()
 
-@app.post("/upload_avatar")
-async def upload_avatar(file: UploadFile = File(...), token: str = Header(...)):
+@app.post("/api/profile/avatar")
+async def upload_avatar(
+    file: UploadFile = File(...),
+    token: str = Header(..., alias="X-Auth-Token")
+):
     """Upload user avatar."""
-    user = await get_user_by_token(token)
+    user = await authenticate_user(token)
     if not user:
         raise HTTPException(401)
     
+    # Validate file type
+    allowed_types = {'image/jpeg', 'image/png', 'image/gif', 'image/webp'}
+    if file.content_type not in allowed_types:
+        raise HTTPException(400, "Only JPEG, PNG, GIF, and WebP images are allowed")
+    
+    # Validate file size (max 5MB)
     data = await file.read()
-    ext = os.path.splitext(file.filename)[1] if file.filename else ".jpg"
+    if len(data) > 5 * 1024 * 1024:
+        raise HTTPException(400, "Image must be less than 5MB")
+    
+    # Generate unique path
+    ext = os.path.splitext(file.filename)[1] if file.filename else '.jpg'
     remote_path = f"avatars/{user['username']}_{uuid.uuid4().hex}{ext}"
     
-    store_file(remote_path, data)
+    try:
+        store_file(remote_path, data)
+    except Exception as e:
+        logger.error(f"Avatar upload failed: {e}")
+        raise HTTPException(500, "Failed to upload avatar")
     
     db = await get_db()
     try:
-        await db.execute("UPDATE users SET avatar_path = ? WHERE id = ?", (remote_path, user["id"]))
+        await db.execute(
+            "UPDATE users SET avatar_path = ? WHERE id = ?",
+            (remote_path, user['id'])
+        )
         await db.commit()
+        user['avatar_path'] = remote_path
     finally:
         await db.close()
     
     return {"avatar_path": remote_path}
 
-# --------------------------- Chunked upload ---------------------------
-upload_sessions: Dict[str, dict] = {}
+# ------------------------------------------------------------------------
+# File Upload (Chunked)
+# ------------------------------------------------------------------------
+upload_sessions: Dict[str, Dict] = {}
 
-@app.post("/upload_chunk")
+@app.post("/api/upload/chunk")
 async def upload_chunk(
     file: UploadFile = File(...),
-    chunk_index: int = Query(...),
-    total_chunks: int = Query(...),
-    file_name: str = Query(...),
-    file_type: str = Query(...),
-    upload_id: str = Query(...),
-    token: str = Header(...)
+    chunk_index: int = Query(..., ge=0),
+    total_chunks: int = Query(..., ge=1),
+    file_name: str = Query(..., max_length=255),
+    file_type: str = Query(..., max_length=100),
+    file_size: int = Query(..., ge=0),
+    upload_id: str = Query(..., min_length=8),
+    token: str = Header(..., alias="X-Auth-Token")
 ):
     """Handle chunked file upload."""
-    user = await get_user_by_token(token)
+    user = await authenticate_user(token)
     if not user:
         raise HTTPException(401)
     
+    # Validate file size (max 100MB total)
+    if file_size > 100 * 1024 * 1024:
+        raise HTTPException(400, "File must be less than 100MB")
+    
+    # Initialize upload session if new
     if upload_id not in upload_sessions:
         upload_sessions[upload_id] = {
             "chunks": {},
             "filename": file_name,
             "file_type": file_type,
+            "file_size": file_size,
             "total_chunks": total_chunks,
-            "user_id": user["id"],
-            "received": 0
+            "received": 0,
+            "user_id": user['id'],
+            "created_at": time.time()
         }
     
-    ses = upload_sessions[upload_id]
+    session = upload_sessions[upload_id]
     
-    if chunk_index in ses["chunks"]:
-        return {"status": "duplicate"}
+    # Verify ownership
+    if session['user_id'] != user['id']:
+        raise HTTPException(403, "Upload session belongs to another user")
     
-    ses["chunks"][chunk_index] = await file.read()
-    ses["received"] += 1
+    # Store chunk
+    chunk_data = await file.read()
+    if chunk_index in session['chunks']:
+        return {"status": "duplicate", "received": session['received'], "total": total_chunks}
     
-    if ses["received"] == ses["total_chunks"]:
-        # All chunks received, reassemble
-        full_data = b"".join(ses["chunks"][i] for i in sorted(ses["chunks"]))
-        remote_path = f"uploads/{user['username']}_{uuid.uuid4().hex}/{file_name}"
+    session['chunks'][chunk_index] = chunk_data
+    session['received'] += 1
+    
+    # Check if upload complete
+    if session['received'] == total_chunks:
+        # Reassemble file
+        full_data = b"".join(session['chunks'][i] for i in sorted(session['chunks'].keys()))
+        
+        # Generate storage path
+        remote_path = f"uploads/{user['username']}/{uuid.uuid4().hex}/{file_name}"
         
         try:
             store_file(remote_path, full_data)
         except Exception as e:
-            logger.error(f"Failed to store file: {e}")
+            logger.error(f"Failed to store uploaded file: {e}")
             del upload_sessions[upload_id]
-            raise HTTPException(500, "Storage error")
+            raise HTTPException(500, "Failed to store file")
         
+        # Cleanup
         del upload_sessions[upload_id]
         
+        # Cleanup old sessions (older than 30 minutes)
+        current_time = time.time()
+        expired = [uid for uid, s in upload_sessions.items() 
+                  if current_time - s['created_at'] > 1800]
+        for uid in expired:
+            del upload_sessions[uid]
+        
         return {
             "status": "complete",
             "file_path": remote_path,
             "file_name": file_name,
-            "file_type": file_type
+            "file_type": file_type,
+            "file_size": file_size
         }
     
-    return {"status": "chunk_ok", "received": ses["received"], "total": ses["total_chunks"]}
+    return {
+        "status": "in_progress",
+        "received": session['received'],
+        "total": total_chunks
+    }
 
-# --------------------------- File download ---------------------------
-@app.get("/download/{file_path:path}")
-async def download(file_path: str, token: str = Header(...)):
-    """Download a file (authenticated proxy)."""
-    user = await get_user_by_token(token)
+# ------------------------------------------------------------------------
+# File Download
+# ------------------------------------------------------------------------
+@app.get("/api/download/{file_path:path}")
+async def download_file(
+    file_path: str,
+    token: str = Header(..., alias="X-Auth-Token")
+):
+    """Download a file with authentication."""
+    user = await authenticate_user(token)
     if not user:
         raise HTTPException(401)
     
+    # Security: prevent path traversal
+    if '..' in file_path or file_path.startswith('/'):
+        raise HTTPException(400, "Invalid file path")
+    
     try:
         data = retrieve_file(file_path)
+    except FileNotFoundError:
+        raise HTTPException(404, "File not found")
     except Exception as e:
-        raise HTTPException(404, "File not found or decryption failed")
+        logger.error(f"File download failed: {file_path} - {e}")
+        raise HTTPException(500, "Failed to retrieve file")
     
     import mimetypes
     media_type = mimetypes.guess_type(file_path)[0] or "application/octet-stream"
-    return Response(content=data, media_type=media_type)
-
-@app.get("/")
-async def root():
-    """Serve the main page."""
-    return FileResponse("static/index.html")
+    filename = os.path.basename(file_path)
+    
+    return Response(
+        content=data,
+        media_type=media_type,
+        headers={
+            "Content-Disposition": f'inline; filename="{filename}"',
+            "Cache-Control": "private, max-age=3600"
+        }
+    )
 
-# --------------------------- WebSocket Manager ---------------------------
+# ------------------------------------------------------------------------
+# WebSocket Connection Manager
+# ------------------------------------------------------------------------
 class ConnectionManager:
+    """Manage WebSocket connections with per-user tracking."""
+    
     def __init__(self):
-        self.active: Dict[str, Set[WebSocket]] = {}
-    
-    async def connect(self, ws: WebSocket, username: str):
-        await ws.accept()
-        self.active.setdefault(username, set()).add(ws)
-    
-    def disconnect(self, ws: WebSocket, username: str):
-        if username in self.active:
-            self.active[username].discard(ws)
-            if not self.active[username]:
-                del self.active[username]
-    
-    async def broadcast(self, msg: dict, exclude: WebSocket = None):
-        for uname, ws_set in self.active.items():
-            for w in list(ws_set):
-                if w == exclude:
-                    continue
+        self.active_connections: Dict[int, Dict[str, Any]] = {}  # user_id -> {username, connections: set}
+    
+    async def connect(self, websocket: WebSocket, user_id: int, username: str):
+        """Add a new connection."""
+        await websocket.accept()
+        
+        if user_id not in self.active_connections:
+            self.active_connections[user_id] = {
+                "username": username,
+                "connections": set()
+            }
+        
+        self.active_connections[user_id]["connections"].add(websocket)
+    
+    def disconnect(self, websocket: WebSocket, user_id: int):
+        """Remove a connection."""
+        if user_id in self.active_connections:
+            self.active_connections[user_id]["connections"].discard(websocket)
+            if not self.active_connections[user_id]["connections"]:
+                del self.active_connections[user_id]
+    
+    async def broadcast(self, message: dict, exclude_user_id: Optional[int] = None):
+        """Broadcast message to all connected users."""
+        for uid, info in self.active_connections.items():
+            if uid == exclude_user_id:
+                continue
+            for ws in list(info["connections"]):
+                try:
+                    await ws.send_json(message)
+                except Exception:
+                    self.disconnect(ws, uid)
+    
+    async def send_to_user(self, user_id: int, message: dict):
+        """Send message to a specific user."""
+        if user_id in self.active_connections:
+            for ws in list(self.active_connections[user_id]["connections"]):
                 try:
-                    await w.send_json(msg)
+                    await ws.send_json(message)
                 except Exception:
-                    self.disconnect(w, uname)
+                    self.disconnect(ws, user_id)
+    
+    def get_online_users(self) -> List[Dict]:
+        """Get list of online users with their info."""
+        users = []
+        for uid, info in self.active_connections.items():
+            users.append({
+                "id": uid,
+                "username": info["username"]
+            })
+        return users
     
-    def online_users(self) -> list:
-        return list(self.active.keys())
+    def is_online(self, user_id: int) -> bool:
+        """Check if a user is online."""
+        return user_id in self.active_connections
 
 manager = ConnectionManager()
 
-# --------------------------- WebSocket endpoint ---------------------------
+# ------------------------------------------------------------------------
+# WebSocket Endpoint
+# ------------------------------------------------------------------------
 @app.websocket("/ws")
-async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
-    """WebSocket endpoint for real-time communication."""
-    if not token:
-        await ws.close(code=4001)
-        return
-    
-    user = await get_user_by_token(token)
+async def websocket_endpoint(websocket: WebSocket, token: str = Query(...)):
+    """WebSocket endpoint for real-time messaging."""
+    # Authenticate
+    user = await authenticate_user(token)
     if not user:
-        await ws.close(code=4001, reason="Invalid token")
+        await websocket.close(code=4001, reason="Authentication failed")
         return
     
-    username = user["username"]
-    uid = user["id"]
+    user_id = user['id']
+    username = user['username']
+    display_name = user['display_name']
+    avatar_path = user['avatar_path']
+    
+    # Connect
+    await manager.connect(websocket, user_id, username)
     
-    await manager.connect(ws, username)
+    # Send initial data
+    online_users = manager.get_online_users()
+    await websocket.send_json({
+        "type": "connection_established",
+        "user_id": user_id,
+        "username": username,
+        "display_name": display_name,
+        "avatar_path": avatar_path,
+        "online_users": online_users
+    })
     
-    # Send online users
-    online = manager.online_users()
-    await ws.send_json({"type": "online_users", "users": online})
-    await manager.broadcast({"type": "user_online", "username": username}, exclude=ws)
+    # Broadcast user online status
+    await manager.broadcast({
+        "type": "user_status",
+        "user_id": user_id,
+        "username": username,
+        "display_name": display_name,
+        "avatar_path": avatar_path,
+        "status": "online"
+    }, exclude_user_id=user_id)
     
     try:
         while True:
-            data = await ws.receive_json()
-            mtype = data.get("type")
+            data = await websocket.receive_json()
+            message_type = data.get("type")
             
-            if mtype == "chat_message":
-                # Handle new message
-                content = data.get("content", "")
-                reply_to = data.get("reply_to_id")
-                client_id = data.get("client_id")
+            # --------------------------------------------------------------------
+            # Send Message
+            # --------------------------------------------------------------------
+            if message_type == "send_message":
+                content = data.get("content", "").strip()
+                reply_to_id = data.get("reply_to_id")
                 file_path = data.get("file_path")
                 file_type = data.get("file_type")
                 file_name = data.get("file_name")
-                ts = int(time.time() * 1000)
+                file_size = data.get("file_size", 0)
+                client_id = data.get("client_id")
+                
+                # Validate content
+                if not content and not file_path:
+                    await websocket.send_json({
+                        "type": "error",
+                        "code": "INVALID_MESSAGE",
+                        "message": "Message must contain text or a file"
+                    })
+                    continue
                 
+                # Validate content length
+                if len(content) > 10000:
+                    await websocket.send_json({
+                        "type": "error",
+                        "code": "MESSAGE_TOO_LONG",
+                        "message": "Message must be less than 10,000 characters"
+                    })
+                    continue
+                
+                # Validate reply exists
+                if reply_to_id:
+                    db = await get_db()
+                    try:
+                        async with db.execute(
+                            "SELECT id FROM messages WHERE id = ? AND is_deleted = 0",
+                            (reply_to_id,)
+                        ) as cursor:
+                            if not await cursor.fetchone():
+                                reply_to_id = None
+                    finally:
+                        await db.close()
+                
+                timestamp_ms = int(time.time() * 1000)
+                
+                # Encrypt content
                 encrypted = encrypt_message(content)
                 
+                # Store in database
                 db = await get_db()
                 try:
-                    cur = await db.execute(
-                        """INSERT INTO messages (sender_id, encrypted_content, timestamp_ms, reply_to_id, file_path, file_type, file_name)
-                           VALUES (?, ?, ?, ?, ?, ?, ?)""",
-                        (uid, encrypted, ts, reply_to, file_path, file_type, file_name)
+                    cursor = await db.execute(
+                        """INSERT INTO messages 
+                           (sender_id, encrypted_content, timestamp_ms, reply_to_id, 
+                            file_path, file_type, file_name, file_size)
+                           VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
+                        (user_id, encrypted, timestamp_ms, reply_to_id,
+                         file_path, file_type, file_name, file_size)
                     )
-                    mid = cur.lastrowid
+                    message_id = cursor.lastrowid
                     await db.commit()
                 finally:
                     await db.close()
                 
-                msg_out = {
+                # Build message object
+                message_obj = {
                     "type": "new_message",
-                    "id": mid,
-                    "sender_id": uid,
-                    "username": username,
-                    "display_name": user["display_name"],
-                    "avatar_path": user["avatar_path"],
-                    "content": content,
-                    "timestamp_ms": ts,
-                    "reply_to_id": reply_to,
-                    "file_path": file_path,
-                    "file_type": file_type,
-                    "file_name": file_name,
-                    "is_edited": False,
-                    "is_deleted": False,
-                    "client_id": client_id,
-                    "status": "sent"
+                    "message": {
+                        "id": message_id,
+                        "sender_id": user_id,
+                        "sender_username": username,
+                        "sender_display_name": display_name,
+                        "sender_avatar_path": avatar_path,
+                        "content": content,
+                        "timestamp_ms": timestamp_ms,
+                        "reply_to_id": reply_to_id,
+                        "file_path": file_path,
+                        "file_type": file_type,
+                        "file_name": file_name,
+                        "file_size": file_size,
+                        "is_edited": False,
+                        "is_deleted": False,
+                        "status": "sent"
+                    }
                 }
                 
-                # Send confirmation to sender
-                await ws.send_json({**msg_out, "status": "delivered"})
-                # Broadcast to others
-                await manager.broadcast(msg_out, exclude=ws)
+                # Send to sender with delivered status
+                await websocket.send_json({
+                    **message_obj,
+                    "message": {**message_obj["message"], "status": "delivered"},
+                    "client_id": client_id
+                })
+                
+                # Broadcast to other users
+                await manager.broadcast(message_obj, exclude_user_id=user_id)
             
-            elif mtype == "load_messages":
-                # Handle message loading with cursor-based pagination
-                before_id = data.get("before_id")
-                limit = min(data.get("limit", 50), 100)
+            # --------------------------------------------------------------------
+            # Load Messages (Cursor-based Pagination)
+            # --------------------------------------------------------------------
+            elif message_type == "load_messages":
+                cursor_id = data.get("cursor")  # Message ID to load before
+                limit = min(int(data.get("limit", 50)), 100)
                 
                 db = await get_db()
                 try:
-                    q = "SELECT * FROM messages WHERE is_deleted = 0"
-                    params = []
-                    if before_id is not None:
-                        q += " AND id < ?"
-                        params.append(before_id)
-                    q += " ORDER BY id DESC LIMIT ?"
-                    params.append(limit)
+                    if cursor_id:
+                        query = """
+                            SELECT m.*, u.username, u.display_name, u.avatar_path
+                            FROM messages m
+                            JOIN users u ON m.sender_id = u.id
+                            WHERE m.is_deleted = 0 AND m.id < ?
+                            ORDER BY m.id DESC
+                            LIMIT ?
+                        """
+                        params = (cursor_id, limit)
+                    else:
+                        query = """
+                            SELECT m.*, u.username, u.display_name, u.avatar_path
+                            FROM messages m
+                            JOIN users u ON m.sender_id = u.id
+                            WHERE m.is_deleted = 0
+                            ORDER BY m.id DESC
+                            LIMIT ?
+                        """
+                        params = (limit,)
                     
-                    rows = await (await db.execute(q, tuple(params))).fetchall()
-                    rows = list(reversed(rows))
+                    rows = await (await db.execute(query, params)).fetchall()
                     
-                    msgs = []
-                    for row in rows:
+                    messages = []
+                    for row in reversed(rows):  # Reverse to get chronological order
                         try:
-                            plain = decrypt_message(row["encrypted_content"])
+                            decrypted = decrypt_message(row['encrypted_content'])
                         except Exception:
-                            plain = "[decryption error]"
+                            decrypted = "[Decryption Error]"
                         
-                        s = await db.execute(
-                            "SELECT id, username, display_name, avatar_path FROM users WHERE id = ?",
-                            (row["sender_id"],)
-                        )
-                        s = await s.fetchone()
+                        # Check read status
+                        read_status = "sent"
+                        async with db.execute(
+                            "SELECT COUNT(*) as count FROM read_receipts WHERE message_id = ?",
+                            (row['id'],)
+                        ) as read_cursor:
+                            read_result = await read_cursor.fetchone()
+                            if read_result and read_result['count'] > 0:
+                                read_status = "read"
                         
-                        msgs.append({
-                            "id": row["id"],
-                            "sender_id": row["sender_id"],
-                            "username": s["username"] if s else "unknown",
-                            "display_name": s["display_name"] if s else "",
-                            "avatar_path": s["avatar_path"] if s else "",
-                            "content": plain,
-                            "timestamp_ms": row["timestamp_ms"],
-                            "reply_to_id": row["reply_to_id"],
-                            "file_path": row["file_path"],
-                            "file_type": row["file_type"],
-                            "file_name": row["file_name"],
-                            "is_edited": bool(row["is_edited"]),
-                            "is_deleted": bool(row["is_deleted"])
+                        messages.append({
+                            "id": row['id'],
+                            "sender_id": row['sender_id'],
+                            "sender_username": row['username'],
+                            "sender_display_name": row['display_name'],
+                            "sender_avatar_path": row['avatar_path'],
+                            "content": decrypted,
+                            "timestamp_ms": row['timestamp_ms'],
+                            "reply_to_id": row['reply_to_id'],
+                            "file_path": row['file_path'],
+                            "file_type": row['file_type'],
+                            "file_name": row['file_name'],
+                            "file_size": row['file_size'],
+                            "is_edited": bool(row['is_edited']),
+                            "is_deleted": False,
+                            "status": read_status
                         })
                     
-                    await ws.send_json({"type": "messages_batch", "messages": msgs, "before_id": before_id})
+                    # Get the oldest message ID for next page
+                    next_cursor = None
+                    if len(messages) > 0:
+                        next_cursor = messages[0]['id']
+                    
+                    await websocket.send_json({
+                        "type": "messages_loaded",
+                        "messages": messages,
+                        "next_cursor": next_cursor,
+                        "has_more": len(messages) == limit
+                    })
                 finally:
                     await db.close()
             
-            elif mtype == "edit_message":
-                # Handle message edit
-                mid, new_text = data.get("message_id"), data.get("content")
+            # --------------------------------------------------------------------
+            # Edit Message
+            # --------------------------------------------------------------------
+            elif message_type == "edit_message":
+                message_id = data.get("message_id")
+                new_content = data.get("content", "").strip()
+                
+                if not message_id or not new_content:
+                    await websocket.send_json({
+                        "type": "error",
+                        "code": "INVALID_EDIT",
+                        "message": "Message ID and new content are required"
+                    })
+                    continue
+                
+                if len(new_content) > 10000:
+                    await websocket.send_json({
+                        "type": "error",
+                        "code": "MESSAGE_TOO_LONG",
+                        "message": "Message must be less than 10,000 characters"
+                    })
+                    continue
+                
                 db = await get_db()
                 try:
-                    row = await db.execute(
+                    # Verify ownership
+                    async with db.execute(
                         "SELECT * FROM messages WHERE id = ? AND sender_id = ? AND is_deleted = 0",
-                        (mid, uid)
+                        (message_id, user_id)
+                    ) as cursor:
+                        if not await cursor.fetchone():
+                            await websocket.send_json({
+                                "type": "error",
+                                "code": "NOT_FOUND",
+                                "message": "Message not found or you don't have permission to edit it"
+                            })
+                            continue
+                    
+                    new_encrypted = encrypt_message(new_content)
+                    await db.execute(
+                        "UPDATE messages SET encrypted_content = ?, is_edited = 1 WHERE id = ?",
+                        (new_encrypted, message_id)
                     )
-                    if await row.fetchone():
-                        new_enc = encrypt_message(new_text)
-                        await db.execute(
-                            "UPDATE messages SET encrypted_content = ?, is_edited = 1 WHERE id = ?",
-                            (new_enc, mid)
-                        )
-                        await db.commit()
-                        await manager.broadcast({"type": "edit_message", "message_id": mid, "content": new_text})
+                    await db.commit()
+                    
+                    # Broadcast edit
+                    await manager.broadcast({
+                        "type": "message_edited",
+                        "message_id": message_id,
+                        "content": new_content,
+                        "editor_id": user_id
+                    })
                 finally:
                     await db.close()
             
-            elif mtype == "delete_message":
-                # Handle permanent message deletion
-                mid = data.get("message_id")
+            # --------------------------------------------------------------------
+            # Delete Message
+            # --------------------------------------------------------------------
+            elif message_type == "delete_message":
+                message_id = data.get("message_id")
+                
+                if not message_id:
+                    continue
+                
                 db = await get_db()
                 try:
-                    row = await db.execute("SELECT * FROM messages WHERE id = ? AND sender_id = ?", (mid, uid))
-                    msg = await row.fetchone()
-                    if msg:
-                        if msg["file_path"]:
-                            try:
-                                delete_file(msg["file_path"])
-                            except Exception:
-                                pass
-                        
-                        await db.execute("DELETE FROM messages WHERE id = ?", (mid,))
-                        await db.execute("DELETE FROM read_receipts WHERE message_id = ?", (mid,))
-                        await db.commit()
-                        await manager.broadcast({"type": "delete_message", "message_id": mid})
+                    # Verify ownership
+                    async with db.execute(
+                        "SELECT * FROM messages WHERE id = ? AND sender_id = ?",
+                        (message_id, user_id)
+                    ) as cursor:
+                        message = await cursor.fetchone()
+                    
+                    if not message:
+                        await websocket.send_json({
+                            "type": "error",
+                            "code": "NOT_FOUND",
+                            "message": "Message not found or you don't have permission to delete it"
+                        })
+                        continue
+                    
+                    # Delete associated file if exists
+                    if message['file_path']:
+                        try:
+                            delete_file(message['file_path'])
+                        except Exception:
+                            pass
+                    
+                    # Remove message and read receipts
+                    await db.execute("DELETE FROM messages WHERE id = ?", (message_id,))
+                    await db.execute("DELETE FROM read_receipts WHERE message_id = ?", (message_id,))
+                    await db.commit()
+                    
+                    # Broadcast deletion
+                    await manager.broadcast({
+                        "type": "message_deleted",
+                        "message_id": message_id
+                    })
                 finally:
                     await db.close()
             
-            elif mtype == "typing":
-                # Handle typing indicator
-                await manager.broadcast(
-                    {"type": "typing_update", "username": username, "is_typing": data.get("is_typing", False)},
-                    exclude=ws
-                )
+            # --------------------------------------------------------------------
+            # Typing Indicator
+            # --------------------------------------------------------------------
+            elif message_type == "typing":
+                is_typing = bool(data.get("is_typing", False))
+                
+                await manager.broadcast({
+                    "type": "typing_indicator",
+                    "user_id": user_id,
+                    "username": username,
+                    "display_name": display_name,
+                    "is_typing": is_typing
+                }, exclude_user_id=user_id)
             
-            elif mtype == "mark_read":
-                # Handle read receipts
-                up_to = data.get("up_to_message_id")
-                if up_to:
-                    db = await get_db()
-                    try:
-                        rows = await db.execute(
-                            """SELECT id, sender_id FROM messages
-                               WHERE id <= ? AND sender_id != ? AND is_deleted = 0
-                               EXCEPT SELECT message_id FROM read_receipts WHERE user_id = ?""",
-                            (up_to, uid, uid)
+            # --------------------------------------------------------------------
+            # Mark as Read
+            # --------------------------------------------------------------------
+            elif message_type == "mark_read":
+                up_to_message_id = data.get("up_to_message_id")
+                
+                if not up_to_message_id:
+                    continue
+                
+                db = await get_db()
+                try:
+                    # Get messages that haven't been read by this user
+                    rows = await db.execute(
+                        """SELECT m.id, m.sender_id 
+                           FROM messages m 
+                           WHERE m.id <= ? 
+                           AND m.sender_id != ? 
+                           AND m.is_deleted = 0 
+                           AND m.id > (
+                               SELECT COALESCE(MAX(message_id), 0) 
+                               FROM read_receipts 
+                               WHERE user_id = ?
+                           )""",
+                        (up_to_message_id, user_id, user_id)
+                    )
+                    
+                    new_reads = []
+                    async for row in rows:
+                        await db.execute(
+                            "INSERT OR IGNORE INTO read_receipts (user_id, message_id, read_at) VALUES (?, ?, ?)",
+                            (user_id, row['id'], int(time.time()))
                         )
-                        
-                        receipts = []
-                        async for r in rows:
-                            await db.execute(
-                                "INSERT OR IGNORE INTO read_receipts (user_id, message_id) VALUES (?, ?)",
-                                (uid, r["id"])
-                            )
-                            receipts.append(r["id"])
-                        
-                        await db.commit()
-                        
-                        for mid in receipts:
-                            await manager.broadcast({
-                                "type": "read_receipt",
-                                "message_id": mid,
-                                "reader_username": username
-                            })
-                    finally:
-                        await db.close()
+                        new_reads.append({
+                            "message_id": row['id'],
+                            "sender_id": row['sender_id']
+                        })
+                    
+                    await db.commit()
+                    
+                    # Notify message senders
+                    for read_info in new_reads:
+                        await manager.send_to_user(read_info['sender_id'], {
+                            "type": "message_read",
+                            "message_id": read_info['message_id'],
+                            "reader_id": user_id,
+                            "reader_username": username
+                        })
+                finally:
+                    await db.close()
+            
+            # --------------------------------------------------------------------
+            # Get Online Users
+            # --------------------------------------------------------------------
+            elif message_type == "get_online_users":
+                await websocket.send_json({
+                    "type": "online_users",
+                    "users": manager.get_online_users()
+                })
     
     except WebSocketDisconnect:
         pass
     except Exception as e:
         logger.error(f"WebSocket error for {username}: {e}")
     finally:
-        manager.disconnect(ws, username)
-        if username not in manager.active:
-            await manager.broadcast({"type": "user_offline", "username": username})
+        # Disconnect
+        manager.disconnect(websocket, user_id)
+        
+        # Update status in database
+        db = await get_db()
+        try:
+            await db.execute(
+                "UPDATE users SET status = 'offline', last_seen = ? WHERE id = ?",
+                (int(time.time()), user_id)
+            )
+            await db.commit()
+        finally:
+            await db.close()
+        
+        # Broadcast offline status
+        await manager.broadcast({
+            "type": "user_status",
+            "user_id": user_id,
+            "username": username,
+            "display_name": display_name,
+            "avatar_path": avatar_path,
+            "status": "offline"
+        })
+
+# ------------------------------------------------------------------------
+# Root endpoint
+# ------------------------------------------------------------------------
+@app.get("/")
+async def root():
+    """Serve the main application."""
+    return FileResponse("static/index.html")