import time
from collections.abc import AsyncGenerator
from typing import Any, cast
from unittest.mock import MagicMock, patch

import pytest
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric.ed448 import Ed448PrivateKey
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from joserfc import jwk as jose_jwk
from joserfc import jwt
from joserfc.jws import JWSRegistry
from joserfc.registry import HeaderParameter
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser
from starlette.requests import Request

from fastmcp import FastMCP
from fastmcp.server.auth.providers.jwt import JWKData, JWKSData, JWTVerifier, RSAKeyPair
from fastmcp.server.dependencies import (
    FastMCPRequestContext,
    fastmcp_request_ctx,
    get_access_token,
)
from fastmcp.utilities.tests import run_server_async
from tests.utilities.httpx2_mock import HTTPXMock

# Standard public IP used for DNS mocking in tests
TEST_PUBLIC_IP = "93.184.216.34"


class SymmetricKeyHelper:
    """Helper class for generating symmetric key JWT tokens for testing."""

    def __init__(self, secret: str):
        """Initialize with a secret key."""
        self.secret = secret

    def create_token(
        self,
        subject: str = "fastmcp-user",
        issuer: str = "https://fastmcp.example.com",
        audience: str | list[str] | None = None,
        scopes: list[str] | None = None,
        expires_in_seconds: int = 3600,
        additional_claims: dict[str, Any] | None = None,
        algorithm: str = "HS256",
    ) -> str:
        """
        Generate a test JWT token using symmetric key for testing purposes.

        Args:
            subject: Subject claim (usually user ID)
            issuer: Issuer claim
            audience: Audience claim - can be a string or list of strings (optional)
            scopes: List of scopes to include
            expires_in_seconds: Token expiration time in seconds
            additional_claims: Any additional claims to include
            algorithm: JWT signing algorithm (HS256, HS384, or HS512)
        """
        # Create header
        header = {"alg": algorithm}

        # Create payload
        payload: dict[str, str | int | list[str]] = {
            "sub": subject,
            "iss": issuer,
            "iat": int(time.time()),
            "exp": int(time.time()) + expires_in_seconds,
        }

        if audience:
            payload["aud"] = audience

        if scopes:
            payload["scope"] = " ".join(scopes)

        if additional_claims:
            payload.update(additional_claims)

        # Create JWT
        signing_key = jose_jwk.import_key(self.secret, "oct")
        token = jwt.encode(header, payload, signing_key, algorithms=[algorithm])

        return token


def create_okp_key_pair(
    private_key: Ed25519PrivateKey | Ed448PrivateKey,
) -> tuple[str, str]:
    """Serialize an EdDSA key pair as PEM strings."""
    private_pem = private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption(),
    ).decode()
    public_pem = (
        private_key.public_key()
        .public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo,
        )
        .decode()
    )
    return private_pem, public_pem


def create_okp_token(
    private_key: str,
    algorithm: str,
    *,
    kid: str | None = None,
) -> str:
    """Create a JWT signed by an OKP key."""
    header = {"alg": algorithm}
    if kid is not None:
        header["kid"] = kid
    return jwt.encode(
        header,
        {
            "sub": "test-user",
            "iss": "https://test.example.com",
            "aud": "https://api.example.com",
            "exp": int(time.time()) + 3600,
        },
        jose_jwk.import_key(private_key, "OKP"),
        algorithms=[algorithm],
    )


@pytest.fixture(scope="module")
def symmetric_key_helper() -> SymmetricKeyHelper:
    """Generate a symmetric key helper for testing."""
    return SymmetricKeyHelper("test-secret-key-for-hmac-signing")


@pytest.fixture(scope="module")
def bearer_token(rsa_key_pair: RSAKeyPair) -> str:
    return rsa_key_pair.create_token(
        subject="test-user",
        issuer="https://test.example.com",
        audience="https://api.example.com",
    )


