import base64
import os

from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.backends import default_backend
from cryptography.fernet import Fernet, InvalidToken


def decrypt(password: str, token: bytes) -> bytes | None:
    """Decrypts a token using a password."""
    try:
        salt = token[:16]
        encrypted_data = token[16:]
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=480000,
            backend=default_backend()
        )
        key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
        f = Fernet(key)
        return f.decrypt(encrypted_data)
    except InvalidToken:
        return None

def decrypt_file(password: str, input_path: str, output_path: str):
    """Reads an encrypted file, decrypts it, and saves the original content."""
    try:
        with open(input_path, 'rb') as f_in:
            encrypted_data = f_in.read()

        decrypted_data = decrypt(password, encrypted_data)

        if decrypted_data is None:
            print("❌ Decryption failed! Wrong password or corrupted file.")
            return

        with open(output_path, 'wb') as f_out:
            f_out.write(decrypted_data)
        print(f"✅ File '{input_path}' decrypted successfully to '{output_path}'.")
    except FileNotFoundError:
        print(f"❌ Error: Input file not found at '{input_path}'.")
    except Exception as e:
        print(f"❌ An unexpected error occurred during decryption: {e}")

def decrypt_system_prompts() -> bool:
    """Decrypts system prompts stored in a dictionary."""
    password = os.environ.get("PROMPT_PASSWORD")
    if not password:
        print("❌ PROMPT_PASSWORD environment variable not set.")
        return False
    try:
        decrypt_file(password, "./system_prompt.enc", "./system_prompt.txt")
        decrypt_file(password, "./system_prompt_thinking.enc", "./system_prompt_thinking.txt")
        return True
    except Exception as e:
        print(f"❌ An error occurred while decrypting system prompts: {e}")
        return False