backend/main.py

"""
Main FastAPI application for AegisLM SaaS Backend.

Production-ready application with proper middleware,
routing, error handling, and startup/shutdown events.
"""

import time
import asyncio
import os
import json
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from fastapi import FastAPI, Request, HTTPException, status, WebSocket, WebSocketDisconnect
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.exceptions import RequestValidationError
from starlette.exceptions import HTTPException as StarletteHTTPException
from typing import List
import sys
import os
import uvicorn
# Fix for nested imports (AegisLM SaaS Structure)
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# Initialize logging first
from logging_config import setup_logging, get_logger

# Setup logging
setup_logging()
logger = get_logger(__name__)

# Load environment variables
from dotenv import load_dotenv

load_dotenv()

from core.config import settings
from core.database import init_db, close_db, check_db_health, check_redis_health
from api.routes import (
    auth_routes,
    evaluation_routes,
    results_routes,
    analytics_routes,
    benchmark_routes,
    report_routes,
    experiment_routes,
    multimodal_routes,
    performance_routes,
    monitoring_routes,
    sqlite_fallback_routes,
    audit_routes,
    dataset_routes
)
from websocket.routes import websocket as websocket_routes
from middleware.rate_limit import (
    RateLimitMiddleware,
    APIKeyAuthMiddleware,
    SecurityHeadersMiddleware,
    RequestLoggingMiddleware
)
from middleware.validation import InputValidationMiddleware
from middleware.error_handler import GlobalErrorHandler
from schemas.response_schema import ErrorResponse, HealthCheckResponse


