from __future__ import annotations

import logging
from typing import Awaitable, Callable

from aiohttp import web

from ember._internal.exceptions import (
    ProviderAPIError,
    ProviderAuthenticationError,
    ProviderRateLimitError,
)

from ..errors import GatewayError
from ..response.errors import json_error

Handler = Callable[[web.Request], Awaitable[web.StreamResponse]]
_LOG = logging.getLogger(__name__)


def _provider_error_code(pe: ProviderAPIError) -> str:
    """Map ProviderAPIError subclass to response code string."""
    if isinstance(pe, ProviderRateLimitError):
        return "provider_rate_limit"
    if isinstance(pe, ProviderAuthenticationError):
        return "provider_auth"
    return "provider_error"


def _provider_error_status(pe: ProviderAPIError) -> int:
    """Map ProviderAPIError to HTTP status code."""
    if type(pe) is not ProviderAPIError:
        return pe.http_status
    if "http_status" in vars(pe):
        return pe.http_status
    error_type = pe.context.get("error_type", "")
    return {"rate_limit": 429, "authentication": 401}.get(error_type, 502)


@web.middleware
async def error_middleware(request: web.Request, handler: Handler) -> web.StreamResponse:
    try:
        return await handler(request)
    except GatewayError as ge:
        return json_error(
            ge.code,
            str(ge),
            status=getattr(ge, "http_status", 500),
            context=getattr(ge, "context", {}),
        )
    except ProviderAPIError as pe:
        status = _provider_error_status(pe)
        code = _provider_error_code(pe)
        headers = None
        if isinstance(pe, ProviderRateLimitError) and pe.retry_after_seconds:
            headers = {"Retry-After": str(pe.retry_after_seconds)}
        _LOG.warning("Provider error: %s (status=%d)", pe, status)
        return json_error(code, str(pe), status=status, context=pe.context, headers=headers)
    except web.HTTPException as he:
        return json_error(he.reason, he.text or "", status=he.status)
    except Exception:
        _LOG.exception("Unhandled gateway error")
        return json_error("internal", "Internal server error", status=500)
