"""
ES256 keypair generation, JWK formatting, and thumbprint calculation.
"""

from __future__ import annotations

import base64
import hashlib
import json

from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ec


def generate_es256_keypair():
    """
    Generate an EC P-256 keypair.

    Returns:
        (private_key, public_jwk_dict) — private key object and JWK public key dict.
    """
    private_key = ec.generate_private_key(ec.SECP256R1())
    public_key = private_key.public_key()
    public_numbers = public_key.public_numbers()

    x_bytes = public_numbers.x.to_bytes(32, byteorder="big")
    y_bytes = public_numbers.y.to_bytes(32, byteorder="big")
    x_b64 = base64.urlsafe_b64encode(x_bytes).rstrip(b"=").decode()
    y_b64 = base64.urlsafe_b64encode(y_bytes).rstrip(b"=").decode()

    jwk = {"kty": "EC", "crv": "P-256", "x": x_b64, "y": y_b64}
    return private_key, jwk


def compute_jwk_thumbprint(jwk):
    """
    Compute RFC 7638 JWK thumbprint (SHA-256, base64url).

    Uses canonical JSON with alphabetically sorted required members for EC keys.
    """
    canonical = json.dumps(
        {"crv": jwk["crv"], "kty": jwk["kty"], "x": jwk["x"], "y": jwk["y"]},
        sort_keys=True,
        separators=(",", ":"),
    )
    digest = hashlib.sha256(canonical.encode()).digest()
    return base64.urlsafe_b64encode(digest).rstrip(b"=").decode()


def serialize_private_key(private_key):
    """Serialize private key to PEM string."""
    return private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption(),
    ).decode()


def load_private_key(pem_str):
    """Load private key from PEM string."""
    return serialization.load_pem_private_key(pem_str.encode(), password=None)
