app/main.py

# ============================================================================
# backend/app/main.py
# ============================================================================


"""
FastAPI Main Application Entry Point (UPDATED)

Banking RAG Chatbot API with JWT Authentication

CHANGES:
- Replaced old chat router with new conversation_routes
- Added conversation management features
"""

from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from contextlib import asynccontextmanager

from app.config import settings
from app.db.mongodb import connect_to_mongo, close_mongo_connection

# ============================================================================
# LIFESPAN MANAGER (Startup & Shutdown)
# ============================================================================

from contextlib import asynccontextmanager
import asyncio

@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    Manage application lifespan events.

    Startup:
    - Connect to MongoDB Atlas
    - Create indexes for conversations after DB ready (retry logic)
    - ML models lazy loaded
    """
    print("\n" + "=" * 80)
    print("🚀 STARTING BANKING RAG CHATBOT API")
    print("=" * 80)
    print(f"Environment: {settings.ENVIRONMENT}")
    print(f"Debug Mode: {settings.DEBUG}")
    print("=" * 80)

    # STEP 1 — Connect to MongoDB
    print("\n🔌 Connecting to MongoDB Atlas...")
    await connect_to_mongo()

    # STEP 2 — Retry index creation until DB is ready
    from app.db.repositories.conversation_repository import conversation_repository

    for attempt in range(5):
        try:
            await conversation_repository.create_indexes()
            print("✅ MongoDB indexes created successfully!")
            break
        except Exception as e:
            print(f"⚠️ Index creation failed (Attempt {attempt + 1}/5): {e}")
            await asyncio.sleep(2)
    else:
        print("❌ Could not create MongoDB indexes after multiple retries")

    print("\n💡 ML Models Info:")
    print("   Policy Network: Loads on first chat request (lazy load)")
    print("   Retriever Model: Loads on first retrieval (lazy load)")
    print("   LLM: Groq (ChatGroq) with HuggingFace fallback")
    print("\n🤖 LLM Configuration:")
    print(f"   Chat Model: {settings.GROQ_CHAT_MODEL} (Llama 3 8B)")
    print(f"   Eval Model: {settings.GROQ_EVAL_MODEL} (Llama 3 70B)")
    print(f"   Groq API Keys: {len(settings.get_groq_api_keys())} configured")
    print(f"   HuggingFace Tokens: {len(settings.get_hf_tokens())} configured")
    print(f"   Fallback: Groq → HuggingFace")

    print("\n✅ Backend startup complete!")
    print("=" * 80)
    print(f"📖 API Docs: https://eeshanyaj-questrag-backend.hf.space/docs")
    print(f"🏥 Health Check: https://eeshanyaj-questrag-backend.hf.space/health")
    print(f"🧠 Backend Link: https://eeshanyaj-questrag-backend.hf.space/")
    print("=" * 80 + "\n")

    yield

    # SHUTDOWN
    print("\n" + "=" * 80)
    print("🛑 SHUTTING DOWN API")
    print("=" * 80)
    await close_mongo_connection()
    print("✅ Shutdown complete")
    print("=" * 80 + "\n")


# ============================================================================
# CREATE FASTAPI APPLICATION
# ============================================================================

app = FastAPI(
    title="Banking RAG Chatbot API",
    description="""
🤖 AI-powered Banking Assistant with:

**Features:**
- 🔐 JWT Authentication (Sign up, Login, Protected routes)
- 💬 RAG (Retrieval-Augmented Generation)
- 🧠 RL-based Policy Network (BERT)
- 🔍 Custom E5 Retriever
- ⚡ Groq LLM with HuggingFace Fallback (Llama 3 models)
- 📝 Conversation Management (List, Search, Archive, Delete)

