import os
import base64
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.fernet import Fernet

def generate_key_from_string(key_str: str, salt: bytes) -> bytes:
    """從字串生成一個32字節的加密密鑰"""
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100000,
    )
    return base64.urlsafe_b64encode(kdf.derive(key_str.encode()))

def encrypt_data(data: str) -> str:
    """使用環境變數中的密鑰加密字串資料"""
    key_str = os.getenv('ENCRYPTION_KEY')
    if key_str is None:
        raise ValueError("ENCRYPTION_KEY environment variable is not set")

    salt = os.urandom(16)
    key = generate_key_from_string(key_str, salt)
    cipher_suite = Fernet(key)
    
    encrypted_data = cipher_suite.encrypt(data.encode('utf-8'))
    
    # 返回加密字串和鹽值，因為解密時需要相同的鹽值
    return base64.urlsafe_b64encode(salt + encrypted_data).decode('utf-8')

def decrypt_data(encrypted_data: str) -> str:
    """使用環境變數中的密鑰解密字串資料"""
    key_str = os.getenv('ENCRYPTION_KEY')
    if key_str is None:
        raise ValueError("ENCRYPTION_KEY environment variable is not set")

    decoded_data = base64.urlsafe_b64decode(encrypted_data.encode('utf-8'))
    
    salt = decoded_data[:16]  # 提取前16字節作為鹽值
    encrypted_data = decoded_data[16:]  # 剩餘部分是加密的資料
    
    key = generate_key_from_string(key_str, salt)
    cipher_suite = Fernet(key)
    
    return cipher_suite.decrypt(encrypted_data).decode('utf-8')