@pytest.fixture
def bearer_provider(rsa_key_pair: RSAKeyPair) -> JWTVerifier:
    return JWTVerifier(
        public_key=rsa_key_pair.public_key,
        issuer="https://test.example.com",
        audience="https://api.example.com",
    )


@pytest.fixture
def symmetric_provider(symmetric_key_helper: SymmetricKeyHelper) -> JWTVerifier:
    """Create JWTVerifier configured for symmetric key verification."""
    return JWTVerifier(
        public_key=symmetric_key_helper.secret,
        issuer="https://test.example.com",
        audience="https://api.example.com",
        algorithm="HS256",
    )


def create_mcp_server(
    public_key: str,
    auth_kwargs: dict[str, Any] | None = None,
) -> FastMCP:
    mcp = FastMCP(
        auth=JWTVerifier(
            public_key=public_key,
            **auth_kwargs or {},
        )
    )

    @mcp.tool
    def add(a: int, b: int) -> int:
        return a + b

    return mcp


@pytest.fixture
async def mcp_server_url(rsa_key_pair: RSAKeyPair) -> AsyncGenerator[str, None]:
    server = create_mcp_server(
        public_key=rsa_key_pair.public_key,
        auth_kwargs=dict(
            issuer="https://test.example.com",
            audience="https://api.example.com",
        ),
    )
    async with run_server_async(server, transport="http") as url:
        yield url


class TestRSAKeyPair:
    def test_generate_key_pair(self):
        """Test RSA key pair generation."""
        key_pair = RSAKeyPair.generate()

        assert key_pair.private_key is not None
        assert key_pair.public_key is not None

        # Check that keys are in PEM format
        private_pem = key_pair.private_key.get_secret_value()
        public_pem = key_pair.public_key

        assert "-----BEGIN PRIVATE KEY-----" in private_pem
        assert "-----END PRIVATE KEY-----" in private_pem
        assert "-----BEGIN PUBLIC KEY-----" in public_pem
        assert "-----END PUBLIC KEY-----" in public_pem

    def test_create_basic_token(self, rsa_key_pair: RSAKeyPair):
        """Test basic token creation."""
        token = rsa_key_pair.create_token(
            subject="test-user",
            issuer="https://test.example.com",
        )

        assert isinstance(token, str)
        assert len(token.split(".")) == 3  # JWT has 3 parts

    def test_create_token_with_scopes(self, rsa_key_pair: RSAKeyPair):
        """Test token creation with scopes."""
        token = rsa_key_pair.create_token(
            subject="test-user",
            issuer="https://test.example.com",
            scopes=["read", "write"],
        )

        assert isinstance(token, str)
        # We'll validate the scopes in the BearerToken tests


class TestJWTVerifierHeaders:
    async def test_non_critical_private_header_is_allowed(
        self, rsa_key_pair: RSAKeyPair
    ):
        signing_key = jose_jwk.import_key(
            rsa_key_pair.private_key.get_secret_value(),
            "RSA",
        )
        token = jwt.encode(
            {
                "alg": "RS256",
                "cat": "cl_example",
            },
            {
                "sub": "test-user",
                "iss": "https://test.example.com",
                "exp": int(time.time()) + 3600,
            },
            signing_key,
            algorithms=["RS256"],
            registry=JWSRegistry(strict_check_header=False),
        )
        verifier = JWTVerifier(
            public_key=rsa_key_pair.public_key,
            issuer="https://test.example.com",
        )

        access_token = await verifier.verify_token(token)

        assert access_token is not None
        assert access_token.client_id == "test-user"

    async def test_critical_private_header_is_rejected(self, rsa_key_pair: RSAKeyPair):
        signing_key = jose_jwk.import_key(
            rsa_key_pair.private_key.get_secret_value(),
            "RSA",
        )
        token = jwt.encode(
            {
                "alg": "RS256",
                "crit": ["cat"],
                "cat": "cl_example",
            },
            {
                "sub": "test-user",
                "iss": "https://test.example.com",
                "exp": int(time.time()) + 3600,
            },
            signing_key,
            algorithms=["RS256"],
            registry=JWSRegistry(
                header_registry={
                    "cat": HeaderParameter("Custom private header", "str")
                },
                strict_check_header=False,
            ),
        )
        verifier = JWTVerifier(
            public_key=rsa_key_pair.public_key,
            issuer="https://test.example.com",
        )

        access_token = await verifier.verify_token(token)

        assert access_token is None


