from fastapi import FastAPI, Request, status
from fastapi.responses import JSONResponse
from fastapi.exceptions import RequestValidationError
from starlette.exceptions import HTTPException as StarletteHTTPException
import logging

logger = logging.getLogger(__name__)

class InferenceError(Exception):
    pass

class ModelLoadError(Exception):
    pass

class ImageProcessingError(Exception):
    pass

async def validation_exception_handler(request: Request, exc: RequestValidationError):
    logger.warning(f"Validation error: {exc.errors()}")
    return JSONResponse(
        status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
        content={"detail": exc.errors(), "body": exc.body}
    )

async def http_exception_handler(request: Request, exc: StarletteHTTPException):
    logger.warning(f"HTTP exception: {exc.status_code} - {exc.detail}")
    return JSONResponse(
        status_code=exc.status_code,
        content={"detail": exc.detail}
    )

async def inference_exception_handler(request: Request, exc: InferenceError):
    logger.error(f"Inference error: {str(exc)}")
    return JSONResponse(
        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        content={"detail": f"Inference error: {str(exc)}"}
    )

async def model_load_exception_handler(request: Request, exc: ModelLoadError):
    logger.error(f"Model load error: {str(exc)}")
    return JSONResponse(
        status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
        content={"detail": f"Model load error: {str(exc)}"}
    )

async def image_processing_exception_handler(request: Request, exc: ImageProcessingError):
    logger.error(f"Image processing error: {str(exc)}")
    return JSONResponse(
        status_code=status.HTTP_400_BAD_REQUEST,
        content={"detail": f"Image processing error: {str(exc)}"}
    )

async def general_exception_handler(request: Request, exc: Exception):
    logger.exception(f"Unhandled exception: {str(exc)}")
    return JSONResponse(
        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        content={"detail": "Internal server error"}
    )

def setup_exception_handlers(app: FastAPI):
    app.add_exception_handler(RequestValidationError, validation_exception_handler)
    app.add_exception_handler(StarletteHTTPException, http_exception_handler)
    app.add_exception_handler(InferenceError, inference_exception_handler)
    app.add_exception_handler(ModelLoadError, model_load_exception_handler)
    app.add_exception_handler(ImageProcessingError, image_processing_exception_handler)
    app.add_exception_handler(Exception, general_exception_handler)