from __future__ import annotations

import json
import logging
from dataclasses import dataclass

from aiohttp import web

from ember._internal.exceptions import (
    ProviderAPIError,
    ProviderRateLimitError,
)
from keiro.gateway.infra.middleware.context import get_request_context

from ..adapters.runtime import reset_request_api_key, set_request_api_key
from ..pipeline.container import Services
from ..domain.errors import RateLimitedError
from .state import METRICS_KEY, SERVICES_KEY

LOG = logging.getLogger(__name__)


@dataclass
class Metrics:
    requests_total: int = 0
    request_errors_total: int = 0

    def as_prometheus(self) -> str:
        return "\n".join(
            [
                f"gateway2_requests_total {self.requests_total}",
                f"gateway2_request_errors_total {self.request_errors_total}",
            ]
        )


@web.middleware
async def metrics_middleware(request: web.Request, handler):
    metrics: Metrics = request.app[METRICS_KEY]
    metrics.requests_total += 1
    try:
        response = await handler(request)
        return response
    except Exception:
        metrics.request_errors_total += 1
        raise


@web.middleware
async def error_middleware(request: web.Request, handler):
    try:
        return await handler(request)
    except RateLimitedError as exc:
        payload = {
            "error": "rate_limit",
            "message": str(exc),
            "scope": exc.scope,
            "phase": exc.phase,
        }
        return web.json_response(payload, status=429)
    except ProviderRateLimitError as exc:
        headers: dict[str, str] = {}
        retry_after = getattr(exc, "retry_after_seconds", None)
        if retry_after is not None:
            headers["Retry-After"] = str(int(retry_after))
        return web.json_response(
            {"error": "upstream_rate_limit", "message": str(exc)},
            status=429,
            headers=headers or None,
        )
    except ProviderAPIError as exc:
        status = exc.http_status
        return web.json_response(
            {"error": "upstream_error", "message": str(exc)},
            status=status,
        )
    except web.HTTPException as exc:
        if exc.status == 404:
            LOG.warning(
                "http_404",
                extra={
                    "method": request.method,
                    "path": request.path,
                    "query": request.rel_url.query_string,
                },
            )
        raise
    except Exception as exc:  # pragma: no cover
        LOG.exception(
            "Unhandled gateway exception",
            extra={"path": request.path},
            exc_info=exc,
        )
        return web.json_response(
            {"error": "internal", "message": "internal server error"},
            status=500,
        )


@web.middleware
async def json_body_middleware(request: web.Request, handler):
    if request.method in {"POST", "PUT", "PATCH"} and request.can_read_body:
        content_type = request.content_type or ""
        if content_type.startswith("application/json"):
            if "json_body" not in request:
                try:
                    request["json_body"] = await request.json()
                except json.JSONDecodeError as exc:
                    raise web.HTTPBadRequest(text="invalid json body") from exc
        elif request.content_length not in {None, 0}:
            raise web.HTTPUnsupportedMediaType(text="expected application/json")
    return await handler(request)


@web.middleware
async def request_id_middleware(request: web.Request, handler):
    services: Services = request.app[SERVICES_KEY]
    request_id = request.headers.get("X-Request-Id") or services.ids.new_request_id()
    request["request_id"] = request_id
    response = await handler(request)
    if isinstance(response, web.StreamResponse) and "X-Request-Id" not in response.headers:
        response.headers["X-Request-Id"] = request_id
    return response


@web.middleware
async def bind_auth_context_middleware(request: web.Request, handler):
    ctx = get_request_context(request)
    api_key_id = ctx.api_key_id if ctx.api_key_id is not None else "anonymous"
    token = set_request_api_key(api_key_id)
    try:
        return await handler(request)
    finally:
        reset_request_api_key(token)


__all__ = [
    "Metrics",
    "metrics_middleware",
    "error_middleware",
    "json_body_middleware",
    "request_id_middleware",
    "bind_auth_context_middleware",
]
