import logging

from starlette import status
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.requests import Request
from starlette.responses import Response

logger = logging.getLogger(__name__)


class SuppressNoResponseReturnedMiddleware(BaseHTTPMiddleware):
    """
    Handle disconnected/cancelled requests.
    https://github.com/encode/starlette/discussions/1527#discussioncomment-2234702
    """

    async def dispatch(
        self, request: Request, call_next: RequestResponseEndpoint
    ) -> Response:
        try:
            return await call_next(request)
        except RuntimeError as exc:
            if str(exc) == "No response returned." and await request.is_disconnected():
                logger.warning(
                    "Error `No response returned` detected. "
                    "At the same time we know that the client is disconnected "
                    "and not waiting for a response."
                )
                return Response(status_code=status.HTTP_204_NO_CONTENT)
            raise