from __future__ import annotations

import json
import logging
from typing import Awaitable, Callable

from aiohttp import web

from ..concurrency import BaseConcurrencyLimiter
from ..config_store import ConfigStore
from ..rate_limit import RateLimiter
from ..tpm_limiter import BaseTPMLimiter
from .context import get_request_context

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


def _resolve_provider(model: str | None) -> str | None:
    """Resolve provider from model string for tier override matching."""
    if not model:
        return None
    # Lightweight import - function is pure string parsing
    from ember.models.providers import resolve_model_id

    provider, _ = resolve_model_id(model)
    return provider if provider != "unknown" else None


def build_limits_applier(
    config: ConfigStore,
    rpm: RateLimiter,
    tpm: BaseTPMLimiter,
    ccy: BaseConcurrencyLimiter,
) -> Middleware:
    log = logging.getLogger("ember.gateway.limits")

    @web.middleware
    async def limits_mw(request: web.Request, handler: Handler) -> web.StreamResponse:
        ctx = get_request_context(request)
        tier = ctx.plan or "T1"
        body = request.get("json_body")
        model: str | None = None
        provider: str | None = None
        if isinstance(body, dict):
            model = body.get("model")
            if isinstance(model, str):
                provider = _resolve_provider(model)
        limits = config.limits_for(tier, provider=provider, model=model)
        log.debug(
            "limits_applied org=%s tier=%s payload=%s",
            ctx.org_id,
            tier,
            json.dumps(
                {
                    "rpm": limits.rpm,
                    "tpm": limits.tpm,
                    "itpm": limits.itpm,
                    "otpm": limits.otpm,
                    "concurrency": limits.concurrency,
                    "provider": provider,
                    "model": model,
                }
            ),
        )
        if limits.rpm is not None:
            rpm.set_limit(ctx.org_id, int(limits.rpm))

        if limits.itpm is not None:
            tpm.set_input_limit(ctx.org_id, int(limits.itpm))

        if limits.otpm is not None:
            tpm.set_output_limit(ctx.org_id, int(limits.otpm))

        if limits.tpm is not None:
            tpm.set_limit(ctx.org_id, int(limits.tpm))

        if limits.concurrency is not None:
            ccy.set_limit(ctx.org_id, int(limits.concurrency))
        return await handler(request)

    return limits_mw
