from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request

from ..errors import error_response
from ..settings import get_settings


class BodyLimitMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        limit = get_settings().body_limit_bytes
        cl = request.headers.get("content-length")
        if cl is not None:
            try:
                if int(cl) > limit:
                    return error_response(request, 413, "payload_too_large", "body too large")
            except ValueError:
                return error_response(request, 400, "validation", "invalid content-length")
            return await call_next(request)

        total = 0
        chunks: list[bytes] = []
        async for chunk in request.stream():
            total += len(chunk)
            if total > limit:
                return error_response(request, 413, "payload_too_large", "body too large")
            chunks.append(chunk)

        async def receive():
            return {"type": "http.request", "body": b"".join(chunks), "more_body": False}

        request._receive = receive  # type: ignore[attr-defined]
        return await call_next(request)