class TestSymmetricKeyJWT:
    """Tests for JWT verification using symmetric keys (HMAC algorithms)."""

    def test_initialization_with_symmetric_key(
        self, symmetric_key_helper: SymmetricKeyHelper
    ):
        """Test JWTVerifier initialization with symmetric key."""
        provider = JWTVerifier(
            public_key=symmetric_key_helper.secret,
            issuer="https://test.example.com",
            algorithm="HS256",
        )

        assert provider.issuer == "https://test.example.com"
        assert provider.public_key == symmetric_key_helper.secret
        assert provider.algorithm == "HS256"
        assert provider.jwks_uri is None

    def test_initialization_rejects_hs_algorithm_with_jwks_uri(self):
        """Test that HMAC algorithms cannot be used with JWKS URI."""
        with pytest.raises(ValueError, match="cannot be used with jwks_uri"):
            JWTVerifier(
                jwks_uri="https://test.example.com/.well-known/jwks.json",
                issuer="https://test.example.com",
                algorithm="HS256",
            )

    def test_initialization_with_different_symmetric_algorithms(
        self, symmetric_key_helper: SymmetricKeyHelper
    ):
        """Test JWTVerifier initialization with different HMAC algorithms."""
        algorithms = ["HS256", "HS384", "HS512"]

        for algorithm in algorithms:
            provider = JWTVerifier(
                public_key=symmetric_key_helper.secret,
                issuer="https://test.example.com",
                algorithm=algorithm,
            )
            assert provider.algorithm == algorithm

    def test_symmetric_algorithm_rejects_jwks_uri(self):
        """HS* algorithms must not be configured with JWKS/public key endpoints."""
        with pytest.raises(ValueError, match="cannot be used with jwks_uri"):
            JWTVerifier(
                jwks_uri="https://test.example.com/.well-known/jwks.json",
                issuer="https://test.example.com",
                algorithm="HS256",
            )

    def test_symmetric_algorithm_rejects_pem_public_key(self, rsa_key_pair: RSAKeyPair):
        """HS* algorithms must use a shared secret, not PEM public key material."""
        with pytest.raises(ValueError, match="require a shared secret"):
            JWTVerifier(
                public_key=rsa_key_pair.public_key,
                issuer="https://test.example.com",
                algorithm="HS256",
            )

    def test_symmetric_algorithm_accepts_bytes_secret(self):
        """HS* algorithms accept bytes secrets without TypeError."""
        verifier = JWTVerifier(
            public_key=b"secret",
            algorithm="HS256",
        )
        assert verifier.algorithm == "HS256"

    async def test_valid_symmetric_token_validation(
        self, symmetric_key_helper: SymmetricKeyHelper, symmetric_provider: JWTVerifier
    ):
        """Test validation of a valid token signed with symmetric key."""
        token = symmetric_key_helper.create_token(
            subject="test-user",
            issuer="https://test.example.com",
            audience="https://api.example.com",
            scopes=["read", "write"],
            algorithm="HS256",
        )

        access_token = await symmetric_provider.load_access_token(token)

        assert access_token is not None
        assert access_token.client_id == "test-user"
        assert "read" in access_token.scopes
        assert "write" in access_token.scopes
        assert access_token.expires_at is not None
        assert access_token.subject == "test-user"

    async def test_symmetric_token_with_different_algorithms(
        self, symmetric_key_helper: SymmetricKeyHelper
    ):
        """Test that different HMAC algorithms work correctly."""
        algorithms = ["HS256", "HS384", "HS512"]

        for algorithm in algorithms:
            provider = JWTVerifier(
                public_key=symmetric_key_helper.secret,
                issuer="https://test.example.com",
                algorithm=algorithm,
            )

            token = symmetric_key_helper.create_token(
                subject="test-user",
                issuer="https://test.example.com",
                algorithm=algorithm,
            )

            access_token = await provider.load_access_token(token)
            assert access_token is not None
            assert access_token.client_id == "test-user"

    async def test_symmetric_token_issuer_validation(
        self, symmetric_key_helper: SymmetricKeyHelper, symmetric_provider: JWTVerifier
    ):
        """Test issuer validation with symmetric key tokens."""
        # Valid issuer
        valid_token = symmetric_key_helper.create_token(
            subject="test-user",
            issuer="https://test.example.com",
            audience="https://api.example.com",
        )
        access_token = await symmetric_provider.load_access_token(valid_token)
        assert access_token is not None

        # Invalid issuer
        invalid_token = symmetric_key_helper.create_token(
            subject="test-user",
            issuer="https://evil.example.com",
            audience="https://api.example.com",
        )
        access_token = await symmetric_provider.load_access_token(invalid_token)
        assert access_token is None

    async def test_symmetric_token_audience_validation(
        self, symmetric_key_helper: SymmetricKeyHelper, symmetric_provider: JWTVerifier
    ):
        """Test audience validation with symmetric key tokens."""
        # Valid audience
        valid_token = symmetric_key_helper.create_token(
            subject="test-user",
            issuer="https://test.example.com",
            audience="https://api.example.com",
        )
        access_token = await symmetric_provider.load_access_token(valid_token)
        assert access_token is not None

        # Invalid audience
        invalid_token = symmetric_key_helper.create_token(
            subject="test-user",
            issuer="https://test.example.com",
            audience="https://wrong-api.example.com",
        )
        access_token = await symmetric_provider.load_access_token(invalid_token)
        assert access_token is None

    async def test_symmetric_token_scope_extraction(
        self, symmetric_key_helper: SymmetricKeyHelper, symmetric_provider: JWTVerifier
    ):
        """Test scope extraction from symmetric key tokens."""
        token = symmetric_key_helper.create_token(
            subject="test-user",
            issuer="https://test.example.com",
            audience="https://api.example.com",
            scopes=["read", "write", "admin"],
        )

        access_token = await symmetric_provider.load_access_token(token)
        assert access_token is not None
        assert set(access_token.scopes) == {"read", "write", "admin"}

    async def test_symmetric_token_expiration(
        self, symmetric_key_helper: SymmetricKeyHelper, symmetric_provider: JWTVerifier
    ):
        """Test expiration validation with symmetric key tokens."""
        # Valid token
        valid_token = symmetric_key_helper.create_token(
            subject="test-user",
            issuer="https://test.example.com",
            audience="https://api.example.com",
            expires_in_seconds=3600,  # 1 hour from now
        )
        access_token = await symmetric_provider.load_access_token(valid_token)
        assert access_token is not None

        # Expired token
        expired_token = symmetric_key_helper.create_token(
            subject="test-user",
            issuer="https://test.example.com",
            audience="https://api.example.com",
            expires_in_seconds=-3600,  # 1 hour ago
        )
        access_token = await symmetric_provider.load_access_token(expired_token)
        assert access_token is None

    async def test_symmetric_token_invalid_signature(
        self, symmetric_key_helper: SymmetricKeyHelper, symmetric_provider: JWTVerifier
    ):
        """Test rejection of tokens with invalid signatures."""
        # Create a token with a different secret
        other_helper = SymmetricKeyHelper("different-secret-key")
        token = other_helper.create_token(
            subject="test-user",
            issuer="https://test.example.com",
            audience="https://api.example.com",
        )

        access_token = await symmetric_provider.load_access_token(token)
        assert access_token is None

    async def test_symmetric_token_algorithm_mismatch(
        self, symmetric_key_helper: SymmetricKeyHelper
    ):
        """Test that tokens with mismatched algorithms are rejected."""
        # Create provider expecting HS256
        provider = JWTVerifier(
            public_key=symmetric_key_helper.secret,
            issuer="https://test.example.com",
            algorithm="HS256",
        )

        # Create token with HS512
        token = symmetric_key_helper.create_token(
            subject="test-user",
            issuer="https://test.example.com",
            algorithm="HS512",
        )

        # Should fail because provider expects HS256
        access_token = await provider.load_access_token(token)
        assert access_token is None


