from __future__ import annotations

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


class _RequestBodyTooLarge(Exception):
    """Internal signal used to stop an application from consuming an oversized body."""


class RequestBodyLimitMiddleware:
    """Enforce a request-body ceiling using Content-Length and streamed byte counts."""

    def __init__(self, app: ASGIApp, max_bytes: int) -> None:
        if max_bytes < 1:
            raise ValueError("max_bytes must be positive")
        self.app = app
        self.max_bytes = max_bytes

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http" or scope["method"] in {"GET", "HEAD", "OPTIONS"}:
            await self.app(scope, receive, send)
            return

        headers = {key.lower(): value for key, value in scope.get("headers", [])}
        declared = headers.get(b"content-length")
        if declared is not None:
            try:
                content_length = int(declared)
            except ValueError:
                await self._send_413(send)
                return
            if content_length < 0:
                await self._send_413(send)
                return
            if content_length > self.max_bytes:
                await self._send_413(send)
                return

        received = 0
        exceeded = False
        response_started = False
        response_body_sent = False

        async def limited_receive() -> Message:
            nonlocal received, exceeded
            message = await receive()
            if message["type"] == "http.request":
                received += len(message.get("body", b""))
                if received > self.max_bytes:
                    exceeded = True
                    raise _RequestBodyTooLarge
            return message

        async def guarded_send(message: Message) -> None:
            nonlocal response_started, response_body_sent
            if not exceeded:
                if message["type"] == "http.response.start":
                    response_started = True
                await send(message)
                return

            # Replace any application response after overflow with one clean 413.
            if message["type"] == "http.response.start" and not response_started:
                response_started = True
                body = b'{"detail":"Request body too large"}'
                await send({
                    "type": "http.response.start",
                    "status": 413,
                    "headers": [
                        (b"content-type", b"application/json"),
                        (b"content-length", str(len(body)).encode("ascii")),
                        (b"cache-control", b"no-store"),
                    ],
                })
            elif message["type"] == "http.response.body" and not response_body_sent:
                response_body_sent = True
                await send({
                    "type": "http.response.body",
                    "body": b'{"detail":"Request body too large"}',
                    "more_body": False,
                })

        try:
            await self.app(scope, limited_receive, guarded_send)
        except _RequestBodyTooLarge:
            if not response_started:
                await self._send_413(send)
            elif not response_body_sent:
                response_body_sent = True
                await send({
                    "type": "http.response.body",
                    "body": b'{"detail":"Request body too large"}',
                    "more_body": False,
                })

    async def _send_413(self, send: Send) -> None:
        body = b'{"detail":"Request body too large"}'
        await send({
            "type": "http.response.start",
            "status": 413,
            "headers": [
                (b"content-type", b"application/json"),
                (b"content-length", str(len(body)).encode("ascii")),
                (b"cache-control", b"no-store"),
            ],
        })
        await send({"type": "http.response.body", "body": body, "more_body": False})