src/main.py

import logging
import os
import sys
from contextlib import asynccontextmanager

from fastapi import FastAPI, Request
from fastapi.staticfiles import StaticFiles
from fastapi.responses import JSONResponse

from .config import OUTPUT_DIR, MAX_AGE_DAYS, MAX_FILES, CLEANUP_INTERVAL
from .controllers.web_controller import router as web_router
from .utils import RateLimitMiddleware, ConcurrencyLimitMiddleware, RequestSizeLimitMiddleware, perform_cleanup

# Ensure registry is initialized
from .models import get_field_registry_manager

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("aibom_generator")

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup
    logger.info("Starting AI SBOM Generator WebApp")
    try:
        get_field_registry_manager() # Ensure registry is loaded
        logger.info("Registry loaded successfully")
    except Exception as e:
        logger.error(f"Failed to load registry: {e}")
        
    # Initial cleanup
    try:
        perform_cleanup(OUTPUT_DIR, MAX_AGE_DAYS, MAX_FILES)
    except Exception as e:
        logger.warning(f"Initial cleanup failed: {e}")
        
    yield
    # Shutdown (if needed)

app = FastAPI(title="AI SBOM Generator", lifespan=lifespan)

# --- Middleware ---
app.add_middleware(
    RateLimitMiddleware,
    rate_limit_per_minute=10,
    rate_limit_window=60,
    protected_routes=["/generate"]
)
app.add_middleware(
    ConcurrencyLimitMiddleware,
    max_concurrent_requests=5,
    timeout=5.0,
    protected_routes=["/generate"]
)
app.add_middleware(
    RequestSizeLimitMiddleware,
    max_content_length=1024*1024  # 1MB
)

# --- Cleanup Middleware ---
request_counter = 0

@app.middleware("http")
async def cleanup_middleware(request: Request, call_next):
    global request_counter
    request_counter += 1
    if request_counter % CLEANUP_INTERVAL == 0:
        try:
            removed = perform_cleanup(OUTPUT_DIR, MAX_AGE_DAYS, MAX_FILES)
            logger.info(f"Scheduled cleanup removed {removed} files")
        except Exception as e:
            logger.error(f"Error during scheduled cleanup: {e}")
            
    response = await call_next(request)
    return response

# --- Static Files ---
os.makedirs(OUTPUT_DIR, exist_ok=True)
app.mount("/output", StaticFiles(directory=OUTPUT_DIR), name="output")
# Mount static files (CSS/JS)
os.makedirs("src/static", exist_ok=True)
app.mount("/static", StaticFiles(directory="src/static"), name="static")

# --- Routes ---
app.include_router(web_router)

if __name__ == "__main__":
    import uvicorn
    # Print clear access URL to avoid 0.0.0.0 confusion
    print("🚀 Application ready! Access it at: http://localhost:8000")
    uvicorn.run("src.main:app", host="0.0.0.0", port=8000, reload=True)