class TestEdDSAJWT:
    """Tests for JWT verification using Edwards-curve keys."""

    @pytest.mark.parametrize("algorithm", ["Ed25519", "Ed448"])
    async def test_static_public_key(self, algorithm: str):
        """Fully specified EdDSA algorithms verify with a static public key."""
        if algorithm == "Ed25519":
            private_key = Ed25519PrivateKey.generate()
        else:
            private_key = Ed448PrivateKey.generate()
        private_pem, public_pem = create_okp_key_pair(private_key)
        verifier = JWTVerifier(
            public_key=public_pem,
            issuer="https://test.example.com",
            audience="https://api.example.com",
            algorithm=algorithm,
        )

        access_token = await verifier.load_access_token(
            create_okp_token(private_pem, algorithm)
        )

        assert access_token is not None
        assert access_token.client_id == "test-user"

    @pytest.mark.filterwarnings(
        "ignore:EdDSA is deprecated via RFC 9864:joserfc.errors.SecurityWarning"
    )
    async def test_legacy_eddsa_jwks(
        self,
        httpx_mock: HTTPXMock,
    ):
        """Legacy EdDSA tokens verify against an Ed25519 JWKS entry."""
        private_pem, public_pem = create_okp_key_pair(Ed25519PrivateKey.generate())
        public_jwk = jose_jwk.import_key(public_pem, "OKP").as_dict()
        public_jwk.update(kid="ed25519-key", alg="EdDSA", use="sig")
        httpx_mock.add_response(json={"keys": [public_jwk]})
        verifier = JWTVerifier(
            jwks_uri="https://test.example.com/.well-known/jwks.json",
            issuer="https://test.example.com",
            audience="https://api.example.com",
            algorithm="EdDSA",
        )

        access_token = await verifier.load_access_token(
            create_okp_token(private_pem, "EdDSA", kid="ed25519-key")
        )

        assert access_token is not None
        assert access_token.client_id == "test-user"


