from __future__ import annotations

from collections.abc import Iterable
from typing import Any

from starlette.responses import JSONResponse


class _RequestBodyTooLarge(RuntimeError):
    pass


class RequestBodyLimitMiddleware:
    def __init__(self, app: Any, *, max_bytes: int) -> None:
        if max_bytes <= 0:
            raise ValueError("max_bytes must be positive")
        self.app = app
        self.max_bytes = max_bytes

    async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None:
        if scope.get("type") != "http":
            await self.app(scope, receive, send)
            return
        declared = _content_length(scope.get("headers", ()))
        if declared is not None and declared > self.max_bytes:
            await self._reject(scope, receive, send)
            return

        received = 0
        response_started = False

        async def limited_receive() -> dict[str, Any]:
            nonlocal received
            message = await receive()
            if message.get("type") == "http.request":
                received += len(message.get("body", b""))
                if received > self.max_bytes:
                    raise _RequestBodyTooLarge
            return message

        async def tracked_send(message: dict[str, Any]) -> None:
            nonlocal response_started
            if message.get("type") == "http.response.start":
                response_started = True
            await send(message)

        try:
            await self.app(scope, limited_receive, tracked_send)
        except _RequestBodyTooLarge:
            if response_started:
                raise
            await self._reject(scope, receive, send)

    async def _reject(self, scope: dict[str, Any], receive: Any, send: Any) -> None:
        response = JSONResponse(
            {"detail": f"Request body exceeds the deployment limit of {self.max_bytes} bytes"},
            status_code=413,
        )
        await response(scope, receive, send)


def _content_length(headers: Iterable[tuple[bytes, bytes]]) -> int | None:
    for key, value in headers:
        if bytes(key).lower() != b"content-length":
            continue
        try:
            parsed = int(bytes(value))
        except (TypeError, ValueError):
            return None
        return parsed if parsed >= 0 else None
    return None
