from Cryptodome.Cipher import AES
from Cryptodome.Protocol.KDF import scrypt
from Cryptodome.Random import get_random_bytes


def encrypt_aes(data: str, password: str) -> str:
    """
    Encrypt data using password.
    Returns: salt (16) + nonce (16) + tag (16) + ciphertext
    """
    bytes_data = data.encode("utf-8")

    salt = get_random_bytes(16)

    key = scrypt(password.encode("utf-8"), salt, key_len=32, N=2**14, r=8, p=1)

    cipher = AES.new(key, AES.MODE_GCM)

    ciphertext, tag = cipher.encrypt_and_digest(bytes_data)

    return (salt + cipher.nonce + tag + ciphertext).hex()


def decrypt_aes(encrypted_data: str, password: str) -> str:
    """
    Decrypt data using password.
    Raises ValueError if the password is incorrect or data is tampered.
    """
    edata = bytes.fromhex(encrypted_data)
    salt = edata[:16]
    nonce = edata[16:32]
    tag = edata[32:48]
    ciphertext = edata[48:]

    key = scrypt(password.encode("utf-8"), salt, key_len=32, N=2**14, r=8, p=1)

    cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)

    plaintext = cipher.decrypt_and_verify(ciphertext, tag)
    return plaintext.decode("utf-8")
