"""KMS signing tests (REQ-337, C-1.1).

Tests :func:`core.kms_signing.der_to_raw_ecdsa` with a known DER
signature and the full :func:`sign_jwt` round-trip with a mocked KMS
client (no real AWS calls — C-1.1 documented as a CI gate in
``docs/kms-provisioning.md``).
"""

from __future__ import annotations

import base64
import json
import sys
from pathlib import Path
from unittest import mock

import pytest

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

import core.kms_signing as kms_signing
from core.kms_signing import der_to_raw_ecdsa, sign_jwt, get_jwk

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


# ---------------------------------------------------------------------------
# Test keypair — generated once per session (P-256).
# ---------------------------------------------------------------------------


@pytest.fixture(scope="module")
def test_keypair():
    priv = ec.generate_private_key(ec.SECP256R1())
    pub = priv.public_key()
    return priv, pub


@pytest.fixture(scope="module")
def test_pub_der(test_keypair):
    _priv, pub = test_keypair
    return pub.public_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PublicFormat.SubjectPublicKeyInfo,
    )


class _MockKmsSignClient:
    """Mock KMS client that signs with a test ECDSA private key (DER)."""

    def __init__(self, priv, pub_der, key_id="alias/nova-oidc-signing"):
        self._priv = priv
        self._pub_der = pub_der
        self._key_id = key_id

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

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


# ---------------------------------------------------------------------------
# der_to_raw_ecdsa — unit test with a known DER signature.
# ---------------------------------------------------------------------------


def test_der_to_raw_ecdsa_known_vector():
    # Minimal DER: SEQUENCE { INTEGER r, INTEGER s }.
    # r=5, s=7 → DER: 30 06 02 01 05 02 01 07
    der = b"\x30\x06\x02\x01\x05\x02\x01\x07"
    raw = der_to_raw_ecdsa(der)
    assert len(raw) == 64  # 32 + 32
    r = int.from_bytes(raw[:32], "big")
    s = int.from_bytes(raw[32:], "big")
    assert r == 5
    assert s == 7


def test_der_to_raw_ecdsa_real_signature(test_keypair):
    priv, _ = test_keypair
    msg = b"test message for der->raw"
    der = priv.sign(msg, ec.ECDSA(hashes.SHA256()))
    raw = der_to_raw_ecdsa(der)
    assert len(raw) == 64
    # Round-trip: raw → (r, s) should verify against the message.
    r = int.from_bytes(raw[:32], "big")
    s = int.from_bytes(raw[32:], "big")
    from cryptography.hazmat.primitives.asymmetric.utils import encode_dss_signature
    der2 = encode_dss_signature(r, s)
    # Verifying with the re-encoded DER proves the raw split is correct.
    priv.public_key().verify(der2, msg, ec.ECDSA(hashes.SHA256()))


def test_der_to_raw_ecdsa_rejects_oversized_coord():
    # r needs 33 bytes (2**256+1) → should raise.
    from cryptography.hazmat.primitives.asymmetric.utils import encode_dss_signature
    der = encode_dss_signature(2**256 + 1, 1)
    with pytest.raises(ValueError):
        der_to_raw_ecdsa(der, coord_len=32)


# ---------------------------------------------------------------------------
# sign_jwt — full round-trip with mock KMS + pyjwt verification.
# ---------------------------------------------------------------------------


def test_sign_jwt_roundtrip_verifies_with_pyjwt(test_keypair, test_pub_der):
    priv, pub = test_keypair
    client = _MockKmsSignClient(priv, test_pub_der)
    kms_signing.set_kms_client_for_testing(client)
    try:
        claims = {
            "sub": "user-1",
            "aud": "nova-cli",
            "iss": "nova-idp",
            "exp": 9999999999,
            "iat": 1700000000,
            "jti": "test-jti",
            "roles": ["developer"],
        }
        token = sign_jwt(claims, key_id="alias/nova-oidc-signing")
        parts = token.split(".")
        assert len(parts) == 3  # header.payload.signature

        # Verify the header.
        header = json.loads(base64.urlsafe_b64decode(parts[0] + "=="))
        assert header["alg"] == "ES256"
        assert header["typ"] == "JWT"
        assert header["kid"] == "alias/nova-oidc-signing"

        # Verify the signature with pyjwt using the test public key.
        pem = pub.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo,
        ).decode("ascii")
        decoded = pyjwt.decode(token, pem, algorithms=["ES256"], options={"verify_aud": False})
        assert decoded["sub"] == "user-1"
        assert decoded["jti"] == "test-jti"
        assert decoded["roles"] == ["developer"]
    finally:
        kms_signing.set_kms_client_for_testing(None)


def test_get_jwk_returns_valid_ec_jwk(test_keypair, test_pub_der):
    priv, pub = test_keypair
    client = _MockKmsSignClient(priv, test_pub_der)
    kms_signing.set_kms_client_for_testing(client)
    try:
        jwk = get_jwk(key_id="alias/nova-oidc-signing")
        assert jwk["kty"] == "EC"
        assert jwk["crv"] == "P-256"
        assert jwk["kid"] == "alias/nova-oidc-signing"
        assert jwk["alg"] == "ES256"
        # x and y are 32 bytes → 43 base64url chars (no padding).
        assert len(jwk["x"]) == 43
        assert len(jwk["y"]) == 43
        # Cross-verify: a JWT signed with the test private key verifies
        # with pyjwt using this JWK as the key.
        client2 = _MockKmsSignClient(priv, test_pub_der)
        kms_signing.set_kms_client_for_testing(client2)
        token = sign_jwt({"sub": "x", "exp": 9999999999, "iat": 1, "jti": "j"})
        key = pyjwt.PyJWK(jwk).key
        decoded = pyjwt.decode(token, key, algorithms=["ES256"], options={"verify_aud": False})
        assert decoded["sub"] == "x"
    finally:
        kms_signing.set_kms_client_for_testing(None)