"""KMS-signed JWT issuance for the Nova IdP (REQ-337, REQ-336).

Signs OIDC tokens with an AWS KMS asymmetric key (``ECC_NIST_P256``,
``ECDSA_SHA_256`` → JWS ``ES256``) and exposes the public key as a JWK
for the JWKS endpoint (REQ-338).

## DER → raw ECDSA conversion (the #1 gotcha, RESEARCH §5)

KMS ``sign()`` returns a **DER-encoded** ASN.1 ECDSA signature. JWS
(RFC 7515 §3.1.3) requires the **raw** ``r‖s`` concatenation, each
coordinate 32 bytes big-endian. :func:`der_to_raw_ecdsa` performs the
conversion via ``cryptography``'s ``decode_dss_signature``. This is the
core of REQ-337 and is verified by the CAP-037 round-trip test.

## Lazy boto3

``boto3.client("kms")`` is constructed lazily so the module imports
without AWS creds (mirrors ``nova_idp_auth.py``). Tests inject a mock
client via :func:`set_kms_client_for_testing`.
"""

from __future__ import annotations

import base64
import json
import os
from typing import Any

import boto3

from cryptography.hazmat.primitives.asymmetric.utils import decode_dss_signature
from cryptography.hazmat.primitives.asymmetric.ec import (
    EllipticCurvePublicKey,
)
from cryptography.hazmat.primitives.serialization import load_der_public_key
from cryptography.hazmat.primitives.asymmetric import ec


# Default KMS key alias for Nova OIDC signing (REQ-337).
DEFAULT_KEY_ID = os.environ.get("NOVA_OIDC_KMS_KEY_ID", "alias/nova-oidc-signing")

_kms_client = None


def _get_kms_client():
    """Lazy boto3 KMS client singleton (mirrors nova_idp_auth.py)."""
    global _kms_client
    if _kms_client is None:
        _kms_client = boto3.client("kms")
    return _kms_client


def set_kms_client_for_testing(client: Any) -> None:
    """Inject a mock KMS client for tests (no real AWS calls)."""
    global _kms_client
    _kms_client = client


def _b64url(data: bytes) -> str:
    """Base64url encode without padding (RFC 7515 §2)."""
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def der_to_raw_ecdsa(der_sig: bytes, coord_len: int = 32) -> bytes:
    """Convert a DER-encoded ECDSA signature to raw ``r‖s`` (JWS format).

    KMS returns DER; JWS requires raw ``r‖s`` concatenation, each
    coordinate ``coord_len`` bytes big-endian (32 for P-256, 48 for
    P-384). Uses ``cryptography``'s ``decode_dss_signature`` to parse
    the DER, then zero-pads each integer to ``coord_len``.

    Raises:
        ValueError: if a coordinate does not fit in ``coord_len`` bytes
            (the integer is larger than the curve allows — indicates a
            malformed signature or wrong ``coord_len``).
    """
    r, s = decode_dss_signature(der_sig)
    if r.bit_length() > coord_len * 8 or s.bit_length() > coord_len * 8:
        raise ValueError(
            f"ECDSA coordinate does not fit in {coord_len} bytes "
            f"(r={r.bit_length()} bits, s={s.bit_length()} bits)"
        )
    return r.to_bytes(coord_len, "big") + s.to_bytes(coord_len, "big")


def sign_jwt(claims: dict, key_id: str = DEFAULT_KEY_ID) -> str:
    """Build + sign a JWT with KMS (REQ-337, REQ-336).

    Args:
        claims: the JWT claims payload (``sub, aud, iss, exp, iat, jti,
            roles`` per REQ-336, plus ``typ`` for PATs).
        key_id: the KMS key ID or alias (default
            ``alias/nova-oidc-signing``).

    Returns:
        The compact JWS (``header.payload.signature``), ``ES256``,
        with the signature in raw ``r‖s`` form (DER→raw converted).
    """
    header = {"alg": "ES256", "typ": "JWT", "kid": key_id}
    signing_input = (
        _b64url(json.dumps(header, separators=(",", ":"), sort_keys=True).encode())
        + "."
        + _b64url(json.dumps(claims, separators=(",", ":"), sort_keys=True).encode())
    )
    resp = _get_kms_client().sign(
        KeyId=key_id,
        Message=signing_input.encode("ascii"),
        MessageType="RAW",
        SigningAlgorithm="ECDSA_SHA_256",
    )
    der_sig = resp["Signature"]
    raw_sig = der_to_raw_ecdsa(der_sig)
    return signing_input + "." + _b64url(raw_sig)


def get_jwk(key_id: str = DEFAULT_KEY_ID) -> dict:
    """Fetch the KMS public key and return it as a JWK (REQ-338).

    Calls ``kms.get_public_key`` → DER SPKI → ``cryptography``'s
    ``load_der_public_key`` → JWK ``{"kty":"EC","crv":"P-256","kid":...,
    "x":...,"y":...}``. The ``x``/``y`` are base64url-encoded
    big-endian 32-byte coordinates.
    """
    resp = _get_kms_client().get_public_key(KeyId=key_id)
    pub = load_der_public_key(resp["PublicKey"])
    if not isinstance(pub, EllipticCurvePublicKey):
        raise ValueError(
            f"KMS public key is not an EC key (got {type(pub).__name__})"
        )
    nums = pub.public_numbers()
    # P-256 coordinates are 32 bytes big-endian.
    x = nums.x.to_bytes(32, "big")
    y = nums.y.to_bytes(32, "big")
    return {
        "kty": "EC",
        "crv": "P-256",
        "kid": key_id,
        "x": _b64url(x),
        "y": _b64url(y),
        "alg": "ES256",
        "use": "sig",
    }


if __name__ == "__main__":  # pragma: no cover - CLI inspection helper
    import sys

    if "--print-jwk" in sys.argv:
        print(json.dumps(get_jwk(), indent=2))
    else:
        print("usage: python3 -m core.kms_signing --print-jwks", file=sys.stderr)