import os
import oqs
import tarfile
import requests
import rfc3161ng
from OpenSSL import crypto
from base64 import b64decode
from cryptography import x509
from shutil import copy2, rmtree
from pyasn1.type.univ import OctetString
from pyasn1_modules import rfc5280, rfc5208
from cryptography.hazmat.backends import default_backend
from pyasn1.codec.der.decoder import decode as der_decoder
from cryptography.hazmat.primitives import asymmetric, hashes, serialization

from .exceptions import CertificateNotValidCustomException


def store_ts_data(data, path, name):
    """
    Function responsible for writing bits of data
    to given path directory and asigning it a
    name from function.
    """
    with open(os.path.join(path, name), "wb") as f:
        f.write(data)


def load_data(path):
    """
    Loads bytes of data from system path
    """
    with open(path, "rb") as f:
        data = f.read()
    return data


def get_timestamp(tsa_info: dict, file_hash):
    """
    Function responsible for getting timstamp
    for bits of data. It is creating connection to
    remote timestamping authority.
    It needs hash data or bytes of small object
    and it needs dictionary with information to
    create connection.
    dictionary example:
    TSA_info: {
        tsa_tsr_url: 'https://freetsa.org/tsr'
        tsa_cert_path: '/home/server/Downloads/tsa.crt'
    }
    Returns bytes with timestamp tsr
    """
    tsa_cert = load_data(tsa_info["tsa_cert_path"])
    remote_tsa = rfc3161ng.RemoteTimestamper(
        url=tsa_info["tsa_tsr_url"], certificate=tsa_cert, hashname="sha512"
    )
    data_ts = remote_tsa(file_hash)
    return data_ts


def verify_timestamp(ts, data, info: dict):
    """
    Function for verifying timestamp respones.
    It will create connection to timestamping authority
    which need dictionary like this
    TSA_info: {
        tsa_tsr_url: 'https://freetsa.org/tsr'
        tsa_cert_path: '/home/server/Downloads/tsa.crt'
    }
    It need data that was timestamped and bytes with ts
    returns true or false if timestamp is corect
    """
    tsa_cert = load_data(info["tsa_cert_path"])
    remote_tsa = rfc3161ng.RemoteTimestamper(
        url=info["tsa_tsr_url"], certificate=tsa_cert, hashname="sha512"
    )
    return remote_tsa.check(ts, data)


def sign_data(data, private_key):
    """
    Function for signig byte data with given private key.
    """
    signature = private_key.sign(
        data,
        asymmetric.padding.PSS(
            mgf=asymmetric.padding.MGF1(hashes.SHA512()),
            salt_length=asymmetric.padding.PSS.MAX_LENGTH,
        ),
        hashes.SHA512(),
    )
    return signature

def pqc_get_private_key(path: str, sigalg: str) -> bytes:
    """
    Function for retrieving private key bytes from
    file.
    """

    # Defines private key size for algorithm variants
    ALGORITHMS = {
        "Dilithium5": 4864,
        "Dilithium3": 4000,
        "Dilithium2": 2528,
    }

    # Read key from file (assuming no passphrase)
    with open(path) as key_file:
        # Strip first and last line
        b64_serialized = ''.join(key_file.readlines()[1:-1])

    # Remove base64 serialisation from bytes
    der_serialized = b64decode(b64_serialized)

    # Remove DER serialisation from bytes (using PKCS#8 (rfc5208) as ASN.1 template)
    private_key_info, rest_of_input = der_decoder(der_serialized, asn1Spec=rfc5208.PrivateKeyInfo())

    # Remove ASN.1 Octet String wraped around serialized private and public keys
    keys, rest_of_input = der_decoder(bytes(private_key_info["privateKey"]), asn1Spec=OctetString())

    # Split private and public key
    # (Note: Keys object contains serialized private and public key bytes togeather
    # for more detail refer to https://datatracker.ietf.org/doc/draft-uni-qsckeys-dilithium/)
    private_key = bytes(keys)[:ALGORITHMS.get(sigalg, "invalid")]

    return private_key

def pqc_get_public_key(path: str) -> bytes:
    """
    Function for retrieving public key bytes from certificate.
    """

    # Read certificate from file
    with open(path) as key_file:
        # Strip first and last line
        b64_serialized = ''.join(key_file.readlines()[1:-1])

    # Remove base64 serialisation from bytes
    der_serialized = b64decode(b64_serialized)

    # Remove DER serialisation from bytes (using X.509 (rfc5280) as ASN.1 template)
    certificate, rest_of_input = der_decoder(der_serialized, asn1Spec=rfc5280.Certificate())

    # Public key is stored as string of bits in certificate
    public_key_bitstring = str(certificate["tbsCertificate"]["subjectPublicKeyInfo"]["subjectPublicKey"])

    # Convert bitstring to bytes
    public_key = bytes(int(public_key_bitstring[i:i+8], 2) for i in range(0, len(public_key_bitstring), 8))

    return public_key

