import datetime
from io import BufferedWriter
import re
import cryptography.x509 as x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives.asymmetric.types import PrivateKeyTypes
import pathlib

def read_private_key(pem_path: str) -> PrivateKeyTypes:
    """
    Read a PEM file and extract the private key in bytes.
    """
    with open(pem_path, "rb") as f:
        pem = f.read()
        key_pem = re.search(b"-----BEGIN (?:RSA )?PRIVATE KEY-----.*?END (?:RSA )?PRIVATE KEY-----", pem, re.S).group(0)
        key = serialization.load_pem_private_key(key_pem, password=None)
        return key

def read_public_certificate(pem_path: str) -> any:
    """
    Read a PEM file and extract the public certificate.
    """
    with open(pem_path, "rb") as f:
        pem = f.read()
        cert_pem = re.search(b"-----BEGIN CERTIFICATE-----.*?END CERTIFICATE-----", pem, re.S).group(0)
        cert = x509.load_pem_x509_certificate(cert_pem)
        return cert

def _write_pem_key(
        writer: BufferedWriter,
        key: PrivateKeyTypes,
        encoding: serialization.Encoding = serialization.Encoding.PEM,
        format: serialization.PrivateFormat = serialization.PrivateFormat.PKCS8,
        password: str | None = None
):
    """
    Write a PEM encoded private key to the given writer.
    Allows optional password protection.
    """
    if password:
        encryption_algorithm = serialization.BestAvailableEncryption(password.encode())
    else:
        encryption_algorithm = serialization.NoEncryption()

    key_pem = key.private_bytes(
        encoding=encoding,
        format=format,
        encryption_algorithm=encryption_algorithm,
    )
    writer.write(key_pem)

def _write_pem_cert(
        writer: BufferedWriter,
        cert: x509.Certificate,
        encoding: serialization.Encoding = serialization.Encoding.PEM
        ):
    cert_pem = cert.public_bytes(encoding=encoding)
    writer.write(cert_pem)

def write_pem_file(
        pem_file: pathlib.Path,
        cert: x509.Certificate | None = None,
        key: PrivateKeyTypes | None = None,
        encoding: serialization.Encoding = serialization.Encoding.PEM,
        key_serialization_format: serialization.PrivateFormat = serialization.PrivateFormat.PKCS8,
        password: str | None = None
        ):
    """
    Write the certificate and/or private key to a PEM file.
    """
    with open(pem_file, "wb") as f:
        if cert:
            _write_pem_cert(
                f,
                cert,
                encoding=encoding
                )
        if key:
            _write_pem_key(
                f,
                key,
                encoding=encoding,
                password=password,
                format=key_serialization_format
                )

def create_self_signed_certificate(
        file_path: str,
        subject_name: str,
        organization_name: str,
        country_name: str,
        valid_days: int = 365,
        key_size: int = 2048
):
    """
    Create a self-signed certificate. It saves the certificate and private key
    in PEM format to the specified file path. Three files are created:

    - <file_path>.crt : The public certificate
    - <file_path>.key : The private key
    - <file_path>.pem : The certificate and private key combined in one file

    :param file_path: Base file path to save the certificate and key.
    :param subject_name: Common Name (CN) for the certificate.
    :param organization_name: Organization Name (O) for the certificate.
    :param country_name: Country Name (C) for the certificate.
    :param valid_days: Number of days the certificate is valid for.
    :param key_size: Size of the RSA key.

    Use the following command to replace any credentials already defined for the 
    App Registration with that certificate:
    
    az ad app credential reset --id <CLIENT_ID> --cert @<file_path>.crt

    Use --append to add the certificate without removing existing credentials.

    > Note: Do not upload the private key file (.key) nor the combined PEM file (.pem).
    """
    key = rsa.generate_private_key(public_exponent=65537, key_size=key_size)

    subject = issuer = x509.Name([
        x509.oid.NameAttribute(x509.oid.NameOID.COUNTRY_NAME, country_name),
        x509.oid.NameAttribute(x509.oid.NameOID.ORGANIZATION_NAME, organization_name),
        x509.oid.NameAttribute(x509.oid.NameOID.COMMON_NAME, subject_name),
    ])

    cert = (
        x509.CertificateBuilder()
        .subject_name(subject)
        .issuer_name(issuer)
        .public_key(key.public_key())
        .serial_number(x509.random_serial_number())
        .not_valid_before(datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(days=1))
        .not_valid_after(datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(days=valid_days))
        .add_extension(x509.BasicConstraints(ca=True, path_length=None), critical=True)
        .sign(key, hashes.SHA256())
    )

    crt_path = pathlib.Path(file_path).with_suffix('.crt')
    key_path = pathlib.Path(file_path).with_suffix('.key')
    pem_path = pathlib.Path(file_path).with_suffix('.pem')

    public_key = cert.public_bytes(serialization.Encoding.PEM)
    private_key = key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.TraditionalOpenSSL,
        encryption_algorithm=serialization.NoEncryption(),
    )

    # Write the certificate
    write_pem_file(crt_path, cert=cert)
    # Write the private key
    write_pem_file(key_path, key=key)
    # Write both to a combined PEM file
    write_pem_file(pem_path, cert=cert, key=key)

    print(f"Certificate created and saved.")
    print(f" - Certificate: {crt_path}")
    print(f" - Private Key: {key_path}")
    print(f" - PEM File: {pem_path}")