import base64
import hashlib
import os
from hashlib import sha384
from credentials cimport Credentials
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import padding

BUFFER_SIZE = 64 * 1024  # 64 KB

cdef class Security:
    @staticmethod
    cdef encrypt_to(input_bytes, key):
        cdef bytes aes_key = hashlib.sha256(key.encode('utf-8')).digest()
        iv = os.urandom(16)

        cipher = Cipher(algorithms.AES(<bytes> aes_key), modes.CBC(iv), backend=default_backend())
        encryptor = cipher.encryptor()
        padder = padding.PKCS7(128).padder()

        padded_plaintext = padder.update(input_bytes) + padder.finalize()
        ciphertext = encryptor.update(padded_plaintext) + encryptor.finalize()

        return iv + ciphertext

    @staticmethod
    cdef decrypt_to(ciphertext_with_iv_bytes, key):
        cdef bytes aes_key = hashlib.sha256(key.encode('utf-8')).digest()
        iv = ciphertext_with_iv_bytes[:16]
        ciphertext_bytes = ciphertext_with_iv_bytes[16:]

        cipher = Cipher(algorithms.AES(<bytes>aes_key), modes.CBC(<bytes>iv), backend=default_backend())
        decryptor = cipher.decryptor()

        decrypted_padded_bytes = decryptor.update(ciphertext_bytes) + decryptor.finalize()

        # Manual PKCS7 unpadding check and removal
        padding_value = decrypted_padded_bytes[-1]  # Get the last byte, which indicates padding length
        if 1 <= padding_value <= 16: # Valid PKCS7 padding value range for AES-128
            padding_length = padding_value
            plaintext_bytes = decrypted_padded_bytes[:-padding_length] # Remove padding bytes
        else:
            plaintext_bytes = decrypted_padded_bytes

        return bytes(plaintext_bytes)

    @staticmethod
    cdef get_hw_hash(str hardware):
        cdef str key = f'Azaion_{hardware}_%$$$)0_'
        return Security.calc_hash(key)

    @staticmethod
    cdef get_api_encryption_key(Credentials creds, str hardware_hash):
        cdef str key = f'{creds.email}-{creds.password}-{hardware_hash}-#%@AzaionKey@%#---'
        return Security.calc_hash(key)

    @staticmethod
    cdef get_model_encryption_key():
        cdef str key = '-#%@AzaionKey@%#---234sdfklgvhjbnn'
        return Security.calc_hash(key)

    @staticmethod
    cdef calc_hash(str key):
        str_bytes = key.encode('utf-8')
        hash_bytes = sha384(str_bytes).digest()
        cdef str h = base64.b64encode(hash_bytes).decode('utf-8')
        return h