from starlette.exceptions import HTTPException
from fastapi.responses import ORJSONResponse


async def http_exception(request, exc):
    if hasattr(exc, 'detail'):
        msg = exc.detail
    else:
        msg = exc
    return ORJSONResponse({"msg": msg}, status_code=exc.status_code)


# async def http_404_exception(request, exc):
#     return ORJSONResponse({"msg": exc.detail}, status_code=exc.status_code)
#
#
# async def http_500_exception(request, exc):
#     return ORJSONResponse({"msg": exc.detail}, status_code=exc.status_code)


# 自定义异常
class UnicornException(Exception):
    def __init__(self, name: str):
        self.name = name


async def unicorn_exception_handler(request, exc: UnicornException):
    return ORJSONResponse(
        status_code=418,
        content={"message": f"Oops! {exc.name} did something. There goes a rainbow..."},
    )


exception_handlers = {
    HTTPException: http_exception,
    # 404: http_404_exception,
    # 500: http_500_exception,
    UnicornException: unicorn_exception_handler
}

__all__ = [exception_handlers]