def _create_token_without_sub(
    rsa_key_pair: RSAKeyPair,
    *,
    issuer: str = "https://test.example.com",
    audience: str | None = None,
) -> str:
    """Sign a JWT with no 'sub' claim, to exercise the missing-subject path."""
    payload: dict[str, str | int] = {
        "iss": issuer,
        "iat": int(time.time()),
        "exp": int(time.time()) + 3600,
    }
    if audience:
        payload["aud"] = audience
    signing_key = jose_jwk.import_key(
        rsa_key_pair.private_key.get_secret_value(), "RSA"
    )
    return jwt.encode({"alg": "RS256"}, payload, signing_key, algorithms=["RS256"])


class TestJWTAccessTokenSubject:
    """Regression tests for issue #4266.

    ``get_access_token().subject`` was always ``None``, even when the
    authorization server supplied a ``sub`` claim, because neither the
    JWT verifier nor the dependency-layer conversion carried it through.
    """

    async def test_subject_populated_from_sub_claim(
        self, bearer_provider: JWTVerifier, rsa_key_pair: RSAKeyPair
    ):
        """A JWT bearing a 'sub' claim results in a populated subject."""
        token = rsa_key_pair.create_token(
            subject="user-42",
            issuer="https://test.example.com",
            audience="https://api.example.com",
        )

        access_token = await bearer_provider.load_access_token(token)

        assert access_token is not None
        assert access_token.subject == "user-42"

    async def test_subject_is_none_without_sub_claim(
        self, bearer_provider: JWTVerifier, rsa_key_pair: RSAKeyPair
    ):
        """A JWT with no 'sub' claim yields subject=None instead of crashing."""
        token = _create_token_without_sub(
            rsa_key_pair, audience="https://api.example.com"
        )

        access_token = await bearer_provider.load_access_token(token)

        assert access_token is not None
        assert access_token.subject is None

    async def test_get_access_token_returns_subject_for_realistic_jwt_request(
        self, bearer_provider: JWTVerifier, rsa_key_pair: RSAKeyPair
    ):
        """End-to-end path a real authenticated request actually takes.

        Mirrors ``mcp.server.auth.middleware.bearer_auth.BearerAuthBackend``,
        which wraps whatever the ``TokenVerifier`` returns directly into
        ``AuthenticatedUser`` and stores it on ``request.scope["user"]``.
        ``get_access_token()`` reads that object first, before ever reaching
        the dependency-layer conversion block - so this is the path that
        must carry ``subject`` through for real JWT-authenticated requests.
        """
        token = rsa_key_pair.create_token(
            subject="user-42",
            issuer="https://test.example.com",
            audience="https://api.example.com",
        )
        access_token = await bearer_provider.load_access_token(token)
        assert access_token is not None

        user = AuthenticatedUser(access_token)
        request = Request({"type": "http", "user": user, "auth": MagicMock()})

        ctx_token = fastmcp_request_ctx.set(
            FastMCPRequestContext(
                session=MagicMock(),
                request_id="0",
                meta=None,
                request=request,
                protocol_version="2025-06-18",
                close_sse_stream=None,
                lifespan_context=MagicMock(),
                _srctx=MagicMock(meta=None),
            )
        )
        try:
            result = get_access_token()
        finally:
            fastmcp_request_ctx.reset(ctx_token)

        assert result is not None
        assert result.subject == "user-42"


