backend/main.py

"""
AegisLM Backend Application

FastAPI application entry point for the multi-agent adversarial LLM evaluation framework.
"""

import asyncio
import logging
from contextlib import asynccontextmanager
from typing import AsyncGenerator

from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles

from backend.api import routes
from backend.api.public import router as public_router
from backend.api.regulatory_routes import router as regulatory_router
from backend.config import settings
from backend.core.exceptions import AegisLMException
from backend.db.session import close_database, init_database
from backend.logging.logger import get_logger


# =============================================================================
# Application Setup
# =============================================================================

# Configure logging
logging.basicConfig(
    level=getattr(logging, settings.log_level.upper()),
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)

# Get logger
logger = get_logger("main", component="api")


@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
    """
    Application lifespan manager.
    
    Handles startup and shutdown events with FAIL-FAST validation.
    In production mode, critical dependencies MUST be available before serving traffic.
    """
    import os
    
    # Check if running in production mode
    is_production = os.getenv("AEGISLM_ENV", "development").lower() == "production"
    
    # Startup
    logger.info("Starting AegisLM backend", metadata={
        "version": "0.1.0",
        "mode": "production" if is_production else "development"
    })
    
    # =====================================================================
    # FAIL-FAST STARTUP VALIDATION (Production Mode)
    # =====================================================================
    
    # 1. Validate Database Connection SYNCHRONOUSLY at startup
    try:
        from backend.db.session import check_database_connection
        db_connected = await check_database_connection()
        
        if not db_connected:
            error_msg = "CRITICAL: Database connection failed at startup. Cannot start without database."
            logger.error(error_msg)
            if is_production:
                raise RuntimeError(error_msg)
    except Exception as e:
        error_msg = f"CRITICAL: Database initialization error: {str(e)}"
        logger.error(error_msg)
        if is_production:
            raise RuntimeError(error_msg)
    
    # 2. Validate Redis Connection SYNCHRONOUSLY at startup (Production only)
    if is_production:
        try:
            import redis.asyncio as redis
            from backend.core.config import settings
            
            if not settings.redis_url:
                error_msg = "CRITICAL: REDIS_URL not configured in production mode. Redis is required for rate limiting."
                logger.error(error_msg)
                raise RuntimeError(error_msg)
            
            redis_client = redis.from_url(
                settings.redis_url,
                encoding="utf-8",
                decode_responses=True
            )
            await redis_client.ping()
            await redis_client.aclose()
            logger.info("Redis connection validated at startup")
            
        except RuntimeError:
            # Re-raise our own RuntimeErrors
            raise
        except Exception as e:
            error_msg = f"CRITICAL: Redis connection failed at startup: {str(e)}. Rate limiting requires Redis in production."
            logger.error(error_msg)
            raise RuntimeError(error_msg)
    
    # 3. Validate Required Secrets at startup (Production only)
    if is_production:
        try:
            from security.secret_manager import get_secret_manager
            secret_mgr = get_secret_manager()
            
            # Check for required secrets
            required_secrets = ["JWT_SECRET_KEY", "DATABASE_URL"]
            missing = []
            for secret in required_secrets:
                if not secret_mgr.get_secret(secret):
                    missing.append(secret)
            
            if missing:
                error_msg = f"CRITICAL: Missing required secrets at startup: {', '.join(missing)}"
                logger.error(error_msg)
                raise RuntimeError(error_msg)
            
            logger.info("Secret manager validated at startup")
            
        except RuntimeError:
            raise
        except Exception as e:
            error_msg = f"CRITICAL: Secret manager validation failed: {str(e)}"
            logger.error(error_msg)
            raise RuntimeError(error_msg)
    
    # 4. Try to initialize database (after validation)
    try:
        await init_database()
        logger.info("Database initialized successfully")
    except Exception as e:
        error_msg = f"CRITICAL: Database initialization failed: {str(e)}"
        logger.error(error_msg)
        if is_production:
            raise RuntimeError(error_msg)
        logger.warning(f"Database initialization warning: {str(e)}")
    
    logger.info("AegisLM startup validation complete - ready to serve traffic")
    
    yield
    
    # Shutdown
    logger.info("Shutting down AegisLM backend")
    
    try:
        # Close database connections
        await close_database()
        logger.info("Database connections closed")
    except Exception as e:
        logger.error(
            f"Error during shutdown: {str(e)}",
            exception=e
        )


# =============================================================================
# FastAPI Application
# =============================================================================

app = FastAPI(
    title="AegisLM",
    description="Multi-Agent Adversarial LLM Evaluation Framework",
    version="0.1.0",
    docs_url="/docs",
    redoc_url="/redoc",
    lifespan=lifespan,
)


# =============================================================================
# Middleware
# =============================================================================

# CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Configure appropriately for production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# =============================================================================
# Exception Handlers
# =============================================================================

@app.exception_handler(AegisLMException)
async def aegislm_exception_handler(request: Request, exc: AegisLMException):
    """Handle custom AegisLM exceptions."""
    logger.error(
        f"AegisLM exception: {exc.message}",
        metadata={"code": exc.code, "path": str(request.url)},
        exception=exc
    )
    
    return JSONResponse(
        status_code=getattr(exc, "status_code", 500),
        content=exc.to_dict(),
    )


