"""
RSA signing key management for local-identity OIDC serve.

The signing key is generated on first invocation and stored at:
  ~/.local-identity/keys/signing_private.pem   (mode 600)

The corresponding public key is never stored separately — it is always
derived from the private key on load.
"""

import hashlib
import os
from pathlib import Path

from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey

from .jwt_utils import _b64url_encode as _b64url
from .store import _store_dir


def _keys_dir() -> Path:
    return _store_dir() / "keys"


def ensure_signing_key() -> RSAPrivateKey:
    """Load the signing key from disk, or generate a new RSA-2048 key if absent."""
    keys_dir = _keys_dir()
    keys_dir.mkdir(mode=0o700, parents=True, exist_ok=True)

    priv_path = keys_dir / "signing_private.pem"

    if priv_path.exists():
        priv_pem = priv_path.read_bytes()
        return serialization.load_pem_private_key(priv_pem, password=None)

    private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)

    priv_pem = private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.TraditionalOpenSSL,
        encryption_algorithm=serialization.NoEncryption(),
    )
    priv_path.write_bytes(priv_pem)
    os.chmod(priv_path, 0o600)

    return private_key


def _public_key_bytes(private_key: RSAPrivateKey) -> tuple[bytes, bytes]:
    """Return (n_bytes, e_bytes) for the public key — shared by key_id and jwk_public."""
    pub = private_key.public_key().public_numbers()
    n_bytes = pub.n.to_bytes((pub.n.bit_length() + 7) // 8, byteorder="big")
    e_bytes = pub.e.to_bytes((pub.e.bit_length() + 7) // 8, byteorder="big")
    return n_bytes, e_bytes


def key_id(private_key: RSAPrivateKey) -> str:
    """Return a stable 16-hex-char key ID derived from the public key modulus."""
    n_bytes, _ = _public_key_bytes(private_key)
    return hashlib.sha256(n_bytes).hexdigest()[:16]


def jwk_public(private_key: RSAPrivateKey) -> dict:
    """Return the RSA public key as a JWK dict (RS256, sig use)."""
    n_bytes, e_bytes = _public_key_bytes(private_key)
    return {
        "kty": "RSA",
        "use": "sig",
        "alg": "RS256",
        "kid": hashlib.sha256(n_bytes).hexdigest()[:16],
        "n": _b64url(n_bytes),
        "e": _b64url(e_bytes),
    }