from __future__ import annotations

from typing import Awaitable, Callable

from aiohttp import web

from ..concurrency import BaseConcurrencyLimiter
from ..response.errors import json_error
from .context import get_request_context

Handler = Callable[[web.Request], Awaitable[web.StreamResponse]]
Middleware = Callable[[web.Request, Handler], Awaitable[web.StreamResponse]]


_CONCURRENCY_BYPASS_PREFIXES = (
    "/healthz", "/v1/health", "/metrics", "/v1/metrics", "/v1/admin/",
)


def build_concurrency_middleware(limiter: BaseConcurrencyLimiter) -> Middleware:
    @web.middleware
    async def concurrency_middleware(request: web.Request, handler: Handler) -> web.StreamResponse:
        if any(request.path.startswith(p) for p in _CONCURRENCY_BYPASS_PREFIXES):
            return await handler(request)
        org_id = get_request_context(request).org_id
        if not limiter.acquire(org_id):
            return json_error(
                "rate_limit",
                "Concurrency limit exceeded",
                status=429,
                context={"scope": "concurrency"},
                headers={"Retry-After": "1"},
            )
        try:
            return await handler(request)
        finally:
            limiter.release(org_id)

    return concurrency_middleware