class TestBearerTokenJWKS:
    """Tests for JWKS URI functionality.

    Note: With SSRF protection, JWKS fetches validate DNS and connect to the
    resolved IP. Tests mock DNS resolution to return a public IP.
    """

    @pytest.fixture
    def jwks_provider(self, rsa_key_pair: RSAKeyPair) -> JWTVerifier:
        """Provider configured with JWKS URI."""
        return JWTVerifier(
            jwks_uri="https://test.example.com/.well-known/jwks.json",
            issuer="https://test.example.com",
            audience="https://api.example.com",
        )

    @pytest.fixture
    def mock_jwks_data(self, rsa_key_pair: RSAKeyPair) -> JWKSData:
        """Create mock JWKS data from RSA key pair."""
        # Create JWK from the RSA public key
        public_key = jose_jwk.import_key(rsa_key_pair.public_key, "RSA")
        public_key_data = public_key.as_dict()
        kty = public_key_data["kty"]
        n = public_key_data["n"]
        e = public_key_data["e"]
        assert isinstance(kty, str)
        assert isinstance(n, str)
        assert isinstance(e, str)
        jwk_data = JWKData(
            kty=kty,
            n=n,
            e=e,
            kid="test-key-1",
            alg="RS256",
        )

        return {"keys": [jwk_data]}

    @pytest.fixture
    def mock_dns(self):
        """Mock DNS resolution to return test public IP."""
        with patch(
            "fastmcp.server.auth.ssrf.resolve_hostname",
            return_value=[TEST_PUBLIC_IP],
        ):
            yield

    async def test_jwks_token_validation(
        self,
        rsa_key_pair: RSAKeyPair,
        jwks_provider: JWTVerifier,
        mock_jwks_data: JWKSData,
        httpx_mock: HTTPXMock,
        mock_dns,
    ):
        """Test token validation using JWKS URI."""
        httpx_mock.add_response(json=mock_jwks_data)

        username = "test-user"
        issuer = "https://test.example.com"
        audience = "https://api.example.com"

        token = rsa_key_pair.create_token(
            subject=username,
            issuer=issuer,
            audience=audience,
        )

        access_token = await jwks_provider.load_access_token(token)
        assert access_token is not None
        assert access_token.client_id == username

        # ensure the raw claims are present - #1398
        assert access_token.claims.get("sub") == username
        assert access_token.claims.get("iss") == issuer
        assert access_token.claims.get("aud") == audience

    async def test_jwks_skips_unusable_keys(
        self,
        rsa_key_pair: RSAKeyPair,
        jwks_provider: JWTVerifier,
        mock_jwks_data: JWKSData,
        httpx_mock: HTTPXMock,
        mock_dns,
    ):
        """An unusable key must not poison the whole key set - #4515."""
        malformed_key = cast(
            "JWKData",
            {
                "kty": "RSA",
                "kid": "malformed-key",
                "use": "sig",
            },
        )
        mock_jwks_data["keys"][0]["kid"] = "test-key-1"
        # Malformed key FIRST, so an unguarded conversion loop would
        # abort before reaching the RSA key the token needs
        mock_jwks_data["keys"].insert(0, malformed_key)
        httpx_mock.add_response(json=mock_jwks_data)

        token = rsa_key_pair.create_token(
            subject="test-user",
            issuer="https://test.example.com",
            audience="https://api.example.com",
            kid="test-key-1",
        )

        access_token = await jwks_provider.load_access_token(token)
        assert access_token is not None
        assert access_token.client_id == "test-user"

    async def test_jwks_ignores_other_algorithm_key_types_without_kid(
        self,
        rsa_key_pair: RSAKeyPair,
        jwks_provider: JWTVerifier,
        mock_jwks_data: JWKSData,
        httpx_mock: HTTPXMock,
        mock_dns,
    ):
        """Unrelated key types do not make a no-kid lookup ambiguous."""
        _, public_pem = create_okp_key_pair(Ed25519PrivateKey.generate())
        okp_key = jose_jwk.import_key(public_pem, "OKP").as_dict()
        okp_key.update(kid="ed25519-key", alg="Ed25519", use="sig")
        mock_jwks_data["keys"].append(cast("JWKData", okp_key))
        httpx_mock.add_response(json=mock_jwks_data)

        token = rsa_key_pair.create_token(
            subject="test-user",
            issuer="https://test.example.com",
            audience="https://api.example.com",
        )

        access_token = await jwks_provider.load_access_token(token)

        assert access_token is not None
        assert access_token.client_id == "test-user"

    async def test_jwks_with_only_unusable_keys_rejects_cleanly(
        self,
        rsa_key_pair: RSAKeyPair,
        jwks_provider: JWTVerifier,
        httpx_mock: HTTPXMock,
        mock_dns,
    ):
        """If every key in the JWKS is unusable, verification fails
        cleanly (returns None) rather than crashing - #4515."""
        unusable_only = {
            "keys": [
                cast(
                    "JWKData",
                    {
                        "kty": "RSA",
                        "kid": "malformed-key",
                        "use": "sig",
                    },
                )
            ]
        }
        httpx_mock.add_response(json=unusable_only)

        token = rsa_key_pair.create_token(
            subject="test-user",
            issuer="https://test.example.com",
            audience="https://api.example.com",
            kid="malformed-key",
        )

        access_token = await jwks_provider.load_access_token(token)
        assert access_token is None

    async def test_jwks_token_validation_with_invalid_key(
        self,
        rsa_key_pair: RSAKeyPair,
        rsa_key_pair_2: RSAKeyPair,
        jwks_provider: JWTVerifier,
        mock_jwks_data: JWKSData,
        httpx_mock: HTTPXMock,
        mock_dns,
    ):
        httpx_mock.add_response(json=mock_jwks_data)
        token = rsa_key_pair_2.create_token(
            subject="test-user",
            issuer="https://test.example.com",
            audience="https://api.example.com",
        )

        access_token = await jwks_provider.load_access_token(token)
        assert access_token is None

    async def test_jwks_token_validation_with_kid(
        self,
        rsa_key_pair: RSAKeyPair,
        jwks_provider: JWTVerifier,
        mock_jwks_data: JWKSData,
        httpx_mock: HTTPXMock,
        mock_dns,
    ):
        mock_jwks_data["keys"][0]["kid"] = "test-key-1"
        httpx_mock.add_response(json=mock_jwks_data)
        token = rsa_key_pair.create_token(
            subject="test-user",
            issuer="https://test.example.com",
            audience="https://api.example.com",
            kid="test-key-1",
        )

        access_token = await jwks_provider.load_access_token(token)
        assert access_token is not None
        assert access_token.client_id == "test-user"

    async def test_jwks_token_validation_with_kid_and_no_kid_in_token(
        self,
        rsa_key_pair: RSAKeyPair,
        jwks_provider: JWTVerifier,
        mock_jwks_data: JWKSData,
        httpx_mock: HTTPXMock,
        mock_dns,
    ):
        mock_jwks_data["keys"][0]["kid"] = "test-key-1"
        httpx_mock.add_response(json=mock_jwks_data)
        token = rsa_key_pair.create_token(
            subject="test-user",
            issuer="https://test.example.com",
            audience="https://api.example.com",
        )

        access_token = await jwks_provider.load_access_token(token)
        assert access_token is not None
        assert access_token.client_id == "test-user"

    async def test_jwks_token_validation_with_no_kid_and_kid_in_jwks(
        self,
        rsa_key_pair: RSAKeyPair,
        jwks_provider: JWTVerifier,
        mock_jwks_data: JWKSData,
        httpx_mock: HTTPXMock,
        mock_dns,
    ):
        mock_jwks_data["keys"][0]["kid"] = "test-key-1"
        httpx_mock.add_response(json=mock_jwks_data)
        token = rsa_key_pair.create_token(
            subject="test-user",
            issuer="https://test.example.com",
            audience="https://api.example.com",
        )

        access_token = await jwks_provider.load_access_token(token)
        assert access_token is not None
        assert access_token.client_id == "test-user"

    async def test_jwks_token_validation_with_kid_mismatch(
        self,
        rsa_key_pair: RSAKeyPair,
        jwks_provider: JWTVerifier,
        mock_jwks_data: JWKSData,
        httpx_mock: HTTPXMock,
        mock_dns,
    ):
        mock_jwks_data["keys"][0]["kid"] = "test-key-1"
        httpx_mock.add_response(json=mock_jwks_data)
        token = rsa_key_pair.create_token(
            subject="test-user",
            issuer="https://test.example.com",
            audience="https://api.example.com",
            kid="test-key-2",
        )

        access_token = await jwks_provider.load_access_token(token)
        assert access_token is None

    async def test_jwks_token_validation_with_multiple_keys_and_no_kid_in_token(
        self,
        rsa_key_pair: RSAKeyPair,
        jwks_provider: JWTVerifier,
        mock_jwks_data: JWKSData,
        httpx_mock: HTTPXMock,
        mock_dns,
    ):
        mock_jwks_data["keys"] = [
            {
                "kid": "test-key-1",
                "alg": "RS256",
            },
            {
                "kid": "test-key-2",
                "alg": "RS256",
            },
        ]

        httpx_mock.add_response(json=mock_jwks_data)
        token = rsa_key_pair.create_token(
            subject="test-user",
            issuer="https://test.example.com",
            audience="https://api.example.com",
        )

        access_token = await jwks_provider.load_access_token(token)
        assert access_token is None