import secrets
import time
from unittest.mock import patch
from urllib.parse import parse_qs, urlsplit

from authlib.integrations.django_client import OAuth
from django.conf import settings
from django.contrib.sessions.backends.db import SessionStore
from django.test import RequestFactory, TestCase
from joserfc import jwt
from joserfc.jwk import RSAKey


class AuthlibFlowTests(TestCase):
    """Exercise real Authlib state, PKCE, nonce and signature validation offline."""

    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        cls.key = RSAKey.generate_key(2048, parameters={"kid": "test-key"})

    def flow(self, **changes):
        issuer = settings.DEX_ISSUER
        remote = OAuth().register("dex", client_id=settings.DEX_CLIENT_ID, client_secret="test-only", issuer=issuer, authorization_endpoint=issuer + "/auth", token_endpoint=issuer + "/token", jwks_uri=issuer + "/keys", id_token_signing_alg_values_supported=["RS256"], client_kwargs={"scope": "openid email profile federated:id", "code_challenge_method": "S256"})
        request = RequestFactory().get("/members/login/", secure=True)
        request.session = SessionStore()
        nonce = secrets.token_urlsafe(32)
        response = remote.authorize_redirect(request, "https://eom.dev/members/oidc/callback/", nonce=nonce)
        auth_params = parse_qs(urlsplit(response.url).query)
        self.assertEqual(auth_params["code_challenge_method"], ["S256"])
        self.assertTrue(auth_params["code_challenge"][0])
        callback = RequestFactory().get("/members/oidc/callback/", {"code": "test-code", "state": auth_params["state"][0]}, secure=True)
        callback.session = request.session
        claims = {"iss": issuer, "sub": "test-sub", "aud": settings.DEX_CLIENT_ID, "iat": int(time.time()), "exp": int(time.time()) + 300, "nonce": nonce, **changes}
        token = {"id_token": jwt.encode({"alg": "RS256", "kid": "test-key"}, claims, self.key)}
        return remote, callback, token

    def exchange(self, remote, callback, token):
        with patch.object(remote, "fetch_access_token", return_value=token) as exchange, patch.object(remote, "fetch_jwk_set", return_value={"keys": [self.key.as_dict(private=False)]}):
            result = remote.authorize_access_token(callback, claims_options={"iss": {"essential": True, "value": settings.DEX_ISSUER}}, leeway=0)
            self.assertTrue(exchange.call_args.kwargs["code_verifier"])
            return result

    def test_real_pkce_nonce_signature_and_one_use_state(self):
        remote, callback, token = self.flow()
        result = self.exchange(remote, callback, token)
        self.assertEqual(result["userinfo"]["sub"], "test-sub")
        with self.assertRaises(Exception):
            self.exchange(remote, callback, token)

    def test_wrong_nonce_issuer_audience_expiry_rejected(self):
        for changes in ({"nonce": "attacker"}, {"iss": "https://evil.example"}, {"aud": "another-client"}, {"exp": int(time.time()) - 100}):
            with self.subTest(changes=changes):
                remote, callback, token = self.flow(**changes)
                with self.assertRaises(Exception):
                    self.exchange(remote, callback, token)

    def test_wrong_state_never_exchanges_code(self):
        remote, callback, token = self.flow()
        callback.GET = callback.GET.copy()
        callback.GET["state"] = "unknown-state"
        with patch.object(remote, "fetch_access_token") as exchange:
            with self.assertRaises(Exception):
                remote.authorize_access_token(callback)
            exchange.assert_not_called()

    def test_bad_signature_rejected(self):
        remote, callback, token = self.flow()
        other = RSAKey.generate_key(2048, parameters={"kid": "test-key"})
        claims = {"iss": settings.DEX_ISSUER, "sub": "attacker", "aud": settings.DEX_CLIENT_ID, "iat": int(time.time()), "exp": int(time.time()) + 300, "nonce": "attacker"}
        token["id_token"] = jwt.encode({"alg": "RS256", "kid": "test-key"}, claims, other)
        with self.assertRaises(Exception):
            self.exchange(remote, callback, token)