"""Server-side token verification, pluggable per identity provider.

Both classes implement the same TokenVerifier protocol (auth/base.py) so
PassportAuthMiddleware never changes when you swap one for the other --
that swap is the entire point of this module. Start with
SharedSecretVerifier against a mock IdP in dev/CI, deploy JWKSVerifier
against your real IdP's JWKS endpoint in production. No other code changes.
"""
from __future__ import annotations

import jwt

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


def _claims_to_verified(claims: dict, audience: str) -> VerifiedClaims:
    scope_claim = claims.get("scope", "")
    scopes = (
        frozenset(scope_claim.split())
        if isinstance(scope_claim, str)
        else frozenset(scope_claim)
    )
    return VerifiedClaims(
        subject=claims.get("sub", ""),
        audience=audience,
        scopes=scopes,
        raw=claims,
    )


class SharedSecretVerifier:
    """HS256 verification against a static shared secret.

    Dev/CI/test only: a shared symmetric secret means anything that can
    verify a token can also forge one. Use this against a2a_passport's mock
    IdP (a2a_passport.testing.mock_idp) or an equivalent test double --
    never against production traffic. Use JWKSVerifier for that.
    """

    def __init__(self, secret: str, *, algorithm: str = "HS256"):
        self._secret = secret
        self._algorithm = algorithm

    async def verify(self, token: str, *, audience: str) -> VerifiedClaims:
        try:
            claims = jwt.decode(
                token, self._secret, algorithms=[self._algorithm], audience=audience
            )
        except jwt.ExpiredSignatureError as exc:
            raise TokenVerificationError(str(exc), error_code="token_expired") from exc
        except jwt.InvalidAudienceError as exc:
            raise TokenVerificationError(str(exc), error_code="invalid_audience") from exc
        except jwt.InvalidTokenError as exc:
            raise TokenVerificationError(str(exc), error_code="invalid_token") from exc
        return _claims_to_verified(claims, audience)


class JWKSVerifier:
    """RS256/ES256 verification against a real IdP's published JWKS.

    Works unmodified against Auth0, Keycloak, Okta, Azure AD, or any OIDC
    provider that exposes a standard `jwks_uri` -- point it at
    `https://<your-tenant>/.well-known/jwks.json` (or your IdP's equivalent
    path) and pass the `issuer` your IdP puts in the `iss` claim. Keys are
    fetched and cached by PyJWT's PyJWKClient, including rotation.
    """

    def __init__(
        self,
        *,
        jwks_uri: str,
        issuer: str | None = None,
        algorithms: list[str] | None = None,
    ):
        self._jwks_client = jwt.PyJWKClient(jwks_uri)
        self._issuer = issuer
        self._algorithms = algorithms or ["RS256"]

    async def verify(self, token: str, *, audience: str) -> VerifiedClaims:
        try:
            signing_key = self._jwks_client.get_signing_key_from_jwt(token)
            claims = jwt.decode(
                token,
                signing_key.key,
                algorithms=self._algorithms,
                audience=audience,
                issuer=self._issuer,
            )
        except jwt.ExpiredSignatureError as exc:
            raise TokenVerificationError(str(exc), error_code="token_expired") from exc
        except jwt.InvalidAudienceError as exc:
            raise TokenVerificationError(str(exc), error_code="invalid_audience") from exc
        except jwt.InvalidIssuerError as exc:
            raise TokenVerificationError(str(exc), error_code="invalid_issuer") from exc
        except jwt.PyJWKClientError as exc:
            raise TokenVerificationError(str(exc), error_code="key_resolution_failed") from exc
        except jwt.InvalidTokenError as exc:
            raise TokenVerificationError(str(exc), error_code="invalid_token") from exc
        return _claims_to_verified(claims, audience)


__all__ = ["SharedSecretVerifier", "JWKSVerifier", "TokenVerifier"]
