from google.cloud import kms
from base64 import b64encode, b64decode
from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
from flask import current_app

client = kms.KeyManagementServiceClient()
key_name = client.crypto_key_path('your-project-id', 'your-key-ring', 'your-crypto-key')


def encrypt_api_key(api_key):
    """Encrypts the API key using the latest version of the KEK."""
    dek = get_random_bytes(32)  # AES 256-bit key
    cipher = AES.new(dek, AES.MODE_GCM)
    ciphertext, tag = cipher.encrypt_and_digest(api_key.encode())

    # Encrypt the DEK using the latest version of the Google Cloud KMS key
    encrypt_response = client.encrypt(
        request={'name': key_name, 'plaintext': dek}
    )
    encrypted_dek = encrypt_response.ciphertext

    # Store the version of the key used
    key_version = encrypt_response.name

    return {
        'key_version': key_version,
        'encrypted_dek': b64encode(encrypted_dek).decode('utf-8'),
        'nonce': b64encode(cipher.nonce).decode('utf-8'),
        'tag': b64encode(tag).decode('utf-8'),
        'ciphertext': b64encode(ciphertext).decode('utf-8')
    }


def decrypt_api_key(encrypted_data):
    """Decrypts the API key using the specified key version."""
    key_version = encrypted_data['key_version']
    encrypted_dek = b64decode(encrypted_data['encrypted_dek'])
    nonce = b64decode(encrypted_data['nonce'])
    tag = b64decode(encrypted_data['tag'])
    ciphertext = b64decode(encrypted_data['ciphertext'])

    # Decrypt the DEK using the specified version of the Google Cloud KMS key
    decrypt_response = client.decrypt(
        request={'name': key_version, 'ciphertext': encrypted_dek}
    )
    dek = decrypt_response.plaintext

    cipher = AES.new(dek, AES.MODE_GCM, nonce=nonce)
    api_key = cipher.decrypt_and_verify(ciphertext, tag)
    return api_key.decode()