main.py

import logging
import asyncio
from routers import health, predict, home
from contextlib import asynccontextmanager

from fastapi import FastAPI

# Import deployment functions
from models.artifacts_loader import get_artifacts
from core.config import settings
from core.logging import setup_logging
from core.exceptions import (
    ModelPredictionError,
    ArtifactLoadError,
    generic_exception_handler,
    model_prediction_exception_handler,
)

# Setup logging before app startup
setup_logging()
logger = logging.getLogger(__name__)

# -------------------------------------------------------------------------
# 1. Global State Management (The safe way to load ML models in APIs)
# -------------------------------------------------------------------------


@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    Lifespan context manager — non-blocking startup.

    Model loading is fired as a background task so the server yields
    immediately and can answer /health/liveness probes right away.
    Traffic is withheld naturally by /health/readiness (returns 503)
    until `artifacts.is_loaded` becomes True.
    """
    logger.info(f"Starting {settings.PROJECT_NAME} v{settings.VERSION}")

    async def _load():
        try:
            await asyncio.to_thread(get_artifacts().load_artifacts)
        except ArtifactLoadError as e:
            logger.error(f"Failed to load artifacts: {e.detail}")
        except Exception as e:
            logger.critical(f"Unhandled error loading artifacts: {e}")

    # Fire-and-forget — the server becomes ready for probes instantly.
    asyncio.create_task(_load())

    yield  # Server accepts connections NOW (before model loading completes)

    # Cleanup on shutdown
    logger.info("Shutting down API...")
    get_artifacts().clear()


# -------------------------------------------------------------------------
# 2. App Initialization
# -------------------------------------------------------------------------
app = FastAPI(
    title=settings.PROJECT_NAME,
    description=settings.DESCRIPTION,
    version=settings.VERSION,
    lifespan=lifespan,
)

# Exception Handlers
app.add_exception_handler(Exception, generic_exception_handler)
app.add_exception_handler(ModelPredictionError, model_prediction_exception_handler)

# -------------------------------------------------------------------------
# 3. API Endpoints / Routers
# -------------------------------------------------------------------------
API_V1_PREFIX = "/api/v1"

app.include_router(home.router)                                          # Serves HTML landing page — no versioning needed
app.include_router(health.router, prefix=API_V1_PREFIX)                  # GET  /api/v1/health
app.include_router(predict.router, prefix=API_V1_PREFIX)                 # POST /api/v1/predict