"""限制请求体大小：声明的 Content-Length 超限直接 413，没声明（分块）或声明得不实的，边读边数

FastAPI 会先把整个请求体收完再调用接口函数，超大的请求体已经占满了带宽和内存
和数据文件上传的 UploadSizeLimit 同一出发点；限制随路径而不同，所以用规则表，第一条命中的生效
"""

import re
from dataclasses import dataclass

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

from app.common.response import error_content


class BodyTooLarge(HTTPException):
    """边读边数时超限；message 由全局 HTTP 异常处理放进信封，和声明超限的响应文案一致"""

    def __init__(self, message: str):
        super().__init__(413)
        self.message = message


@dataclass(frozen=True)
class BodyLimitRule:
    methods: frozenset[str]
    path: re.Pattern[str]
    max_bytes: int
    message: str


class BodyLimit:
    def __init__(self, app: ASGIApp, rules: list[BodyLimitRule]):
        self.app = app
        self.rules = rules

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        rule = self._match(scope) if scope["type"] == "http" else None
        if rule is None:
            await self.app(scope, receive, send)
            return
        length = dict(scope["headers"]).get(b"content-length", b"")
        if length.isdigit() and int(length) > rule.max_bytes:
            response = JSONResponse(error_content(413, rule.message), 413)
            await response(scope, receive, send)
            return
        received = 0

        async def counted() -> Message:
            nonlocal received
            message = await receive()
            if message["type"] == "http.request":
                received += len(message.get("body", b""))
                if received > rule.max_bytes:
                    # 抛出的异常由应用里统一的 HTTP 异常处理转成 413 信封
                    raise BodyTooLarge(rule.message)
            return message

        await self.app(scope, counted, send)

    def _match(self, scope: Scope) -> BodyLimitRule | None:
        return next(
            (
                rule
                for rule in self.rules
                if scope["method"] in rule.methods
                and rule.path.fullmatch(scope["path"])
            ),
            None,
        )