from starlite import Request
from starlite.response import Response
from starlite.types import ASGIApp, Message, Receive, Scope, Send

from starlite_starception.exception_handler import exception_handler


class StarceptionMiddleware:
    """
    Handles returning 500 responses when a server error occurs.

    If 'debug' is set, then traceback responses will be returned,
    otherwise the designated 'handler' will be called.

    This middleware class should generally be used to wrap *everything*
    else up, so that unhandled exceptions anywhere in the stack
    always result in an appropriate 500 response.
    """

    def __init__(self, app: ASGIApp) -> None:
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http":  # pragma: nocover
            await self.app(scope, receive, send)
            return

        response_started = False

        async def _send(message: Message) -> None:
            nonlocal response_started, send

            if message["type"] == "http.response.start":  # pragma: nocover
                response_started = True
            await send(message)

        try:
            await self.app(scope, receive, _send)
        except Exception as exc:
            if not response_started:
                request = Request(scope)
                if request.app.debug:
                    # In debug mode, return traceback responses.
                    response = self.debug_response(request, exc)
                else:
                    # Use our default 500 error handler.
                    response = self.error_response(request, exc)

                await response(scope, receive, send)

            # We always continue to raise the exception.
            # This allows servers to log the error, or allows test clients
            # to optionally raise the error within the test case.
            raise exc

    def debug_response(self, request: Request, exc: Exception) -> Response:
        return exception_handler(request, exc)

    def error_response(self, request: Request, exc: Exception) -> Response:
        return Response("Internal Server Error", status_code=500, media_type="text/plain")
