from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request

from axon.config import settings

_EXEMPT = {"/health", "/metrics", "/docs", "/openapi.json", "/redoc"}


class APIKeyMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        if not settings.api_key_list:
            return await call_next(request)

        if request.url.path in _EXEMPT:
            return await call_next(request)

        key = request.headers.get("X-API-Key", "")
        if key not in settings.api_key_list:
            return JSONResponse(
                status_code=401,
                content={"detail": "Invalid or missing API key"},
            )

        return await call_next(request)