"""Early, bounded ingress for the batch-check multipart upload."""
from __future__ import annotations

from starlette.types import ASGIApp, Message, Receive, Scope, Send

_TOO_LARGE = b'{"ok":false,"message":"That file is too large to check."}'
# The route still enforces an exact 8 MiB file.  This separate ingress allowance leaves
# bounded room for multipart field names, headers, boundary markers, and the final delimiter.
DEFAULT_MAX_BODY_BYTES = 8_000_000 + 64 * 1024


class _UploadTooLarge(Exception):
    pass


class BatchCheckUploadLimitMiddleware:
    """Stop an oversized batch-check body before Starlette spools multipart parts."""
    def __init__(self, app: ASGIApp, *, max_bytes: int = DEFAULT_MAX_BODY_BYTES) -> None:
        self.app, self.max_bytes = app, max_bytes

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http" or scope.get("method") != "POST" or scope.get("path") != "/api/batch-check":
            await self.app(scope, receive, send)
            return
        declared = dict(scope.get("headers") or ()).get(b"content-length", b"")
        if declared.isdigit() and int(declared) > self.max_bytes:
            await self._send_refusal(send)
            return
        total = 0
        overflow = False
        async def limited_receive() -> Message:
            nonlocal total, overflow
            event = await receive()
            if event["type"] == "http.request":
                total += len(event.get("body", b""))
                if total > self.max_bytes:
                    overflow = True
                    raise _UploadTooLarge
            return event
        async def limited_send(event: Message) -> None:
            # Starlette turns the receive exception into a multipart 400.  Once overflow is
            # known, discard that response and preserve the route's existing JSON contract.
            if not overflow:
                await send(event)
        try:
            await self.app(scope, limited_receive, limited_send)
        except _UploadTooLarge:
            overflow = True
        if overflow:
            await self._send_refusal(send)

    @staticmethod
    async def _send_refusal(send: Send) -> None:
        await send({"type": "http.response.start", "status": 200, "headers": [(b"content-type", b"application/json"), (b"content-length", str(len(_TOO_LARGE)).encode())]})
        await send({"type": "http.response.body", "body": _TOO_LARGE})