"""Unit tests for the pluggable auth layer: both verifiers (mock + JWKS) and
the generic OAuth2 client-credentials provider, in isolation from any
running agent process.
"""
from __future__ import annotations

import time
from unittest.mock import patch

import httpx
import jwt
import pytest
import respx
from cryptography.hazmat.primitives.asymmetric import rsa

from a2a_passport.auth.oauth2 import OAuth2ClientCredentialsProvider
from a2a_passport.auth.verifiers import JWKSVerifier, SharedSecretVerifier
from a2a_passport.exceptions import TokenAcquisitionError, TokenVerificationError

SECRET = "unit-test-secret"


def _mint_hs256(audience: str, scope: str, exp_delta: int = 60, secret: str = SECRET) -> str:
    now = int(time.time())
    claims = {
        "sub": "test-client",
        "aud": audience,
        "scope": scope,
        "iat": now,
        "exp": now + exp_delta,
    }
    return jwt.encode(claims, secret, algorithm="HS256")


class TestSharedSecretVerifier:
    async def test_valid_token_verifies(self):
        verifier = SharedSecretVerifier(SECRET)
        token = _mint_hs256(audience="weather-agent", scope="weather:read")
        claims = await verifier.verify(token, audience="weather-agent")
        assert claims.subject == "test-client"
        assert claims.has_scope("weather:read")

    async def test_wrong_audience_rejected(self):
        verifier = SharedSecretVerifier(SECRET)
        token = _mint_hs256(audience="some-other-agent", scope="weather:read")
        with pytest.raises(TokenVerificationError) as exc_info:
            await verifier.verify(token, audience="weather-agent")
        assert exc_info.value.error_code == "invalid_audience"

    async def test_expired_token_rejected(self):
        verifier = SharedSecretVerifier(SECRET)
        token = _mint_hs256(audience="weather-agent", scope="weather:read", exp_delta=-10)
        with pytest.raises(TokenVerificationError) as exc_info:
            await verifier.verify(token, audience="weather-agent")
        assert exc_info.value.error_code == "token_expired"

    async def test_wrong_secret_rejected(self):
        verifier = SharedSecretVerifier(SECRET)
        token = _mint_hs256(audience="weather-agent", scope="weather:read", secret="not-the-secret")
        with pytest.raises(TokenVerificationError) as exc_info:
            await verifier.verify(token, audience="weather-agent")
        assert exc_info.value.error_code == "invalid_token"

    async def test_multiple_scopes_parsed(self):
        verifier = SharedSecretVerifier(SECRET)
        token = _mint_hs256(audience="weather-agent", scope="weather:read weather:admin")
        claims = await verifier.verify(token, audience="weather-agent")
        assert claims.has_scope("weather:read")
        assert claims.has_scope("weather:admin")
        assert not claims.has_scope("nope")


class TestJWKSVerifier:
    """Exercises real RS256 signature/audience/issuer/expiry checks without
    any network call: get_signing_key_from_jwt is patched to return the key
    we signed with, which is the only network-touching call PyJWKClient makes.
    """

    @pytest.fixture
    def rsa_keypair(self):
        private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
        return private_key

    def _mint_rs256(
        self, private_key, *, audience, scope, issuer="https://idp.example.test/", exp_delta=60
    ):
        now = int(time.time())
        claims = {
            "sub": "test-client",
            "aud": audience,
            "iss": issuer,
            "scope": scope,
            "iat": now,
            "exp": now + exp_delta,
        }
        return jwt.encode(claims, private_key, algorithm="RS256")

    async def test_valid_token_verifies(self, rsa_keypair):
        verifier = JWKSVerifier(
            jwks_uri="https://idp.example.test/.well-known/jwks.json",
            issuer="https://idp.example.test/",
        )
        token = self._mint_rs256(rsa_keypair, audience="translator-agent", scope="translator:use")

        fake_signing_key = jwt.PyJWK.from_json(
            jwt.algorithms.RSAAlgorithm.to_jwk(rsa_keypair.public_key())
        )
        with patch.object(
            verifier._jwks_client, "get_signing_key_from_jwt", return_value=fake_signing_key
        ):
            claims = await verifier.verify(token, audience="translator-agent")
        assert claims.subject == "test-client"
        assert claims.has_scope("translator:use")

    async def test_wrong_issuer_rejected(self, rsa_keypair):
        verifier = JWKSVerifier(
            jwks_uri="https://idp.example.test/.well-known/jwks.json",
            issuer="https://idp.example.test/",
        )
        token = self._mint_rs256(
            rsa_keypair,
            audience="translator-agent",
            scope="translator:use",
            issuer="https://someone-else.test/",
        )

        fake_signing_key = jwt.PyJWK.from_json(
            jwt.algorithms.RSAAlgorithm.to_jwk(rsa_keypair.public_key())
        )
        with patch.object(
            verifier._jwks_client, "get_signing_key_from_jwt", return_value=fake_signing_key
        ):
            with pytest.raises(TokenVerificationError) as exc_info:
                await verifier.verify(token, audience="translator-agent")
        assert exc_info.value.error_code == "invalid_issuer"


class TestOAuth2ClientCredentialsProvider:
    @respx.mock
    async def test_acquires_and_returns_token(self):
        route = respx.post("https://idp.example.test/token").mock(
            return_value=httpx.Response(
                200,
                json={"access_token": "abc123", "expires_in": 300, "token_type": "Bearer"},
            )
        )
        provider = OAuth2ClientCredentialsProvider(
            token_url="https://idp.example.test/token", client_id="cid", client_secret="csecret"
        )
        token = await provider.get_token("weather-agent")
        assert token == "abc123"
        assert route.called
        sent_form = dict(httpx.QueryParams(route.calls.last.request.content.decode()))
        assert sent_form["audience"] == "weather-agent"
        assert sent_form["grant_type"] == "client_credentials"

    @respx.mock
    async def test_caches_token_per_audience(self):
        route = respx.post("https://idp.example.test/token").mock(
            return_value=httpx.Response(
                200, json={"access_token": "cached-token", "expires_in": 300}
            )
        )
        provider = OAuth2ClientCredentialsProvider(
            token_url="https://idp.example.test/token", client_id="cid", client_secret="csecret"
        )
        await provider.get_token("weather-agent")
        await provider.get_token("weather-agent")
        assert route.call_count == 1

    @respx.mock
    async def test_separate_cache_per_audience(self):
        respx.post("https://idp.example.test/token").mock(
            return_value=httpx.Response(200, json={"access_token": "tok", "expires_in": 300})
        )
        provider = OAuth2ClientCredentialsProvider(
            token_url="https://idp.example.test/token", client_id="cid", client_secret="csecret"
        )
        await provider.get_token("weather-agent")
        await provider.get_token("translator-agent")
        assert respx.calls.call_count == 2

    @respx.mock
    async def test_non_200_raises_token_acquisition_error(self):
        respx.post("https://idp.example.test/token").mock(
            return_value=httpx.Response(401, json={"error": "invalid_client"})
        )
        provider = OAuth2ClientCredentialsProvider(
            token_url="https://idp.example.test/token", client_id="cid", client_secret="wrong"
        )
        with pytest.raises(TokenAcquisitionError):
            await provider.get_token("weather-agent")
