from __future__ import annotations

from json import JSONDecodeError
from typing import Awaitable, Callable

from aiohttp import web
from aiohttp.client_exceptions import ContentTypeError

from ..response.errors import json_error
from ..tokenizer import estimate_chat_prompt_tokens, estimate_responses_prompt_tokens
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]]

_TPM_ENFORCED_PATHS = frozenset({"/v1/chat/completions", "/v1/responses"})


def build_tpm_middleware(limiter: BaseTPMLimiter) -> Middleware:
    @web.middleware
    async def tpm_middleware(request: web.Request, handler: Handler) -> web.StreamResponse:
        if request.method == "POST" and request.path in _TPM_ENFORCED_PATHS:
            try:
                body = request.get("json_body")
                if body is None:
                    body = await request.json()
                    request["json_body"] = body

                # Estimate prompt tokens based on endpoint
                if request.path == "/v1/responses":
                    est_prompt = estimate_responses_prompt_tokens(body or {})
                else:
                    est_prompt = estimate_chat_prompt_tokens(body or {})

                org_id = get_request_context(request).org_id
                ok = limiter.consume_input(org_id, est_prompt)
                if not ok:
                    return json_error(
                        "rate_limit",
                        "TPM limit exceeded (prompt)",
                        status=429,
                        context={"scope": "tpm", "phase": "prompt_precharge"},
                        headers={"Retry-After": "1"},
                    )
                # tag the precharged amount for observability
                request["tpm_precharged"] = int(est_prompt)
            except (JSONDecodeError, ContentTypeError, ValueError, web.HTTPBadRequest):
                # Let the handler surface the parse error.
                return await handler(request)
        return await handler(request)

    return tpm_middleware
