"""Bound security-sensitive request bodies before JSON/multipart parsing."""

from starlette.exceptions import HTTPException
from starlette.responses import JSONResponse
from starlette.types import ASGIApp, Message, Receive, Scope, Send


# Encrypted backup limit is 32 MiB. Allow a bounded margin for the multipart
# envelope; count streamed chunks as well as checking the untrusted header.
RESTORE_BODY_LIMIT = 34 * 1024 * 1024
BOOTSTRAP_BODY_LIMIT = 16 * 1024


class InstallationBodyLimitMiddleware:
    def __init__(self, app: ASGIApp) -> None:
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http" or scope.get("method") != "POST":
            await self.app(scope, receive, send)
            return
        path = scope.get("path", "").rstrip("/")
        limit = {
            "/admin/backups/restore": RESTORE_BODY_LIMIT,
            "/admin/backups/export": BOOTSTRAP_BODY_LIMIT,
            "/setup/bootstrap": BOOTSTRAP_BODY_LIMIT,
        }.get(path)
        if limit is None:
            await self.app(scope, receive, send)
            return
        headers = dict(scope.get("headers", []))
        try:
            length = int(headers.get(b"content-length", b"0"))
        except ValueError:
            length = -1
        if length < 0 or length > limit:
            await JSONResponse({"detail": "Request body is too large or has an invalid length."}, status_code=413)(scope, receive, send)
            return
        received = 0

        async def bounded_receive() -> Message:
            nonlocal received
            message = await receive()
            if message["type"] == "http.request":
                received += len(message.get("body", b""))
                if received > limit:
                    raise HTTPException(status_code=413, detail="Request body is too large.")
            return message

        await self.app(scope, bounded_receive, send)