import os
import hashlib
from cryptography.hazmat.primitives.ciphers.aead import AESGCM


# ===== Key (same idea as Node) =====
def get_key():
    raw_key = os.environ["TOKEN_KEY_ENCRYPTION"].encode()
    return hashlib.sha256(raw_key).digest()  # 32 bytes


# ===== Encrypt =====
def encrypt_token_to_json(token: str) -> dict:
    key = get_key()
    aesgcm = AESGCM(key)

    iv = os.urandom(12)  # same as crypto.randomBytes(12)

    encrypted = aesgcm.encrypt(iv, token.encode(), None)

    ciphertext = encrypted[:-16]
    tag = encrypted[-16:]

    return {
        "iv": iv.hex(),
        "data": ciphertext.hex(),
        "tag": tag.hex(),
    }


# ===== Decrypt =====
def decrypt_token_from_json(enc: dict) -> str:
    key = get_key()
    aesgcm = AESGCM(key)

    iv = bytes.fromhex(enc["iv"])
    ciphertext = bytes.fromhex(enc["data"])
    tag = bytes.fromhex(enc["tag"])

    encrypted = ciphertext + tag

    decrypted = aesgcm.decrypt(iv, encrypted, None)

    return decrypted.decode()