import base64

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

from . import ty


class KeyPair(ty.BaseModel):
    public_pem: str
    private_pem: str

    @property
    def public_pem_b64(self) -> str:
        return base64.b64encode(self.public_pem.encode("utf8")).decode("utf8")

    @property
    def private_pem_b64(self) -> str:
        return base64.b64encode(self.private_pem.encode("utf8")).decode("utf8")


def generate_key_pair(password: ty.Optional[str] = None) -> KeyPair:
    """Generate the Public/Private KeyPair for asymmetric JWK PAM server."""
    private_key = generate_private_key(public_exponent=65537, key_size=2048)
    password_bytes = password.encode("utf-8") if password else b""
    encryption = (
        serialization.BestAvailableEncryption(password_bytes)
        if password
        else serialization.NoEncryption()
    )
    pem = private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=encryption,
    )
    public_key = private_key.public_key()
    public_pem = public_key.public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo,
    )
    return KeyPair(private_pem=pem.decode("utf8"), public_pem=public_pem.decode("utf8"))