**Capabilities:**
- Intelligent document retrieval
- Context-aware responses
- Conversation persistence & history
- Auto-generated conversation titles
- Real-time chat with RAG pipeline
- User authentication & authorization
- Multi-provider LLM with automatic fallback
    """,
    version="2.0.0",
    docs_url="/docs",
    redoc_url="/redoc",
    lifespan=lifespan
)

# ============================================================================
# CORS MIDDLEWARE
# ============================================================================

allowed_origins = settings.get_allowed_origins()
print("\n🌐 CORS Configuration:")
print(f"   Allowed Origins: {allowed_origins}")

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

# ============================================================================
# INCLUDE API ROUTERS (UPDATED)
# ============================================================================

from app.api.v1 import auth
from app.api.v1 import conversation_routes  # ✅ NEW IMPORT
from app.api.v1 import file_routes # new file routes

# Auth router (public endpoints - register, login)
app.include_router(
    auth.router,
    prefix="/api/v1/auth",
    tags=["🔐 Authentication"]
)

# Conversation & Chat router (protected endpoints - requires JWT token)
app.include_router(
    conversation_routes.router,  # ✅ NEW ROUTER
    prefix="/api/v1/chat",
    tags=["💬 Chat & Conversations"]
)
# File Upload router (protected endpoints - requires JWT token)
app.include_router(
    file_routes.router,
    prefix="/api/v1",
    tags=["📁 File Upload"]
)

# ============================================================================
# ROOT ENDPOINTS
# ============================================================================

@app.get("/", tags=["📍 Root"])
async def root():
    """
    Root endpoint - API information and available endpoints
    """
    return {
        "message": "Banking RAG Chatbot API with Authentication & Conversation Management",
        "version": "2.0.0",
        "status": "online",
        "authentication": "JWT Bearer Token Required for chat endpoints",
        "llm_provider": "Groq (ChatGroq) with HuggingFace fallback",
        "models": {
            "chat": settings.GROQ_CHAT_MODEL,
            "evaluation": settings.GROQ_EVAL_MODEL
        },
        "documentation": {
            "swagger_ui": "/docs",
            "redoc": "/redoc"
        },
        "endpoints": {
            "auth": {
                "register": "POST /api/v1/auth/register",
                "login": "POST /api/v1/auth/login",
                "me": "GET /api/v1/auth/me (requires token)",
                "logout": "POST /api/v1/auth/logout (requires token)"
            },
            "chat": {
                "send_message": "POST /api/v1/chat/ (requires token)",
                "create_conversation": "POST /api/v1/chat/conversation (requires token)",
                "list_conversations": "GET /api/v1/chat/conversations (requires token)",
                "get_conversation": "GET /api/v1/chat/conversation/{id} (requires token)",
                "update_conversation": "PATCH /api/v1/chat/conversation/{id} (requires token)",
                "delete_conversation": "DELETE /api/v1/chat/conversation/{id} (requires token)",
                "search_conversations": "GET /api/v1/chat/conversations/search (requires token)",
                "conversation_stats": "GET /api/v1/chat/conversations/stats (requires token)"
            },
            "files": {
                "upload_image": "POST /api/v1/files/upload/image (requires token)",
                "upload_pdf": "POST /api/v1/files/upload/pdf (requires token)",
                "upload_document": "POST /api/v1/files/upload/document (requires token)",
                "upload_audio": "POST /api/v1/files/upload/audio (requires token)",
                "delete_file": "DELETE /api/v1/files/delete (requires token)",
                "health": "GET /api/v1/files/health"
            },
            "health": "GET /health"
        }
    }

@app.get("/health", tags=["🏥 Health"])
async def health_check():
    """
    Comprehensive health check endpoint
    
    Checks status of:
    - API service
    - MongoDB connection
    - ML models (lazy loaded)
    - Authentication system
    - LLM providers (Groq & HuggingFace)
    
    Returns:
        dict: Health status of all components
    """
    from app.db.mongodb import get_database
    
    # Check MongoDB
    mongodb_status = "connected" if get_database() is not None else "disconnected"
    
    # Check ML models (don't load them, just check readiness)
    ml_models_status = {
        "policy_network": "ready (lazy load)",
        "retriever": "ready (lazy load)",
        "llm": "ready (API-based)"
    }
    
    # Check LLM providers
    llm_providers = {
        "groq": {
            "enabled": settings.is_groq_enabled(),
            "api_keys_configured": len(settings.get_groq_api_keys()),
            "chat_model": settings.GROQ_CHAT_MODEL,
            "eval_model": settings.GROQ_EVAL_MODEL
        },
        "huggingface": {
            "enabled": settings.is_hf_enabled(),
            "tokens_configured": len(settings.get_hf_tokens()),
            "chat_model": settings.HF_CHAT_MODEL,
            "eval_model": settings.HF_EVAL_MODEL
        }
    }
    
    # Check authentication
    auth_status = {
        "jwt_enabled": bool(settings.SECRET_KEY and settings.SECRET_KEY != "your-secret-key-change-in-production"),
        "algorithm": settings.ALGORITHM,
        "token_expiry_minutes": settings.ACCESS_TOKEN_EXPIRE_MINUTES
    }
    
    # Overall health
    is_healthy = (
        mongodb_status == "connected" and 
        auth_status["jwt_enabled"] and
        (llm_providers["groq"]["enabled"] or llm_providers["huggingface"]["enabled"])
    )
    
    return {
        "status": "healthy" if is_healthy else "degraded",
        "api": "online",
        "version": "2.0.0",
        "mongodb": mongodb_status,
        "authentication": auth_status,
        "llm_providers": llm_providers,
        "ml_models": ml_models_status,
        "environment": settings.ENVIRONMENT,
        "debug_mode": settings.DEBUG
    }

# ============================================================================
# GLOBAL EXCEPTION HANDLER
# ============================================================================

@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    """
    Global exception handler for unhandled errors
    """
    print(f"\n❌ Unhandled Exception:")
    print(f"   Path: {request.url.path}")
    print(f"   Error: {str(exc)}")
    
    if settings.DEBUG:
        import traceback
        traceback.print_exc()
    
    return JSONResponse(
        status_code=500,
        content={
            "error": "Internal Server Error",
            "detail": str(exc) if settings.DEBUG else "An unexpected error occurred",
            "path": str(request.url.path)
        }
    )

# ============================================================================
# MAIN ENTRY POINT (for direct execution)
# ============================================================================

if __name__ == "__main__":
    import uvicorn
    
    print("\n🚀 Starting server directly...")
    print("   Note: For production, use: uvicorn app.main:app --host 0.0.0.0 --port 8000")
    
    uvicorn.run(
        "app.main:app",
        host="0.0.0.0",
        port=8000,
        reload=settings.DEBUG  # Auto-reload only in debug mode
    )


# ============================================================================================================================
# OLD CODE - KEEP FOR REFERENCE
# ============================================================================================================================







# ============================================================================
# LIFESPAN MANAGER (Startup & Shutdown)
# ============================================================================

# @asynccontextmanager
# async def lifespan(app: FastAPI):
#     """
#     Manage application lifespan events.
    