@app.exception_handler(ValueError)
async def value_error_handler(request: Request, exc: ValueError):
    """Handle ValueError exceptions."""
    logger.error(
        f"Value error: {str(exc)}",
        metadata={"path": str(request.url)},
        exception=exc
    )
    
    return JSONResponse(
        status_code=400,
        content={
            "error": "ValidationError",
            "code": "VALIDATION_ERROR",
            "message": str(exc),
        },
    )


@app.exception_handler(Exception)
async def generic_exception_handler(request: Request, exc: Exception):
    """Handle generic exceptions."""
    logger.error(
        f"Unexpected error: {str(exc)}",
        metadata={"path": str(request.url)},
        exception=exc
    )
    
    return JSONResponse(
        status_code=500,
        content={
            "error": "InternalServerError",
            "code": "INTERNAL_ERROR",
            "message": "An unexpected error occurred",
        },
    )


# =============================================================================
# Include Routers
# =============================================================================

app.include_router(routes.router)
app.include_router(public_router)
app.include_router(regulatory_router)


# =============================================================================
# Root Endpoints
# =============================================================================

@app.get("/")
async def root():
    """Root endpoint."""
    return {
        "name": "AegisLM",
        "version": "0.1.0",
        "description": "Multi-Agent Adversarial LLM Evaluation Framework",
        "docs": "/docs",
    }


@app.get("/ping")
async def ping():
    """Ping endpoint for health checks."""
    return {"status": "ok"}


# =============================================================================
# Observability Endpoints (Week 6 Day 5)
# =============================================================================

@app.get("/metrics")
async def metrics():
    """
    Prometheus-style metrics endpoint.
    
    Returns metrics in Prometheus text format for scraping.
    Compatible with Prometheus, Grafana, and other observability tools.
    """
    try:
        from observability.exporters import export_prometheus_metrics
        from fastapi.responses import PlainTextResponse
        
        metrics_output = export_prometheus_metrics()
        
        return PlainTextResponse(
            content=metrics_output,
            media_type="text/plain; version=0.0.4; charset=utf-8"
        )
    except Exception as e:
        logger.error(f"Failed to export metrics: {str(e)}")
        return JSONResponse(
            status_code=500,
            content={"error": "Failed to export metrics", "detail": str(e)}
        )


@app.get("/health/live")
async def liveness():
    """
    Kubernetes-style liveness check.
    
    Confirms:
    - Process is running
    - No fatal crash
    
    Returns 200 if alive, 503 if not.
    """
    return {
        "status": "alive",
        "service": "aegislm"
    }


@app.get("/health/ready")
async def readiness():
    """
    Kubernetes-style readiness check.
    
    Confirms:
    - Database reachable
    - Redis reachable
    - Worker pool active
    - Model loaded
    - Secret manager valid
    
    Returns detailed status of all dependencies.
    """
    from backend.db.session import check_database_connection
    
    # Check database
    db_ready = await check_database_connection()
    
    # Check Redis (if available)
    redis_ready = True
    try:
        from backend.db.session import get_redis_client
        redis_client = get_redis_client()
        if redis_client:
            await redis_client.ping()
    except Exception:
        # Redis is optional - if not configured, mark as ready
        redis_ready = True
    
    # Check workers - In development mode, workers are optional
    workers_ready = True
    active_workers = 0
    import os
    is_production = os.getenv("AEGISLM_ENV", "development").lower() == "production"
    try:
        from backend.queue.worker_registry import get_worker_registry
        registry = get_worker_registry()
        metrics = await registry.get_worker_metrics()
        active_workers = metrics.active_workers
    except Exception:
        # Registry not available - workers are not running
        active_workers = 0
    
    # In production, workers are required; in dev, they're optional
    if is_production:
        workers_ready = active_workers > 0
    else:
        workers_ready = True  # Dev mode - workers are optional
    
    # Check model loaded (simplified - model loading is handled by orchestrator)
    model_loaded = True
    try:
        from backend.core.model_registry import get_model_registry
        model_registry = get_model_registry()
        model_loaded = len(model_registry.list_models()) >= 0
    except Exception:
        model_loaded = True  # Default to True if model loading is handled separately
    
    # Check secret manager
    secret_manager_valid = True
    try:
        from security.secret_manager import get_secret_manager
        secret_mgr = get_secret_manager()
        secret_manager_valid = secret_mgr is not None
    except Exception:
        secret_manager_valid = True  # Default to True if not critical
    
    # Determine overall readiness
    ready = db_ready and redis_ready and workers_ready and model_loaded and secret_manager_valid
    
    return JSONResponse(
        status_code=200 if ready else 503,
        content={
            "status": "ready" if ready else "not_ready",
            "checks": {
                "database": db_ready,
                "redis": redis_ready,
                "workers": workers_ready,
                "model_loaded": model_loaded,
                "secret_manager": secret_manager_valid
            },
            "active_workers": active_workers,
        }
    )


# =============================================================================
# Main Entry Point
# =============================================================================

if __name__ == "__main__":
    import uvicorn
    
    uvicorn.run(
        "backend.main:app",
        host=settings.api_host,
        port=settings.api_port,
        reload=False,
        log_level=settings.log_level.lower(),
    )