import base64
import hashlib
import os

from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes


class AESCipher:
    def __init__(self, key_string):
        """
        initialize AES cipher
        :param key_string: key string, will be converted to 32 bytes key
        """
        # use SHA-256 to convert arbitrary length string to 32 bytes key
        self.key = hashlib.sha256(key_string.encode()).digest()
        self.backend = default_backend()

    def encrypt(self, plaintext):
        """
        encrypt string
        :param plaintext: string to encrypt
        :return: base64 encoded encrypted string
        """
        # generate random IV (initialization vector)
        iv = os.urandom(16)

        # convert string to bytes
        plaintext_bytes = plaintext.encode("utf-8")

        # create AES-CBC encryptor
        encryptor = Cipher(algorithms.AES(self.key), modes.CBC(iv), backend=self.backend).encryptor()

        # ensure plaintext length is a multiple of 16 (AES block size)
        padder = self._create_padder()
        padded_plaintext = padder.update(plaintext_bytes) + padder.finalize()

        # encrypt
        ciphertext = encryptor.update(padded_plaintext) + encryptor.finalize()

        # combine IV and ciphertext, then base64 encode
        encrypted_data = base64.b64encode(iv + ciphertext).decode("utf-8")

        return encrypted_data

    def decrypt(self, encrypted_data):
        """
        decrypt string
        :param encrypted_data: base64 encoded encrypted string
        :return: decrypted string
        """
        try:
            # base64 decode
            encrypted_bytes: bytes = base64.b64decode(encrypted_data)

            # extract IV (first 16 bytes)
            iv = encrypted_bytes[:16]
            ciphertext = encrypted_bytes[16:]

            # create AES-CBC decryptor
            decryptor = Cipher(algorithms.AES(self.key), modes.CBC(iv), backend=self.backend).decryptor()

            # decrypt
            padded_plaintext = decryptor.update(ciphertext) + decryptor.finalize()

            # remove padding
            unpadder = self._create_unpadder()
            plaintext_bytes = unpadder.update(padded_plaintext) + unpadder.finalize()

            # convert to string
            plaintext = plaintext_bytes.decode("utf-8")

            return plaintext
        except Exception as e:
            raise ValueError(f"decrypt failed: {str(e)}")

    def _create_padder(self):
        """create PKCS7 padder"""
        from cryptography.hazmat.primitives.padding import PKCS7

        return PKCS7(algorithms.AES.block_size).padder()

    def _create_unpadder(self):
        """create PKCS7 unpadder"""
        from cryptography.hazmat.primitives.padding import PKCS7

        return PKCS7(algorithms.AES.block_size).unpadder()


if __name__ == "__main__":
    # this is a Cipher test
    cipher = AESCipher("kb_ai_agent")

    # text for test
    # plaintext = "This is plain text"
    # encrypted = cipher.encrypt(plaintext)
    # print(f"加密后: {encrypted}")

    # decrypt
    decrypted = cipher.decrypt(
        "PX784ZDszIUME3nd9G4I+S+PlojNEfDbV4OwGX2O9+0="
    )
    print(f"decrypted: {decrypted}")
