"""Server-side enforcement middleware.

Takes any TokenVerifier (mock or real IdP -- see verifiers.py) and enforces
it on every request except the AgentCard discovery endpoint, which must stay
public or nothing can ever discover this agent in the first place.
"""
from __future__ import annotations

from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.requests import Request
from starlette.responses import JSONResponse, Response

from a2a_passport.auth.base import TokenVerifier
from a2a_passport.exceptions import TokenVerificationError

DEFAULT_CARD_PATH = "/.well-known/agent-card.json"

_ERROR_STATUS = {
    "token_expired": 401,
    "invalid_audience": 401,
    "invalid_issuer": 401,
    "key_resolution_failed": 401,
    "invalid_token": 401,
}


class PassportAuthMiddleware(BaseHTTPMiddleware):
    """Validates the inbound Bearer token's signature, audience, and scope
    before letting the request reach A2A route handlers.

    `verifier` is the only thing that changes between a dev deployment
    (SharedSecretVerifier against a mock IdP) and a production one
    (JWKSVerifier against your real IdP) -- this class is identical either way.
    """

    def __init__(
        self,
        app,
        *,
        verifier: TokenVerifier,
        expected_audience: str,
        required_scope: str | None = None,
        card_path: str = DEFAULT_CARD_PATH,
    ):
        super().__init__(app)
        self._verifier = verifier
        self._audience = expected_audience
        self._required_scope = required_scope
        self._card_path = card_path

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        if request.url.path == self._card_path:
            return await call_next(request)

        header = request.headers.get("authorization", "")
        if not header.startswith("Bearer "):
            return JSONResponse({"error": "missing_bearer_token"}, status_code=401)
        token = header[len("Bearer ") :]

        try:
            claims = await self._verifier.verify(token, audience=self._audience)
        except TokenVerificationError as exc:
            status = _ERROR_STATUS.get(exc.error_code, 401)
            return JSONResponse({"error": exc.error_code, "detail": str(exc)}, status_code=status)

        if self._required_scope and not claims.has_scope(self._required_scope):
            return JSONResponse(
                {"error": "insufficient_scope", "required": self._required_scope},
                status_code=403,
            )

        request.state.caller_identity = claims.subject
        request.state.caller_scopes = claims.scopes
        return await call_next(request)
