#!/usr/bin/env python3
"""
ASGI middleware for streaming body size limits.
Prevents RAM blow-ups from chunked uploads without Content-Length.
"""
from starlette.types import ASGIApp, Receive, Scope, Send
from starlette.responses import JSONResponse


class StreamingBodyLimitMiddleware:
    """
    ASGI middleware that enforces MAX_BODY limit on streaming request bodies.

    Unlike Content-Length header checks, this actively counts bytes as they
    arrive in chunks and returns 413 immediately when limit exceeded, without
    buffering the entire payload into RAM.
    """
    def __init__(self, app: ASGIApp, max_bytes: int):
        self.app = app
        self.max_bytes = max_bytes

    async def __call__(self, scope: Scope, receive: Receive, send: Send):
        if scope["type"] != "http" or scope["method"] not in ("POST", "PUT", "PATCH"):
            return await self.app(scope, receive, send)

        consumed = 0
        done = False

        async def limited_receive():
            nonlocal consumed, done
            if done:
                return {"type": "http.request", "body": b"", "more_body": False}

            message = await receive()
            if message["type"] != "http.request":
                return message

            body = message.get("body", b"")
            consumed += len(body)

            if consumed > self.max_bytes:
                done = True
                resp = JSONResponse(
                    {"detail": "Payload too large", "max_bytes": self.max_bytes},
                    status_code=413
                )
                await resp(scope, receive, send)
                return {"type": "http.request", "body": b"", "more_body": False}

            return message

        await self.app(scope, limited_receive, send)
