import os
from datetime import datetime

from OpenSSL.crypto import PKey, X509, TYPE_RSA

TYPE_RSA = TYPE_RSA


def create_key_pair(bits: int) -> PKey:
    """
    Create a public/private key pair.

    Args:
        bits: Number of bits to use in the key

    Returns:
        The public/private key pair in a PKey object

    """
    pkey = PKey()
    pkey.generate_key(TYPE_RSA, bits)
    return pkey


def create_certificate(key_pair: PKey, digest: str = 'sha256') -> X509:
    """
    Generate a certificate given a certificate request.

    Args:
        key_pair: A public/private key pair generated by OpenSSL.crypto.
        digest: Digest method to use for signing. Default is sha256.

    Returns:
        The signed certificate in an X509 object.

    """
    current_year = datetime.utcnow().strftime('%Y')
    next_year = str(int(current_year) + 1)
    now = current_year + datetime.utcnow().strftime('%m%d%H%M%SZ')
    expires = next_year + datetime.utcnow().strftime('%m%d%H%M%SZ')

    cert = X509()
    cert.set_notBefore(now.encode('utf-8'))
    cert.set_notAfter(expires.encode('utf-8'))
    cert.set_pubkey(key_pair)
    cert.set_serial_number(get_serial_number())
    cert.get_subject().C = 'US'
    cert.get_subject().ST = 'Colorado'
    cert.get_subject().L = 'Colorado'
    cert.get_subject().O = 'TDA Client'
    cert.get_subject().CN = 'localhost'
    cert.set_issuer(cert.get_subject())
    cert.sign(key_pair, digest)
    return cert


def get_serial_number():
    """
    Checks a for a version file in the source code and increments for a new certificate. (Solves
    SEC_ERROR_REUSED_ISSUER_AND_SERIAL exception during access token exchange).

    Returns:
        Serial number to use for the new certificate

    """
    version_file_path = os.path.join(os.path.dirname(__file__), 'certificate_version.txt')
    if os.path.exists(version_file_path):
        with open(version_file_path, 'r') as f:
            content = f.read()
            version = int(content.split(' ')[1])
    else:
        version = 10000
    with open(version_file_path, 'w') as f:
        f.write(f'version: {version + 1}')

    return version
