"""CAP-037 KMS round-trip test (REQ-350).

Sign a test JWT via ``core.kms_signing.sign_jwt()`` (mock KMS with a
test keypair) → fetch JWKS via ``nova_idp_jwks.lambda_handler()`` (mock
KMS) → verify the JWT with ``pyjwt`` using the JWKS key. Round-trip
succeeds — proves the DER→raw conversion + JWK export are mutually
consistent (the #1 gotcha from RESEARCH §5).
"""

from __future__ import annotations

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

import pytest

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

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

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

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, priv, pub_der):
        self._priv = priv
        self._pub_der = pub_der

    def sign(self, KeyId, Message, MessageType, SigningAlgorithm):
        return {"Signature": self._priv.sign(Message, ec.ECDSA(hashes.SHA256()))}

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


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


def test_cap037_kms_roundtrip():
    """Sign JWT → JWKS → pyjwt verify. The full KMS round-trip (REQ-350)."""
    priv = ec.generate_private_key(ec.SECP256R1())
    pub = priv.public_key()
    pub_der = pub.public_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PublicFormat.SubjectPublicKeyInfo,
    )
    kms_signing.set_kms_client_for_testing(_MockKms(priv, pub_der))

    # 1. Sign a JWT via kms_signing.sign_jwt (uses DER→raw conversion).
    claims = {
        "sub": "roundtrip-user", "aud": "nova-cli", "iss": "nova-idp",
        "exp": 9999999999, "iat": 1700000000, "jti": "rt-jti",
        "roles": ["developer"], "typ": "nova_oidc_token",
    }
    token = kms_signing.sign_jwt(claims, key_id="alias/nova-oidc-signing")

    # 2. Fetch the JWKS via the JWKS Lambda (mock KMS get_public_key).
    resp = jwks_mod.lambda_handler({}, None)
    assert resp["statusCode"] == 200
    jwks_body = json.loads(resp["body"])
    jwk = jwks_body["keys"][0]
    assert jwk["kty"] == "EC" and jwk["crv"] == "P-256"

    # 3. Verify the JWT with pyjwt using the JWKS key.
    key = pyjwt.PyJWK(jwk).key
    decoded = pyjwt.decode(token, key, algorithms=["ES256"], audience="nova-cli")
    assert decoded["sub"] == "roundtrip-user"
    assert decoded["jti"] == "rt-jti"
    assert decoded["roles"] == ["developer"]
    assert decoded["typ"] == "nova_oidc_token"