src/main.py

"""FastAPI application entrypoint for imageO_v3."""

from __future__ import annotations

from contextlib import asynccontextmanager
from typing import Optional

from asgi_correlation_id import CorrelationIdMiddleware
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from fastapi_limiter import FastAPILimiter

from .api.v1.router import api_router as v1_router
from .config import settings
from .core.exceptions import ImageTooLargeError, InferenceServiceError, InvalidImageError, UnsupportedImageFormatError
from .core.logging import get_logger, setup_logging
from .core.middleware import AccessLogMiddleware
from .core.services.ml_inference import initialize_models
from .core.services.redis_client import redis_manager
from .repository import DeviceRepository

setup_logging()
logger = get_logger()

device_repo: Optional[DeviceRepository] = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Startup/shutdown lifecycle hooks."""
    global device_repo
    logger.info("Starting imageO_v3 application", app_env=settings.app_env)

    # Redis and device repository startup.
    try:
        if redis_manager.client:
            await FastAPILimiter.init(redis_manager.client)
            logger.info("Rate limiter initialized with Redis backend")
            device_repo = DeviceRepository()
            await device_repo._load_credentials()
        else:
            logger.warning("Redis unavailable at startup; auth lock/rate limiting will be disabled")
    except Exception as exc:
        logger.error("Redis/bootstrap startup error", error=str(exc), exc_info=True)
        redis_manager.client = None
        if settings.app_env == "production":
            raise

    # Model startup.
    try:
        initialize_models()
        logger.info("Inference models initialized")
    except Exception as exc:
        logger.error("Model initialization error", error=str(exc), exc_info=True)
        if settings.app_env == "production":
            raise

    logger.info(
        "Application startup complete",
        model_version=settings.model_version,
        log_level=settings.log_level,
        max_image_size_mb=settings.max_image_size_mb,
    )
    yield

    logger.info("Shutting down imageO_v3 application")
    if redis_manager.client:
        await FastAPILimiter.close()
        await redis_manager.client.aclose()
        logger.info("Redis connection closed")


def create_app() -> FastAPI:
    """Create and configure FastAPI app instance."""
    app = FastAPI(
        title="imageO_v3",
        description="Production-grade FastAPI service for secure ML image inference",
        version=settings.model_version,
        lifespan=lifespan,
        docs_url="/docs" if settings.app_env != "production" else None,
        redoc_url="/redoc" if settings.app_env != "production" else None,
    )

    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_credentials=True,
        allow_methods=["GET", "POST"],
        allow_headers=["*"],
    )
    app.add_middleware(AccessLogMiddleware)
    app.add_middleware(
        CorrelationIdMiddleware,
        header_name="X-Request-ID",
        update_request_header=True,
    )

    app.include_router(v1_router)

    @app.exception_handler(InvalidImageError)
    async def invalid_image_handler(request: Request, exc: InvalidImageError) -> JSONResponse:
        logger.warning("Invalid image error", error=str(exc))
        return JSONResponse(status_code=400, content={"error": "Invalid or corrupted image file."})

    @app.exception_handler(UnsupportedImageFormatError)
    async def unsupported_format_handler(request: Request, exc: UnsupportedImageFormatError) -> JSONResponse:
        logger.warning("Unsupported image format", error=str(exc))
        return JSONResponse(status_code=400, content={"error": "Unsupported image format. Allowed: JPEG, PNG, WEBP, BMP, TIFF."})

    @app.exception_handler(ImageTooLargeError)
    async def image_too_large_handler(request: Request, exc: ImageTooLargeError) -> JSONResponse:
        logger.warning("Image too large", error=str(exc))
        return JSONResponse(
            status_code=400,
            content={"error": f"Image exceeds maximum size of {settings.max_image_size_mb} MB."},
        )

    @app.exception_handler(InferenceServiceError)
    async def inference_error_handler(request: Request, exc: InferenceServiceError) -> JSONResponse:
        logger.error("Inference service error", error=str(exc))
        return JSONResponse(
            status_code=500,
            content={"error": "An error occurred during prediction. Please try again."},
        )

    @app.exception_handler(Exception)
    async def general_exception_handler(request: Request, exc: Exception) -> JSONResponse:
        logger.error("Unhandled exception", error=str(exc), exc_info=True)
        return JSONResponse(
            status_code=500,
            content={"error": "An error occurred during prediction. Please try again."},
        )

    @app.get("/", tags=["root"])
    async def root() -> dict[str, str]:
        return {
            "service": "imageO_v3",
            "version": settings.model_version,
            "status": "running",
            "docs": "/docs" if settings.app_env != "production" else "disabled",
        }

    @app.get("/healthz", tags=["health"])
    async def healthz() -> dict[str, str]:
        return {"status": "ok"}

    return app


app = create_app()


def get_device_repo() -> DeviceRepository:
    """Dependency provider for initialized DeviceRepository."""
    if device_repo is None:
        raise RuntimeError("Device repository is not initialized.")
    return device_repo