"""ASGI body limit before JSON parsing, including chunked uploads."""
import asyncio
import os
from threading import Lock
from starlette.responses import JSONResponse


class BodyLimitMiddleware:
    def __init__(self, app, max_bytes=None, timeout=None):
        self.app = app
        self.active = 0
        self.lock = Lock()
        self.max_concurrency = max(1, int(os.getenv("NEEDLE_BODY_CONCURRENCY", "16")))
        self.max_bytes = max_bytes or int(os.getenv("NEEDLE_MAX_REQUEST_BYTES", str(8 * 1024 * 1024)))
        self.timeout = timeout or float(os.getenv("NEEDLE_BODY_TIMEOUT", "30"))

    async def __call__(self, scope, receive, send):
        if scope["type"] != "http" or scope["method"] not in {"POST", "PUT", "PATCH"}:
            return await self.app(scope, receive, send)
        with self.lock:
            admitted = self.active < self.max_concurrency
            if admitted:
                self.active += 1
        if not admitted:
            return await JSONResponse({"detail": "Request concurrency limit"}, 429)(scope, receive, send)
        try:
            await self.handle(scope, receive, send)
        finally:
            with self.lock:
                self.active -= 1

    async def handle(self, scope, receive, send):
        chunks, size = [], 0
        try:
            async with asyncio.timeout(self.timeout):
                while True:
                    message = await receive()
                    if message["type"] == "http.disconnect":
                        return
                    body = message.get("body", b"")
                    size += len(body)
                    if size > self.max_bytes:
                        return await JSONResponse({"detail": "Request body exceeds byte limit"}, 413)(scope, receive, send)
                    chunks.append(body)
                    if not message.get("more_body", False):
                        break
        except TimeoutError:
            return await JSONResponse({"detail": "Request body timed out"}, 408)(scope, receive, send)
        scope["needle.body_bytes"] = size
        body = b"".join(chunks)
        delivered = False
        async def replay():
            nonlocal delivered
            if not delivered:
                delivered = True
                return {"type": "http.request", "body": body, "more_body": False}
            return await receive()
        await self.app(scope, replay, send)