#     Startup:
#     - Connect to MongoDB Atlas
#     - Create indexes for conversations
#     - ML models load lazily on first use
    
#     Shutdown:
#     - Close MongoDB connection
#     - Cleanup resources
#     """
#     # ========================================================================
#     # STARTUP
#     # ========================================================================
#     print("\n" + "=" * 80)
#     print("🚀 STARTING BANKING RAG CHATBOT API")
#     print("=" * 80)
#     print(f"Environment: {settings.ENVIRONMENT}")
#     print(f"Debug Mode: {settings.DEBUG}")
#     print("=" * 80)
    
#     # Connect to MongoDB
#     # await connect_to_mongo()
    
#     # # Create indexes for conversations (async)
#     # try:
#     #     from app.db.repositories.conversation_repository import conversation_repository
#     #     await conversation_repository.create_indexes()
#     # except Exception as e:
#     #     print(f"⚠️ Failed to create conversation indexes: {e}")
#     # Connect to MongoDB first
#     await connect_to_mongo()

#     # Create indexes after connection success
#     from app.db.repositories.conversation_repository import conversation_repository

#     for attempt in range(5):
#         try:
#             await conversation_repository.create_indexes()
#             print("✅ Conversation indexes created successfully")
#             break
#         except Exception as e:
#             print(f"⚠️ Index creation failed (Attempt {attempt + 1}/5): {e}")
#             await asyncio.sleep(2)
#     else:
#         print("❌ Failed to create indexes after retries")

    
#     print("\n💡 ML Models Info:")
#     print("   Policy Network: Loads on first chat request (lazy loading)")
#     print("   Retriever Model: Loads on first retrieval (lazy loading)")
#     print("   LLM: Groq (ChatGroq) with HuggingFace fallback")
#     print("\n🤖 LLM Configuration:")
#     print(f"   Chat Model: {settings.GROQ_CHAT_MODEL} (Llama 3 8B)")
#     print(f"   Eval Model: {settings.GROQ_EVAL_MODEL} (Llama 3 70B)")
#     print(f"   Groq API Keys: {len(settings.get_groq_api_keys())} configured")
#     print(f"   HuggingFace Tokens: {len(settings.get_hf_tokens())} configured")
#     print(f"   Fallback: Groq → HuggingFace")
    
#     print("\n✅ Backend startup complete!")
#     print("=" * 80)
#     print(f"📖 API Docs: https://eeshanyaj-questrag-backend.hf.space/docs")
#     print(f"🏥 Health Check: https://eeshanyaj-questrag-backend.hf.space/health")
#     print(f"🧠 Backend Link: https://eeshanyaj-questrag-backend.hf.space/")
#     print("=" * 80 + "\n")
    
#     yield  # Application runs here
    
#     # ========================================================================
#     # SHUTDOWN
#     # ========================================================================
#     print("\n" + "=" * 80)
#     print("🛑 SHUTTING DOWN API")
#     print("=" * 80)
    
#     # Close MongoDB connection
#     await close_mongo_connection()
    
#     print("✅ Shutdown complete")
#     print("=" * 80 + "\n")