main.py

"""
FastAPI Backend for Multi-PDF RAG System with Per-User Storage
===============================================================
Secure multi-user API with:
- API key authentication
- Per-user storage isolation
- PDF upload and management
- RAG-based question answering
- HF persistent storage
"""

import os
import time
import asyncio
import hashlib
from typing import List, Optional, Dict
from fastapi import FastAPI, UploadFile, File, HTTPException, Depends, BackgroundTasks, Header, Request
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from threading import Lock
from rag_engine import RAGEngine
from hf_storage import create_hf_storage_manager
from auth import get_current_user
from user_management import create_hf_user_manager
from job_worker import job_manager
from dotenv import load_dotenv

# ============================================
# CONFIGURATION
# ============================================

load_dotenv()

GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
HF_TOKEN = os.environ.get("HF_TOKEN")
HF_REPO = os.environ.get("HF_REPO", "Hamza4100/multi-pdf-storage")

if not GEMINI_API_KEY:
    raise RuntimeError("❌ GEMINI_API_KEY not set")

hf_storage = create_hf_storage_manager(hf_token=HF_TOKEN, hf_repo=HF_REPO)
hf_user_manager = create_hf_user_manager(hf_token=HF_TOKEN, hf_repo=HF_REPO)

app = FastAPI(
    title="Multi-PDF RAG System",
    description="Secure multi-user RAG API with persistent storage",
    version="2.0.0"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ============================================
# PER-USER RAG ENGINE MANAGER
# ============================================

class UserRAGManager:
    """Manages per-user RAG engine instances with lazy loading."""
    
    def __init__(self):
        self.engines: Dict[str, RAGEngine] = {}
        self.locks: Dict[str, Lock] = {}
        self.global_lock = Lock()
    
    def get_user_lock(self, user_id: str) -> Lock:
        """Get or create lock for user."""
        with self.global_lock:
            if user_id not in self.locks:
                self.locks[user_id] = Lock()
            return self.locks[user_id]
    
    async def get_engine(self, user_id: str) -> RAGEngine:
        """Get or create RAG engine for user (lazy loading)."""
        if user_id in self.engines:
            return self.engines[user_id]
        
        user_lock = self.get_user_lock(user_id)
        
        with user_lock:
            if user_id in self.engines:
                return self.engines[user_id]
            
            print(f"🔧 Initializing RAG for user {user_id}...")
            
            # Sync from HF
            await asyncio.to_thread(hf_storage.sync_storage_from_hf, user_id)
            
            # User-specific paths
            base_dir = os.path.dirname(os.path.abspath(__file__))
            user_storage_dir = os.path.join(base_dir, "users", user_id, "storage")
            
            # Initialize engine
            engine = await asyncio.to_thread(
                RAGEngine,
                gemini_api_key=GEMINI_API_KEY,
                storage_dir=user_storage_dir
            )
            
            self.engines[user_id] = engine
            print(f"✅ RAG ready for user {user_id}")
            
            return engine


rag_manager = UserRAGManager()

# ============================================
# MODELS
# ============================================

class UploadResponse(BaseModel):
    document_id: str
    filename: str
    status: str
    message: str
    pages: Optional[int] = None
    chunks: Optional[int] = None


class QueryRequest(BaseModel):
    question: str
    top_k: Optional[int] = 5
    doc_id: Optional[str] = None


class QueryResponse(BaseModel):
    answer: str
    sources: List[dict]


class DocumentInfo(BaseModel):
    doc_id: str
    filename: str
    upload_timestamp: str
    num_chunks: int
    num_pages: int


class StatsResponse(BaseModel):
    total_documents: int
    total_chunks: int
    index_size: int


class DeleteResponse(BaseModel):
    status: str
    message: str


class RegisterRequest(BaseModel):
    username: str
    password: str
    email: str


class LoginRequest(BaseModel):
    username: str
    password: str


class AuthResponse(BaseModel):
    success: bool
    message: str
    api_key: Optional[str] = None
    username: Optional[str] = None


# ============================================
# HELPER FUNCTIONS
# ============================================

def hash_password(password: str) -> str:
    """Hash password using SHA256."""
    return hashlib.sha256(password.encode()).hexdigest()


def get_available_api_key() -> Optional[str]:
    """
    Get an available API key from the configured pool.
    Returns None if no keys available.
    """
    if not hf_user_manager.enabled:
        return None
    
    # Get API_KEYS from environment
    api_keys_str = os.environ.get("API_KEYS", "")
    available_keys = set(k.strip() for k in api_keys_str.split(",") if k.strip())
    
    if not available_keys:
        return None
    
    # Get already assigned keys
    all_users = hf_user_manager.get_all_users()
    assigned_keys = {user_data.get("api_key") for user_data in all_users.values() if user_data.get("api_key")}
    
    # Find first available key
    for key in available_keys:
        if key not in assigned_keys:
            return key
    
    return None


# ============================================
# STARTUP
# ============================================

@app.on_event("startup")
async def startup_event():
    print("🚀 Multi-PDF RAG System v2.0")
    print(f"📦 HF Storage: {'Enabled' if hf_storage.enabled else 'Disabled'}")
    
    # Load and display user count from HF
    if hf_user_manager.enabled:
        users_count = len(hf_user_manager.get_all_users())
        print(f"📥 Loaded {users_count} user(s) from HF repository")
        print(f"✅ User management: Enabled (HF-based)")
    else:
        print(f"⚠️ User management: Disabled (check HF_TOKEN and HF_REPO)")
    
    print("✅ Server ready (per-user lazy loading)")

    # Start persistent job manager
    try:
        job_manager.start()
        print("🔁 Background job manager started")
    except Exception as e:
        print(f"⚠️ Failed to start job manager: {e}")


# ============================================
# ENDPOINTS
# ============================================

@app.get("/health")
async def health_check():
    """Health check (no auth required)."""
    return {"status": "ok"}


# ============================================
# AUTHENTICATION ENDPOINTS
# ============================================

@app.post("/auth/register", response_model=AuthResponse)
async def register_user(request: RegisterRequest):
    """
    Register a new user account.
    No authentication required.
    """
    if not hf_user_manager.enabled:
        raise HTTPException(503, "User management is not enabled on this server")
    
    # Validate input
    if len(request.username) < 3:
        return AuthResponse(success=False, message="Username must be at least 3 characters")
    
    if len(request.password) < 6:
        return AuthResponse(success=False, message="Password must be at least 6 characters")
    
    # Check if user already exists
    if hf_user_manager.user_exists(request.username):
        return AuthResponse(success=False, message="Username already exists")
    
    # Get available API key
    api_key = get_available_api_key()
    if not api_key:
        return AuthResponse(success=False, message="No API keys available. Please try again later.")
    
    # Create user
    success = hf_user_manager.create_user(
        username=request.username,
        password_hash=hash_password(request.password),
        email=request.email,
        api_key=api_key
    )
    
    if success:
        print(f"✅ New user registered: {request.username}")
        return AuthResponse(
            success=True,
            message="Registration successful",
            api_key=api_key,
            username=request.username
        )
    else:
        return AuthResponse(success=False, message="Failed to create user account")


@app.post("/auth/login", response_model=AuthResponse)
async def login_user(request: LoginRequest):
    """
    Login with username and password.
    No authentication required.
    Returns API key on success.
    """
    if not hf_user_manager.enabled:
        raise HTTPException(503, "User management is not enabled on this server")
    
    # Get user
    user_data = hf_user_manager.get_user(request.username)
    if not user_data:
        return AuthResponse(success=False, message="Invalid username or password")
    
    # Verify password
    if hf_user_manager.verify_password(request.username, hash_password(request.password)):
        api_key = user_data.get("api_key")
        print(f"✅ User logged in: {request.username}")
        return AuthResponse(
            success=True,
            message="Login successful",
            api_key=api_key,
            username=request.username
        )
    else:
        return AuthResponse(success=False, message="Invalid username or password")


@app.get("/auth/check")
async def check_auth(user_id: str = Depends(get_current_user)):
    """
    Check if API key is valid.
    Requires: X-API-KEY header
    """
    return {"authenticated": True, "user_id": user_id}


# ============================================
# DOCUMENT ENDPOINTS
# ============================================

@app.post("/upload", response_model=UploadResponse)
async def upload_pdf(
    request: Request,
    file: UploadFile = File(...),
    user_id: str = Depends(get_current_user),
    x_api_key: Optional[str] = Header(None, alias="X-API-KEY"),
    background_tasks: BackgroundTasks = None
):
    """
    Upload PDF for authenticated user.
    Requires: X-API-KEY header
    """
    # Validate PDF
    if not file.filename.lower().endswith('.pdf'):
        raise HTTPException(400, "Only PDF files allowed")
    
    if file.content_type and file.content_type not in ['application/pdf']:
        raise HTTPException(400, "Invalid MIME type")
    
    # Read content quickly and persist to a temp file, then enqueue a job for background processing
    content = await file.read()

    # Size limit (10MB)
    if len(content) > 10 * 1024 * 1024:
        raise HTTPException(413, "File too large (max 10MB)")

    try:
        # Debug: log a subset of incoming headers to help diagnose authentication issues
        try:
            hdrs = dict(request.headers)
            # Only print a few header keys to avoid leaking secrets in logs
            interesting = {k: hdrs.get(k) for k in ["x-api-key", "authorization", "host", "user-agent"] if hdrs.get(k)}
            print(f"🔍 Incoming upload headers (filtered): {interesting}")
        except Exception:
            pass
        # Save to a temp file so background worker can access it
        base_dir = os.path.dirname(os.path.abspath(__file__))
        tmp_dir = os.path.join(base_dir, 'tmp_uploads')
        os.makedirs(tmp_dir, exist_ok=True)
        tmp_path = os.path.join(tmp_dir, f"{int(time.time())}_{file.filename}")
        with open(tmp_path, 'wb') as f:
            f.write(content)

        # Try to resolve username from HF user manager (if available)
        username = None
        try:
            if hf_user_manager and hf_user_manager.enabled and x_api_key:
                found = hf_user_manager.get_user_by_api_key(x_api_key)
                if found:
                    username = found[0]
        except Exception:
            username = None

        # Create persistent job and return immediately
        job = job_manager.create_job(user_id=user_id, filename=file.filename, file_path=tmp_path, username=username)

        return UploadResponse(
            document_id=job['id'],
            filename=file.filename,
            status="processing",
            message="Upload accepted and is being processed in background"
        )

    except HTTPException:
        raise
    except Exception as e:
        print(f"❌ Upload enqueue error (user {user_id}): {e}")
        raise HTTPException(500, "Failed to enqueue upload")


@app.get('/upload-status/{job_id}')
async def upload_status(job_id: str):
    job = job_manager.get_job(job_id)
    if not job:
        raise HTTPException(404, "Job not found")
    return job


@app.post("/query", response_model=QueryResponse)
async def query_documents(
    request: QueryRequest,
    user_id: str = Depends(get_current_user)
):
    """
    Query user's documents using RAG.
    Requires: X-API-KEY header
    """
    try:
        engine = await rag_manager.get_engine(user_id)
        
        result = await asyncio.to_thread(
            engine.ask,
            query=request.question,
            top_k=request.top_k,
            doc_id=request.doc_id
        )
        
        print(f"✅ Query success for user {user_id}")
        
        return QueryResponse(
            answer=result["answer"],
            sources=result.get("sources", [])
        )
    
    except Exception as e:
        print(f"❌ Query error (user {user_id}): {e}")
        raise HTTPException(500, "Query failed")


@app.get("/documents", response_model=List[DocumentInfo])
async def get_documents(user_id: str = Depends(get_current_user)):
    """
    Get all documents for authenticated user.
    Requires: X-API-KEY header
    """
    try:
        engine = await rag_manager.get_engine(user_id)
        documents = await asyncio.to_thread(engine.get_all_documents)
        
        return [
            DocumentInfo(
                doc_id=doc["doc_id"],
                filename=doc["filename"],
                upload_timestamp=doc["upload_timestamp"],
                num_chunks=doc["num_chunks"],
                num_pages=doc["num_pages"]
            )
            for doc in documents
        ]
    
    except Exception as e:
        print(f"❌ Get documents error (user {user_id}): {e}")
        raise HTTPException(500, "Failed to retrieve documents")


@app.delete("/documents/{doc_id}", response_model=DeleteResponse)
async def delete_document(
    doc_id: str,
    user_id: str = Depends(get_current_user)
):
    """
    Delete document for authenticated user.
    Requires: X-API-KEY header
    """
    try:
        engine = await rag_manager.get_engine(user_id)
        user_lock = rag_manager.get_user_lock(user_id)
        
        with user_lock:
            result = await asyncio.to_thread(engine.delete_document, doc_id)
            
            if result["status"] == "success":
                await asyncio.to_thread(
                    hf_storage.push_storage_to_hf,
                    user_id,
                    f"Delete {doc_id}"
                )
                
                print(f"✅ Delete success for user {user_id}: {doc_id}")
                
                return DeleteResponse(
                    status="success",
                    message=result["message"]
                )
            else:
                raise HTTPException(404, result["message"])
    
    except HTTPException:
        raise
    except Exception as e:
        print(f"❌ Delete error (user {user_id}): {e}")
        raise HTTPException(500, "Deletion failed")


@app.get("/stats", response_model=StatsResponse)
async def get_stats(user_id: str = Depends(get_current_user)):
    """
    Get stats for authenticated user.
    Requires: X-API-KEY header
    """
    try:
        engine = await rag_manager.get_engine(user_id)
        stats = await asyncio.to_thread(engine.get_stats)
        
        return StatsResponse(
            total_documents=stats["total_documents"],
            total_chunks=stats["total_chunks"],
            index_size=stats["index_size"]
        )
    
    except Exception as e:
        print(f"❌ Stats error (user {user_id}): {e}")
        raise HTTPException(500, "Failed to retrieve stats")


# ============================================
# MAIN
# ============================================

if __name__ == "__main__":
    import uvicorn
    port = int(os.environ.get("PORT", 8000))
    uvicorn.run("main:app", host="0.0.0.0", port=port, reload=True)