def pqc_sign_data(hash: bytes, private_key: bytes, sigalg: str) -> bytes:
    """
    Function for signing data with PQC signature.
    """

    signer = oqs.Signature(sigalg, private_key)
    signature = signer.sign(hash)

    return signature

def pqc_validate_signature(hash: bytes, signature: bytes, public_key: bytes, sigalg: str) -> bool:

    signer = oqs.Signature(sigalg)
    valid = signer.verify(hash, signature, public_key)

    return valid

def get_private_key(path, password):
    """
    Function for loading private key from path
    with password
    returning private key
    """
    with open(path, "rb") as pk:
        private_key = serialization.load_pem_private_key(
            pk.read(), password=password.encode(), backend=default_backend()
        )
    return private_key


def get_certificate(path):
    """
    Function for loading PEM certificate from given
    path.
    It will return x509 object with certificate
    """
    with open(path, "rb") as cert:
        crt = cert.read()
        certificate = x509.load_pem_x509_certificate(crt, default_backend())
    return certificate


def get_public_key(path):
    """
    Function for loading stored public key
    from given path
    """
    with open(path, "rb") as pk:
        k = pk.read()
        public_key = serialization.load_pem_public_key(
            k, backend=default_backend()
        )
    return public_key


def store_signature(path, signature):
    """
    Function to store bytes of signature
    to given directory path at gives it name
    signature.sig
    """
    path = os.path.join(path, "signature.sig")
    with open(path, "wb") as f:
        f.write(signature)
    return path


def validate_signature(data, signature, public_key):
    """
    Function for validating signatature of data.
    It need public key, signature and signed bytes data
    returns true if signature is valid else it will raise exception
    """

    public_key.verify(
        signature,
        data,
        asymmetric.padding.PSS(
            mgf=asymmetric.padding.MGF1(hashes.SHA512()),
            salt_length=asymmetric.padding.PSS.MAX_LENGTH,
        ),
        hashes.SHA512(),
    )
    return True


def create_tar_file_from_dir(dir_path, tar_path):
    """
    This will create tar archive, without any compression,
    from files within directory and deletes original files
    to remove duplicits. It will returns path to created
    tar archive
    """
    with tarfile.open(tar_path, "w:") as tarf:  # uncompressed mode
        for f in os.listdir(dir_path):
            fp = os.path.join(dir_path, f)
            if not fp == tar_path:
                tarf.add(fp, arcname=f)
                delete_file(fp)


def delete_file(path):
    """
    This will delete directory or file from
    given path
    """
    if os.path.isdir(path):
        rmtree(path)
    if os.path.isfile(path):
        os.remove(path)


def get_file_hash(hash, file_path):
    """
    Function for hashing file on given path.
    It needs to get specify hash function which
    will be applied (example haslib.sha512)
    It will hash even big files without killing
    ram memmory since it uses buffering.
    Returning bytes with digest
    """
    file_hash = hash()
    with open(file_path, "rb") as f:
        buffer = f.read(8192)
        while buffer != b"":
            file_hash.update(buffer)
            buffer = f.read(8192)
    return file_hash.digest()

def validate_certificate(crl_content, ca_file_path):
    """
    Function for validating if certificate was correct
    at given time. It is validating it with its crl content
    It need crl bytes and path to certificate
    return true if valid
    """
    crl = crypto.load_crl(crypto.FILETYPE_PEM, crl_content)

    crl_crypto = crl.to_cryptography()

    with open(ca_file_path, "rb") as f:
        ca_cert = f.read()

    ca = crypto.load_certificate(crypto.FILETYPE_PEM, ca_cert)

    valid = crl_crypto.is_signature_valid(
        ca.get_pubkey().to_cryptography_key()
    )
    revoked = crl_crypto.get_revoked_certificate_by_serial_number(
        ca.get_serial_number()
    )
    if revoked:
        raise CertificateNotValidCustomException("certificate was revoked")
    if not (valid):
        raise CertificateNotValidCustomException(
            "certificate or crl is invalid"
        )

def get_current_crl(tsa_crl_url):
    """
    Function for downloading crl from tsa_crl_url and returning
    content.
    """
    res = requests.get(tsa_crl_url)
    return res.content


def copy_file_to_dir(file_path_to_copy, dst_dir, name):
    """
    Function for copying file to other directory
    It needs path to file which has to be copied
    and destination directory and name under which
    it will be stored
    """
    dst = os.path.join(dst_dir, name)
    copy2(src=file_path_to_copy, dst=dst)
    return dst


def create_new_dir_in_location(new_dir_path, name):
    """
    Function for creating new directory in given
    directory location with given name
    Parrent directory must exist
    Return path to new directory
    """
    path = os.path.join(new_dir_path, name)
    os.makedirs(path)
    return path