@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    Application lifespan manager.
    
    Handles startup and shutdown events.
    """
    # Startup
    print("Starting AegisLM Red Team API...")
    
    try:
        # Skip production startup validation in debug mode
        if not settings.DEBUG:
          if not hasattr(settings, "STARTUP_VALIDATION") or settings.DEBUG:
            # TODO: Re-enable when startup_validation module is available
            print("[!] Production validation module not available, skipping...")
        else:
            print("[DEBUG] Debug mode enabled - skipping production validation")
        
        # Initialize database (skip if not available)
        try:
            await init_db()
            print("✅ Database initialized successfully")
            
            # Setup SQLite fallback if enabled
            if settings.ENABLE_SQLITE_FALLBACK:
                try:
                    from core.sqlite_fallback_manager import setup_sqlite_fallback
                    fallback_setup = await setup_sqlite_fallback()
                    if fallback_setup:
                        print("[INFO] SQLite fallback database setup completed")
                    else:
                        print("[!] SQLite fallback database setup failed")
                except Exception as e:
                    print(f"[!] SQLite fallback setup skipped: {str(e)}")
            
        except Exception as e:
            print(f"[!] Database initialization skipped: {str(e)}")
            print("   (This is expected if PostgreSQL is not running)")
        
        # Skip redundant health checks since production validation already verified services
        print("[SKIP] Skipping redundant health checks (production validation already passed)")
        
        # Set startup time
        app.state.start_time = time.time()
        
        print(f"[INFO] Server running on {settings.API_V1_STR}")
        print(f"[INFO] Debug mode: {settings.DEBUG}")
        print(f"[INFO] Environment: {'Production' if not settings.DEBUG else 'Development'}")
        
    except Exception as e:
        print(f"Startup failed: {str(e)}")
        raise
    
    yield
    
    # Shutdown
    print("🛑 Shutting down AegisLM Red Team API...")
    
    try:
        await close_db()
        print("✅ Database connections closed")
    except Exception as e:
        print(f"❌ Shutdown error: {str(e)}")


# Create FastAPI application
app = FastAPI(
    title=settings.APP_NAME,
    version=settings.APP_VERSION,
    description="Production-ready API for AI red team evaluations and security assessments",
    docs_url="/docs" if settings.DEBUG else None,
    redoc_url="/redoc" if settings.DEBUG else None,
    openapi_url="/openapi.json" if settings.DEBUG else None,
    lifespan=lifespan
)

# Include API routes
# Add custom middleware (order matters)
app.add_middleware(GlobalErrorHandler)  # Add first to catch all errors
app.add_middleware(InputValidationMiddleware)  # Add input validation
app.add_middleware(SecurityHeadersMiddleware)
app.add_middleware(RequestLoggingMiddleware)
app.add_middleware(APIKeyAuthMiddleware)
app.add_middleware(RateLimitMiddleware)

# Add CORS middleware LAST so it runs FIRST on incoming requests
# This is critical for handling OPTIONS (preflight) requests before auth/rate-limiting
app.add_middleware(
    CORSMiddleware,
    allow_origins=settings.BACKEND_CORS_ORIGINS,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Include API routes
app.include_router(
    auth_routes.router,
    prefix=settings.API_V1_STR,
    tags=["authentication"]
)
# 
app.include_router(
    evaluation_routes.router,
    prefix=settings.API_V1_STR,
    tags=["evaluations"]
)

app.include_router(
    results_routes.router,
    prefix=settings.API_V1_STR,
    tags=["results"]
)

app.include_router(
    analytics_routes.router,
    prefix=settings.API_V1_STR,
    tags=["analytics"]
)

app.include_router(
    benchmark_routes.router,
    prefix=settings.API_V1_STR,
    tags=["benchmarks"]
)

app.include_router(
    report_routes.router,
    prefix=settings.API_V1_STR,
    tags=["reports"]
)

app.include_router(
    experiment_routes.router,
    prefix=settings.API_V1_STR,
    tags=["experiments"]
)

app.include_router(
    multimodal_routes.router,
    prefix=settings.API_V1_STR,
    tags=["multimodal"]
)

app.include_router(
    performance_routes.router,
    prefix=settings.API_V1_STR,
    tags=["performance"]
)
# 
app.include_router(
    monitoring_routes.router,
    prefix=settings.API_V1_STR,
    tags=["monitoring"]
)

app.include_router(
    sqlite_fallback_routes.router,
    prefix=settings.API_V1_STR,
    tags=["sqlite-fallback"]
)

app.include_router(
    audit_routes.router,
    prefix=settings.API_V1_STR,
    tags=["audit"]
)

app.include_router(
    dataset_routes.router,
    prefix=settings.API_V1_STR,
    tags=["datasets"]
)


# Include WebSocket routes
app.include_router(
    websocket_routes,
    tags=["websocket"]
)


# Root endpoint
@app.get("/", tags=["root"])
async def root():
    """
    Root endpoint.
    
    Returns:
        dict: API information
    """
    return {
        "name": settings.APP_NAME,
        "version": settings.APP_VERSION,
        "status": "running",
        "docs_url": "/docs" if settings.DEBUG else None,
        "api_v1": settings.API_V1_STR
    }


@app.get("/api/v1/ws/health", tags=["health"])
async def ws_health_check():
    """
    Health check for WebSocket connections.
    Used by the frontend to monitor connection status.
    """
    from datetime import datetime, timezone
    return {
        "status": "healthy",
        "websocket": "available",
        "timestamp": datetime.now(timezone.utc).isoformat()
    }


@app.get("/api/v1/health", tags=["health"])
async def api_health_check():
    """
    API version of the health check.
    """
    return await health_check()


# Health check endpoint
@app.get("/health", response_model=HealthCheckResponse, tags=["health"])
async def health_check():
    """
    Comprehensive health check endpoint.
    
    Returns:
        HealthCheckResponse: Health status of all services
    """
    # Check database health
    db_healthy = await check_db_health()
    redis_healthy = await check_redis_health()
    
    # Check worker status (via Celery)
    worker_healthy = False
    try:
        from workers.celery_worker import celery_app
        # Check if Celery worker is responsive
        inspect = celery_app.control.inspect()
        active_tasks = inspect.active()
        worker_healthy = active_tasks is not None
    except Exception:
        worker_healthy = False
    
    # Overall status (only critical services affect overall health)
    overall_status = "healthy" if db_healthy and redis_healthy else "unhealthy"
    
    # Service statuses
    services = {
        "database": "healthy" if db_healthy else "unhealthy",
        "redis": "healthy" if redis_healthy else "unhealthy",
        "worker": "healthy" if worker_healthy else "unhealthy",
        "api": "healthy"
    }
    
    # Calculate uptime
    uptime_seconds = int(time.time() - app.state.start_time) if hasattr(app.state, "start_time") else 0
    
    # Additional health information
    health_info = {
        "environment": "production" if not settings.DEBUG else "development",
        "learning_enabled": settings.ENABLE_LEARNING,
        "judge_enabled": settings.ENABLE_LLM_JUDGE,
        "rate_limiting": f"{settings.RATE_LIMIT_PER_MINUTE}/min"
    }
    
    return HealthCheckResponse(
        status=overall_status,
        version=settings.APP_VERSION,
        timestamp=datetime.now(timezone.utc).isoformat(),
        services=services,
        uptime_seconds=uptime_seconds,
        **health_info
    )


# Metrics endpoint
@app.get("/metrics", tags=["monitoring"])
async def metrics():
    """
    Basic metrics endpoint.
    
    Returns:
        dict: Application metrics
    """
    # This is a basic metrics endpoint
    # In production, you might want to use Prometheus or similar
    return {
        "status": "ok",
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "uptime_seconds": int(time.time() - app.state.start_time) if hasattr(app.state, "start_time") else 0,
        "version": settings.APP_VERSION,
        "debug": settings.DEBUG
    }


# Exception handlers
@app.exception_handler(StarletteHTTPException)
async def http_exception_handler(request: Request, exc: StarletteHTTPException):
    """
    Handle HTTP exceptions.
    
    Args:
        request: HTTP request
        exc: HTTP exception
        
    Returns:
        JSONResponse: Error response
    """
    from fastapi.responses import JSONResponse
    from datetime import datetime, timezone
    
    return JSONResponse(
        status_code=exc.status_code,
        content=ErrorResponse(
            success=False,
            message=exc.detail,
            error_code=f"HTTP_{exc.status_code}",
            timestamp=datetime.now(timezone.utc).isoformat()
        ).dict()
    )


@app.exception_handler(RequestValidationError)
async def validation_exception_handler(request: Request, exc: RequestValidationError):
    """
    Handle validation exceptions.
    
    Args:
        request: HTTP request
        exc: Validation exception
        
    Returns:
        JSONResponse: Validation error response
    """
    from fastapi.responses import JSONResponse
    from datetime import datetime, timezone
    from logging_config import get_logger
    
    module_logger = get_logger(__name__)
    
    # Log the validation errors for debugging
    module_logger.error(f"Validation error occurred: {exc.errors()}")
    
    # Clean validation errors to remove non-serializable objects
    cleaned_errors = []
    for error in exc.errors():
        cleaned_error = error.copy()
        # Remove non-serializable objects from ctx
        if 'ctx' in cleaned_error and 'error' in cleaned_error['ctx']:
            cleaned_error['ctx']['error'] = str(cleaned_error['ctx']['error'])
        cleaned_errors.append(cleaned_error)
    
    return JSONResponse(
        status_code=422,
        content=ErrorResponse(
            success=False,
            message="Validation failed",
            error_code="VALIDATION_ERROR",
            error_details={
                "validation_errors": cleaned_errors
            },
            timestamp=datetime.now(timezone.utc).isoformat()
        ).dict()
    )


@app.exception_handler(Exception)
async def general_exception_handler(request: Request, exc: Exception):
    """
    Handle general exceptions.
    
    Args:
        request: HTTP request
        exc: General exception
        
    Returns:
        JSONResponse: Error response
    """
    from fastapi.responses import JSONResponse
    from datetime import datetime, timezone
    from logging_config import get_logger
    import traceback
    
    module_logger = get_logger(__name__)
    
    # Log the full exception 
    error_id = f"INTERNAL_{id(exc)}"
    module_logger.error(
        f"Unhandled exception [{error_id}]: {str(exc)}\n"
        f"Traceback: {traceback.format_exc()}"
    )
    
    return JSONResponse(
        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        content=ErrorResponse(
            success=False,
            message="An internal server error occurred. Please try again later.",
            error_code="INTERNAL_ERROR",
            error_details={"error_id": error_id},
            timestamp=datetime.now(timezone.utc).isoformat()
        ).model_dump()
    )


# Startup event to set start time
@app.on_event("startup")
async def startup_event():
    """
    Set application start time.
    """
    app.state.start_time = time.time()


# Additional utility endpoints
@app.get("/api/v1/info", tags=["info"])
async def api_info():
    """
    Get API information.
    
    Returns:
        dict: API information
    """
    return {
        "name": settings.APP_NAME,
        "version": settings.APP_VERSION,
        "api_version": "v1",
        "description": "Production-ready API for AI red team evaluations",
        "endpoints": {
            "authentication": f"{settings.API_V1_STR}/auth",
            "evaluations": f"{settings.API_V1_STR}/evaluations",
            "results": f"{settings.API_V1_STR}/results",
            "benchmarks": f"{settings.API_V1_STR}/benchmarks"
        },
        "documentation": "/docs" if settings.DEBUG else "Documentation not available in production"
    }


async def check_celery_health():
    """Check if Celery workers are active."""
    try:
        from workers.celery_worker import celery_app
        # Use ping() to check worker connectivity
        # Timeout after 1 second to avoid blocking
        ping_result = celery_app.control.ping(timeout=1.0)
        return ping_result is not None and len(ping_result) > 0
    except Exception:
        return False


@app.get("/api/v1/status", tags=["status"])
async def api_status():
    """
    Get API status.
    
    Returns:
        dict: API status
    """
    db_healthy = await check_db_health()
    redis_healthy = await check_redis_health()
    celery_healthy = await check_celery_health()
    
    return {
        "status": "operational" if db_healthy and redis_healthy and celery_healthy else "degraded",
        "services": {
            "database": "online" if db_healthy else "offline",
            "redis": "online" if redis_healthy else "offline",
            "celery": "online" if celery_healthy else "offline"
        },
        "uptime_seconds": int(time.time() - app.state.start_time) if hasattr(app.state, "start_time") else 0,
        "version": settings.APP_VERSION
    }


# Development-only endpoints
if settings.DEBUG:
    @app.get("/debug/routes", tags=["debug"])
    async def debug_routes():
        """
        Debug endpoint to list all routes.
        
        Returns:
            dict: Route information
        """
        routes = []
        for route in app.routes:
            routes.append({
                "path": route.path,
                "name": route.name,
                "methods": getattr(route, "methods", None),
                "include_in_schema": getattr(route, "include_in_schema", None)
            })
        
        return {"routes": routes}
    
    @app.get("/debug/config", tags=["debug"])
    async def debug_config():
        """
        Debug endpoint to show configuration (masked).
        
        Returns:
            dict: Configuration information
        """
        from core.security import mask_sensitive_data
        
        return {
            "app_name": settings.APP_NAME,
            "app_version": settings.APP_VERSION,
            "debug": settings.DEBUG,
            "database_url": mask_sensitive_data(settings.DATABASE_URL),
            "redis_url": settings.REDIS_URL,
            "cors_origins": settings.BACKEND_CORS_ORIGINS,
            "rate_limit_per_minute": settings.RATE_LIMIT_PER_MINUTE
        }


if __name__ == "__main__":
    import uvicorn
    
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=8001,
        reload=settings.DEBUG,
        log_level="info" if not settings.DEBUG else "debug"
    )