# -------------------------------------------------------------------
# This source file is available under the terms of the
# Pimcore Open Core License (POCL)
# Full copyright and license information is available in
# LICENSE.md which is distributed with this source code.
#
#  @copyright  Copyright (c) Pimcore GmbH (https://www.pimcore.com)
#  @license    Pimcore Open Core License (POCL)
# -------------------------------------------------------------------

"""ASGI middleware guarding against oversized request bodies."""

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


class MaxBodySizeMiddleware:
    """Reject request bodies larger than ``max_body_size`` bytes with ``413`` before a handler
    allocates memory for them.

    A batch embedding request carries base64-encoded images, so a single request can be large;
    without a ceiling a caller could post a body big enough to exhaust memory. A declared
    ``Content-Length`` over the limit is rejected before the body is read; the streamed bytes are
    also counted, so an absent or understated ``Content-Length`` (e.g. chunked transfer) still
    trips the limit.

    The two paths report the same 413 but cannot report it the same way. ``add_middleware`` puts
    this middleware *outside* Starlette's ``ExceptionMiddleware``, so an exception raised in
    ``__call__`` itself would surface as a 500 — hence the ``Content-Length`` check writes its
    response directly. ``limited_receive`` runs the other way round: it is called from the route
    handler, i.e. inside ``ExceptionMiddleware``, and its exception passes through FastAPI's body
    parser, which converts everything except an ``HTTPException`` into ``400 "There was an error
    parsing the body"``.
    """

    def __init__(self, app: ASGIApp, max_body_size: int) -> None:
        self._app = app
        self._max_body_size = max_body_size

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http":
            await self._app(scope, receive, send)
            return

        for name, value in scope.get("headers", []):
            if name == b"content-length" and value.isdigit() and int(value) > self._max_body_size:
                # Raising here would escape ExceptionMiddleware and become a 500; write the
                # response ourselves. Nothing has been sent yet, so this is the whole response.
                await self._send_413(send)
                return

        received = 0

        async def limited_receive() -> Message:
            nonlocal received
            message = await receive()
            if message["type"] == "http.request":
                received += len(message.get("body", b""))
                if received > self._max_body_size:
                    # Must be an HTTPException: FastAPI re-raises those out of its body parser and
                    # ExceptionMiddleware renders them, while any other exception is swallowed
                    # into a 400 "There was an error parsing the body".
                    raise HTTPException(status_code=413, detail="Request body too large")
            return message

        await self._app(scope, limited_receive, send)

    @staticmethod
    async def _send_413(send: Send) -> None:
        await send(
            {
                "type": "http.response.start",
                "status": 413,
                "headers": [(b"content-type", b"application/json")],
            }
        )
        await send(
            {
                "type": "http.response.body",
                "body": b'{"detail":"Request body too large"}',
            }
        )