"""JWKS endpoint tests (REQ-338, D-230).

Mocks ``kms.get_public_key`` with a test ECDSA P-256 public key DER →
asserts the Lambda returns 200 + the right headers + a valid JWK.
Cross-verifies: a JWT signed with the test private key verifies with
pyjwt using the JWKS key.
"""

from __future__ import annotations

import importlib.util
import json
import os
import sys
from pathlib import Path
from unittest import mock

import pytest

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

os.environ.setdefault("NOVA_LAMBDA_LOCAL_BYPASS", "1")

_SOURCE_PATH = (
    Path(__file__).resolve().parent.parent / "core" / "lambda" / "nova_idp_jwks.py"
)
_spec = importlib.util.spec_from_file_location("nova_idp_jwks", _SOURCE_PATH)
jwks = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(jwks)

import core.kms_signing as kms_signing
import jwt as pyjwt
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives import hashes, serialization


class _MockKms:
    def __init__(self, pub_der):
        self._pub_der = pub_der

    def get_public_key(self, KeyId):
        return {"PublicKey": self._pub_der, "KeyId": KeyId}

    def sign(self, KeyId, Message, MessageType, SigningAlgorithm):
        # Provided so sign_jwt works in the cross-verify test.
        return {"Signature": self._priv.sign(Message, ec.ECDSA(hashes.SHA256()))}


@pytest.fixture
def test_keypair():
    priv = ec.generate_private_key(ec.SECP256R1())
    pub = priv.public_key()
    pub_der = pub.public_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PublicFormat.SubjectPublicKeyInfo,
    )
    return priv, pub, pub_der


@pytest.fixture(autouse=True)
def _reset():
    yield
    kms_signing.set_kms_client_for_testing(None)


def test_jwks_returns_200_and_headers(test_keypair):
    _priv, _pub, pub_der = test_keypair
    kms_signing.set_kms_client_for_testing(_MockKms(pub_der))
    resp = jwks.lambda_handler({}, None)
    assert resp["statusCode"] == 200
    headers = resp["headers"]
    assert headers["Content-Type"] == "application/json"
    assert headers["Cache-Control"] == "public, max-age=3600"
    assert headers["Access-Control-Allow-Origin"] == "*"


def test_jwks_returns_valid_ec_jwk(test_keypair):
    _priv, _pub, pub_der = test_keypair
    kms_signing.set_kms_client_for_testing(_MockKms(pub_der))
    resp = jwks.lambda_handler({}, None)
    body = json.loads(resp["body"])
    assert "keys" in body
    assert len(body["keys"]) == 1
    jwk = body["keys"][0]
    assert jwk["kty"] == "EC"
    assert jwk["crv"] == "P-256"
    assert "kid" in jwk
    assert "x" in jwk and "y" in jwk
    assert len(jwk["x"]) == 43  # 32 bytes → 43 base64url chars
    assert len(jwk["y"]) == 43


def test_jwks_cross_verifies_jwt(test_keypair):
    """A JWT signed with the test private key verifies with the JWKS key."""
    priv, _pub, pub_der = test_keypair
    # Mock KMS that can both sign (for sign_jwt) and serve the public key.
    mock_kms = _MockKms(pub_der)
    mock_kms._priv = priv
    kms_signing.set_kms_client_for_testing(mock_kms)

    # Sign a JWT via kms_signing.sign_jwt.
    token = kms_signing.sign_jwt(
        {"sub": "user-1", "exp": 9999999999, "iat": 1, "jti": "j", "aud": "nova-cli"},
        key_id="alias/nova-oidc-signing",
    )
    # Fetch the JWKS via the Lambda.
    resp = jwks.lambda_handler({}, None)
    jwk = json.loads(resp["body"])["keys"][0]
    # Verify the JWT with pyjwt using the JWK.
    key = pyjwt.PyJWK(jwk).key
    decoded = pyjwt.decode(token, key, algorithms=["ES256"], options={"verify_aud": False})
    assert decoded["sub"] == "user-1"
    assert decoded["jti"] == "j"


def test_jwks_500_on_kms_error():
    kms_signing.set_kms_client_for_testing(None)
    # Force get_jwk to raise by using a broken client.
    broken = mock.MagicMock()
    broken.get_public_key.side_effect = RuntimeError("KMS down")
    kms_signing.set_kms_client_for_testing(broken)
    resp = jwks.lambda_handler({}, None)
    assert resp["statusCode"] == 500