"""
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 base64
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 .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 _b64url(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def key_id(private_key: RSAPrivateKey) -> str:
    """Return a stable 16-hex-char key ID derived from the public key modulus."""
    n = private_key.public_key().public_numbers().n
    n_bytes = n.to_bytes((n.bit_length() + 7) // 8, byteorder="big")
    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)."""
    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 {
        "kty": "RSA",
        "use": "sig",
        "alg": "RS256",
        "kid": key_id(private_key),
        "n": _b64url(n_bytes),
        "e": _b64url(e_bytes),
    }