from collections.abc import Callable

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


SESSION_BODY_LIMIT = 16 * 1024
API_MUTATION_BODY_LIMIT = 64 * 1024
ISSUE_ATTACHMENT_BODY_LIMIT = 2 * 1024 * 1024 + 64 * 1024
LEGACY_JSON_ATTACHMENT_BODY_LIMIT = 3 * 1024 * 1024
UNFILED_DRAFT_SYNC_BODY_LIMIT = 17 * 1024 * 1024
MUTATION_METHODS = frozenset({"POST", "PUT", "PATCH"})


def request_body_limit(method: str, path: str) -> int | None:
    """Return the admission limit for request bodies that FastAPI will parse."""
    normalized_method = method.upper()
    if normalized_method == "POST" and path == "/api/v1/session":
        return SESSION_BODY_LIMIT
    if normalized_method == "PUT" and path == "/api/v1/unfiled-drafts":
        return UNFILED_DRAFT_SYNC_BODY_LIMIT
    if (
        normalized_method == "POST"
        and (
            (path.startswith("/api/v1/repos/") and ("/issues/" in path or "/pulls/" in path))
            or path.startswith("/api/v1/notifications/")
        )
        and path.endswith("/attachments")
    ):
        return ISSUE_ATTACHMENT_BODY_LIMIT
    if normalized_method in MUTATION_METHODS and path.startswith("/api/v1/"):
        return API_MUTATION_BODY_LIMIT
    return None


class RequestBodyLimitMiddleware:
    """Reject oversized API bodies before framework parsing or endpoint work."""

    def __init__(
        self,
        app: ASGIApp,
        limit_for: Callable[[str, str], int | None] = request_body_limit,
    ) -> None:
        self.app = app
        self.limit_for = limit_for

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

        path = scope.get("path", "")
        root_path = scope.get("root_path", "").rstrip("/")
        if root_path and (path == root_path or path.startswith(root_path + "/")):
            path = path[len(root_path) :] or "/"
        limit = self.limit_for(scope.get("method", "GET"), path)
        if limit is None:
            await self.app(scope, receive, send)
            return
        if (
            limit == ISSUE_ATTACHMENT_BODY_LIMIT
            and self._header(scope, b"content-type").startswith(b"application/json")
        ):
            limit = LEGACY_JSON_ATTACHMENT_BODY_LIMIT

        declared_length = self._content_length(scope)
        if declared_length is not None and declared_length > limit:
            await self._reject(send)
            return

        messages: list[Message] = []
        received = 0
        while True:
            message = await receive()
            messages.append(message)
            if message["type"] != "http.request":
                break
            received += len(message.get("body", b""))
            if received > limit:
                await self._reject(send)
                return
            if not message.get("more_body", False):
                break

        async def replay() -> Message:
            if messages:
                return messages.pop(0)
            return {"type": "http.request", "body": b"", "more_body": False}

        await self.app(scope, replay, send)

    @staticmethod
    def _header(scope: Scope, wanted: bytes) -> bytes:
        for name, value in scope.get("headers", []):
            if name.lower() == wanted:
                return value.lower()
        return b""

    @staticmethod
    def _content_length(scope: Scope) -> int | None:
        for name, value in scope.get("headers", []):
            if name.lower() != b"content-length":
                continue
            try:
                parsed = int(value)
            except ValueError:
                return None
            return max(0, parsed)
        return None

    @staticmethod
    async def _reject(send: Send) -> None:
        body = b'{"detail":"Request body too large"}'
        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.start", "status": 413, "headers": headers})
        await send({"type": "http.response.body", "body": body})