from __future__ import annotations

import logging
from typing import Awaitable, Callable

from aiohttp import web

from ..rate_limit import RateLimiter
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]]


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


def build_ratelimit_middleware(
    limiter: RateLimiter | None = None,
) -> Middleware:
    limiter = limiter or RateLimiter()
    log = logging.getLogger("ember.gateway.rate_limit")

    @web.middleware
    async def ratelimit_middleware(request: web.Request, handler: Handler) -> web.StreamResponse:
        if any(request.path.startswith(p) for p in _RATELIMIT_BYPASS_PREFIXES):
            return await handler(request)
        ctx = get_request_context(request)
        if not limiter.check(ctx.org_id, cost=1):
            snapshot = limiter.snapshot(ctx.org_id)
            log.warning(
                "rate_limit_blocked org=%s scope=rpm capacity=%s tokens=%s refill_rate=%s",
                ctx.org_id,
                snapshot.get("capacity"),
                snapshot.get("tokens"),
                snapshot.get("refill_rate"),
            )
            # Include a simple Retry-After hint (seconds)
            return json_error(
                "rate_limit",
                "RPM limit exceeded",
                status=429,
                context={"scope": "rpm"},
                headers={"Retry-After": "1"},
            )
        return await handler(request)

    return ratelimit_middleware
