"""
A multi-layer encryption system with key management, substitution ciphers,
XOR operations, and dynamic key derivation.
"""

import base64
import json
import os
import time
from datetime import datetime


class EncryptionSystem:
    """Main encryption engine with key-based transformation layers."""

    def __init__(self, key_file=None, key_password=None):
        """
        Initialize the encryption system with an optional key file.

        Args:
            key_file: Path to the encryption key file.
            key_password: Password for decrypting the key file.
        """
        self.key_file = key_file
        self.key_password = key_password
        self.keys_loaded = False

        # Base64 special character mappings for safe transport
        self.special_encrypt = {
            '+': '.',
            '/': "'",
            '=': ''
        }
        self.special_decrypt = {
            '.': '+',
            "'": '/'
        }

        self.flip_pattern = None

        # Try loading the key file if provided or find default
        if key_file:
            if os.path.exists(key_file):
                if key_password is None:
                    key_password = input(f"Enter password for key file {key_file}: ")
                self._load_keys(key_file, key_password)
                self.keys_loaded = True
            else:
                print(f"⚠️ Key file {key_file} not found")
                self.keys_loaded = False
        else:
            default_key = "encryption.key"
            if os.path.exists(default_key):
                self.key_file = default_key
                if key_password is None:
                    key_password = input(f"Enter password for key file {default_key}: ")
                self._load_keys(default_key, key_password)
                self.keys_loaded = True
            else:
                print("=" * 60)
                print("First launch detected. Please generate a key file first.")
                print("=" * 60)
                self.keys_loaded = False

    # ----------------------------------------------------------------------
    # Key generation utilities
    # ----------------------------------------------------------------------

    def _generate_random_alphabet(self, lowercase=False):
        """Generate a shuffled alphabet string."""
        chars = list("abcdefghijklmnopqrstuvwxyz" if lowercase else "ABCDEFGHIJKLMNOPQRSTUVWXYZ")
        n = len(chars)
        # Fisher-Yates shuffle using secure random bytes
        for i in range(n - 1, 0, -1):
            j = int.from_bytes(os.urandom(1), 'big') % (i + 1)
            chars[i], chars[j] = chars[j], chars[i]
        return ''.join(chars)

    def _generate_random_digit_mapping(self):
        """Create a random mapping for digits 0-9 to alphabet characters."""
        digits = list("0123456789")
        mapping_chars = list("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz")
        n = len(mapping_chars)
        for i in range(n - 1, 0, -1):
            j = int.from_bytes(os.urandom(1), 'big') % (i + 1)
            mapping_chars[i], mapping_chars[j] = mapping_chars[j], mapping_chars[i]
        mapping = {}
        for i, d in enumerate(digits):
            mapping[d] = mapping_chars[i]
        return mapping

    def _generate_random_equal_mapping(self):
        """Generate mapping for Base64 padding count (0-3)."""
        chars = list("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz")
        n = len(chars)
        for i in range(n - 1, 0, -1):
            j = int.from_bytes(os.urandom(1), 'big') % (i + 1)
            chars[i], chars[j] = chars[j], chars[i]
        return {
            '0': chars[0],
            '1': chars[1],
            '2': chars[2],
            '3': chars[3]
        }

    def _generate_random_flip_pattern(self):
        """Generate a 10-bit pattern for case-flipping."""
        return [int.from_bytes(os.urandom(1), 'big') % 2 for _ in range(10)]

    def _generate_long_key(self, length=4096):
        """Generate a long hex key (4096 chars by default)."""
        chars = "0123456789abcdef"
        return ''.join(chars[int.from_bytes(os.urandom(1), 'big') % 16] for _ in range(length))

    def _generate_short_key(self, length=512):
        """Generate a short hex key (512 chars by default)."""
        chars = "0123456789abcdef"
        return ''.join(chars[int.from_bytes(os.urandom(1), 'big') % 16] for _ in range(length))

    # ----------------------------------------------------------------------
    # Key obfuscation and persistence
    # ----------------------------------------------------------------------

    def _obfuscate_keys(self, keys_data):
        """Obfuscate key data using base64 + reversal + Caesar shift."""
        json_str = json.dumps(keys_data)
        b64 = base64.b64encode(json_str.encode()).decode()
        reversed_b64 = b64[::-1]
        shifted = ''.join(chr((ord(c) + 1) % 128) for c in reversed_b64)
        return base64.b64encode(shifted.encode()).decode()

    def _deobfuscate_keys(self, obfuscated_data):
        """Reverse the obfuscation to recover original key data."""
        try:
            shifted = base64.b64decode(obfuscated_data.encode()).decode()
            reversed_b64 = ''.join(chr((ord(c) - 1) % 128) for c in shifted)
            b64 = reversed_b64[::-1]
            json_str = base64.b64decode(b64.encode()).decode()
            return json.loads(json_str)
        except Exception as exc:
            raise Exception("Key file corrupted") from exc

    def _xor_encrypt_data(self, data, password):
        """XOR encrypt data with a password, returning hex string."""
        key_len = len(password)
        return ''.join(
            f"{ord(char) ^ ord(password[i % key_len]):02x}"
            for i, char in enumerate(data)
        )

    def _xor_decrypt_data(self, hex_data, password):
        """XOR decrypt hex data with a password."""
        try:
            key_len = len(password)
            result = []
            for i in range(0, len(hex_data), 2):
                if i + 1 < len(hex_data):
                    hex_byte = hex_data[i:i+2]
                    xor_result = int(hex_byte, 16) ^ ord(password[(i//2) % key_len])
                    result.append(chr(xor_result))
            return ''.join(result)
        except Exception:
            return None

    def generate_keys(self, save_path=None, key_password=None):
        """
        Generate a fresh set of encryption keys and save to a key file.

        Args:
            save_path: Path to save the key file (default: encryption.key).
            key_password: Password to protect the key file.

        Returns:
            The password used, or None if generation failed.
        """
        if key_password is None:
            key_password = input("Set key file password: ")
            confirm = input("Confirm password: ")
            if key_password != confirm:
                print("❌ Passwords do not match")
                return None

        print("=" * 60)
        print("Generating random keys...")
        print("=" * 60)

        # Generate all key components
        self.upper_mapping = self._generate_random_alphabet(lowercase=False)
        self.lower_mapping = self._generate_random_alphabet(lowercase=True)
        self.digit_mapping = self._generate_random_digit_mapping()
        self.equal_mapping = self._generate_random_equal_mapping()
        self.flip_pattern = self._generate_random_flip_pattern()
        self.long_key = self._generate_long_key(4096)
        self.short_key = self._generate_short_key(512)

        # Build reverse mappings
        self.digit_reverse = {v: k for k, v in self.digit_mapping.items()}
        self.equal_reverse = {v: k for k, v in self.equal_mapping.items()}

        # Bundle keys into a dictionary
        keys_data = {
            'upper_mapping': self.upper_mapping,
            'lower_mapping': self.lower_mapping,
            'digit_mapping': self.digit_mapping,
            'equal_mapping': self.equal_mapping,
            'flip_pattern': self.flip_pattern,
            'long_key': self.long_key,
            'short_key': self.short_key,
            'generated_at': datetime.now().isoformat()
        }

        # Obfuscate and encrypt the key data
        obfuscated = self._obfuscate_keys(keys_data)
        encrypted = self._xor_encrypt_data(obfuscated, key_password)

        if save_path is None:
            save_path = "encryption.key"

        with open(save_path, 'w') as f:
            f.write(encrypted)

        self.key_file = save_path
        self.key_password = key_password
        self.keys_loaded = True
        self._build_maps()

        print(f"✅ Keys generated and saved to: {save_path}")
        print("=" * 60)

        return save_path

    # ----------------------------------------------------------------------
    # Internal helpers
    # ----------------------------------------------------------------------

    def _build_maps(self):
        """Build encryption and decryption maps from shuffled alphabets."""
        self.upper_original = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
        self.lower_original = "abcdefghijklmnopqrstuvwxyz"

        self.encrypt_map = {}
        self.decrypt_map = {}
        for i in range(26):
            self.encrypt_map[self.upper_original[i]] = self.upper_mapping[i]
            self.encrypt_map[self.lower_original[i]] = self.lower_mapping[i]
            self.decrypt_map[self.upper_mapping[i]] = self.upper_original[i]
            self.decrypt_map[self.lower_mapping[i]] = self.lower_original[i]

    def _load_keys(self, key_file, key_password):
        """Load and decrypt keys from a key file."""
        try:
            with open(key_file, 'r') as f:
                encrypted_data = f.read()

            decrypted = self._xor_decrypt_data(encrypted_data, key_password)
            if decrypted is None:
                print("❌ Incorrect password")
                self.keys_loaded = False
                return

            keys_data = self._deobfuscate_keys(decrypted)

            self.upper_mapping = keys_data['upper_mapping']
            self.lower_mapping = keys_data['lower_mapping']
            self.digit_mapping = keys_data['digit_mapping']
            self.equal_mapping = keys_data['equal_mapping']
            self.flip_pattern = keys_data['flip_pattern']
            self.long_key = keys_data['long_key']
            self.short_key = keys_data['short_key']

            self.digit_reverse = {v: k for k, v in self.digit_mapping.items()}
            self.equal_reverse = {v: k for k, v in self.equal_mapping.items()}

            self.key_password = key_password
            self._build_maps()

            self.keys_loaded = True
            print(f"✅ Keys loaded from {key_file}")

        except Exception:
            print("❌ Loading failed")
            self.keys_loaded = False

    def _ensure_printable(self, text):
        """Force all characters into printable ASCII range (32-126)."""
        result = []
        for c in text:
            val = ord(c)
            if val == 124:  # '|' is used as a separator, keep it
                result.append('|')
            elif val < 32 or val > 126:
                val = val % 95 + 32
                result.append(chr(val))
            else:
                result.append(c)
        return ''.join(result)

    # ----------------------------------------------------------------------
    # Key derivation
    # ----------------------------------------------------------------------

    def _mix_keys(self, user_password, target_length):
        """
        Mix the user password with the short key using bit operations and swaps.

        Args:
            user_password: User-provided password.
            target_length: Desired output length.

        Returns:
            A mixed printable string.
        """
        if not self.keys_loaded:
            raise Exception("Keys not loaded")

        # Interleave user password and short key
        mixed = []
        max_len = max(len(user_password), len(self.short_key))
        for i in range(max_len):
            if i < len(user_password):
                mixed.append(ord(user_password[i]))
            if i < len(self.short_key):
                mixed.append(ord(self.short_key[i]))

        # Apply bit transformations
        for i in range(len(mixed)):
            if i % 3 == 0:
                mixed[i] = (mixed[i] << 1) & 0xFF
            elif i % 3 == 1:
                mixed[i] = (mixed[i] >> 1) & 0xFF
            else:
                mixed[i] = mixed[i] ^ 0x5A

        # Swap pairs
        for i in range(0, len(mixed) - 3, 4):
            mixed[i], mixed[i+3] = mixed[i+3], mixed[i]
            mixed[i+1], mixed[i+2] = mixed[i+2], mixed[i+1]

        mixed.reverse()

        # Convert to printable characters
        result = []
        for x in mixed:
            if x < 32 or x > 126:
                x = x % 95 + 32
            result.append(chr(x))

        result_str = ''.join(result)

        # Extend if needed by repeating with variations
        if len(result_str) < target_length:
            base_key = result_str
            final_key = base_key
            iteration = 0
            while len(final_key) < target_length:
                if iteration % 3 == 0:
                    next_chunk = base_key[::-1]
                elif iteration % 3 == 1:
                    next_chunk = base_key[::-1]
                    next_chunk = ''.join(chr((ord(c) + iteration) % 95 + 32) for c in next_chunk)
                else:
                    next_chunk = base_key[::-1]
                    next_chunk = ''.join(chr((ord(c) ^ iteration) % 95 + 32) for c in next_chunk)
                final_key += next_chunk
                iteration += 1
            result_str = final_key[:target_length]

        return self._ensure_printable(result_str)

    def _derive_final_key(self, user_password, text_length):
        """
        Derive a final encryption key from the user password and short key.

        The process includes mixing, hex conversion, reversal, and extension.
        """
        mixed_key = self._mix_keys(user_password, text_length * 2)

        # Convert to hex and reverse
        hex_key = ''.join(f"{ord(c):02x}" for c in mixed_key)
        reversed_hex = hex_key[::-1]

        # Convert back to printable characters
        final_key = []
        for i in range(0, len(reversed_hex), 2):
            if i + 1 < len(reversed_hex):
                hex_byte = reversed_hex[i:i+2]
                try:
                    val = int(hex_byte, 16)
                    if val < 32 or val > 126:
                        val = val % 95 + 32
                    final_key.append(chr(val))
                except ValueError:
                    final_key.append('x')

        final_key_str = ''.join(final_key)

        # Extend if needed
        if len(final_key_str) < text_length:
            base_key = final_key_str
            iteration = 0
            while len(final_key_str) < text_length:
                if iteration % 3 == 0:
                    next_chunk = base_key[::-1]
                elif iteration % 3 == 1:
                    next_chunk = base_key[::-1]
                    next_chunk = ''.join(chr((ord(c) + iteration) % 95 + 32) for c in next_chunk)
                else:
                    next_chunk = base_key[::-1]
                    next_chunk = ''.join(chr((ord(c) ^ iteration) % 95 + 32) for c in next_chunk)
                final_key_str += next_chunk
                iteration += 1

        final_key_str = final_key_str[:text_length]
        return self._ensure_printable(final_key_str)

    # ----------------------------------------------------------------------
    # Core encryption operations
    # ----------------------------------------------------------------------

    def _apply_flip(self, text):
        """Apply case-flipping based on the flip pattern."""
        if self.flip_pattern is None:
            return text

        result = []
        for i, ch in enumerate(text):
            idx = i % len(self.flip_pattern)
            if ch.isalpha() and self.flip_pattern[idx] == 1:
                result.append(ch.swapcase())
            else:
                result.append(ch)
        return ''.join(result)

    def _xor_encrypt_with_final_key(self, text, user_password):
        """XOR encrypt with the derived final key, returning hex."""
        final_key = self._derive_final_key(user_password, len(text))
        key_len = len(final_key)
        return ''.join(
            f"{ord(char) ^ ord(final_key[i % key_len]):02x}"
            for i, char in enumerate(text)
        )

    def _xor_decrypt_with_final_key(self, hex_text, user_password):
        """XOR decrypt with the derived final key."""
        try:
            text_length = len(hex_text) // 2
            final_key = self._derive_final_key(user_password, text_length)
            key_len = len(final_key)
            result = []
            for i in range(0, len(hex_text), 2):
                if i + 1 < len(hex_text):
                    hex_byte = hex_text[i:i+2]
                    try:
                        xor_result = int(hex_byte, 16) ^ ord(final_key[(i//2) % key_len])
                        if xor_result < 32 or xor_result > 126:
                            xor_result = xor_result % 95 + 32
                        result.append(chr(xor_result))
                    except ValueError:
                        result.append('?')
            return ''.join(result)
        except Exception:
            return None

    def _substitute_letters(self, text, mapping):
        """Apply a substitution map to letters."""
        return ''.join(mapping.get(char, char) for char in text)

    def _encode_digit(self, num_str):
        """Encode a digit string using the digit mapping."""
        return ''.join(self.digit_mapping.get(char, char) for char in num_str)

    def _decode_digit(self, encoded_str):
        """Decode a digit string using the reverse digit mapping."""
        return ''.join(self.digit_reverse.get(char, char) for char in encoded_str)

    def _get_dynamic_key(self, text_length):
        """Generate a dynamic key from the long key, repeated as needed."""
        key = self.long_key
        if len(key) < text_length:
            repeats = (text_length // len(key)) + 1
            key = (key * repeats)[:text_length]
        return key[:text_length]

    def _xor_encrypt_with_key(self, text, key):
        """XOR encrypt with a fixed key, returning hex."""
        key_len = len(key)
        return ''.join(
            f"{ord(char) ^ ord(key[i % key_len]):02x}"
            for i, char in enumerate(text)
        )

    def _xor_decrypt_with_key(self, hex_text, key):
        """XOR decrypt with a fixed key."""
        try:
            key_len = len(key)
            result = []
            for i in range(0, len(hex_text), 2):
                if i + 1 < len(hex_text):
                    hex_byte = hex_text[i:i+2]
                    try:
                        xor_result = int(hex_byte, 16) ^ ord(key[(i//2) % key_len])
                        if xor_result < 32 or xor_result > 126:
                            xor_result = xor_result % 95 + 32
                        result.append(chr(xor_result))
                    except ValueError:
                        result.append('?')
            return ''.join(result)
        except Exception:
            return None

    # ----------------------------------------------------------------------
    # Public API
    # ----------------------------------------------------------------------

    def encrypt(self, plaintext, user_password):
        """
        Encrypt plaintext with the given password.

        Workflow:
        1. Base64 encode the plaintext.
        2. Apply letter substitution and case-flipping.
        3. Reverse the string and append padding count.
        4. XOR encrypt with dynamic key (long key).
        5. XOR encrypt with derived final key (user password + short key).

        Returns:
            Encrypted ciphertext as a hex string.
        """
        if not self.keys_loaded:
            return "❌ Error: Keys not loaded"

        # Step 1: Base64 encode and handle special chars
        b64 = base64.b64encode(plaintext.encode('utf-8')).decode('utf-8')
        equal_count = b64.count('=')

        processed = b64
        for old, new in self.special_encrypt.items():
            processed = processed.replace(old, new)
        processed = processed.rstrip('=')

        # Step 2: Letter substitution + flip + reverse
        sub = self._substitute_letters(processed, self.encrypt_map)
        flipped = self._apply_flip(sub)
        reversed_text = flipped[::-1]

        # Step 3: Append padding count
        equal_char = self.equal_mapping[str(equal_count)]
        with_equal = f"{reversed_text}|{equal_char}"

        # Step 4: Dynamic key (long key) encryption
        key_length_str = str(len(with_equal))
        key_length_encoded = self._encode_digit(key_length_str)
        dynamic_key = self._get_dynamic_key(len(with_equal))
        encrypted_by_dynamic = self._xor_encrypt_with_key(with_equal, dynamic_key)

        # Step 5: Final encryption with derived key
        combined = f"{encrypted_by_dynamic}|{key_length_encoded}"
        final_encrypted = self._xor_encrypt_with_final_key(combined, user_password)

        return final_encrypted

    def decrypt(self, ciphertext, user_password):
        """
        Decrypt ciphertext with the given password.

        Returns:
            The original plaintext, or an error message on failure.
        """
        if not self.keys_loaded:
            return "❌ Decryption failed"

        try:
            # Step 1: Decrypt with derived final key
            combined = self._xor_decrypt_with_final_key(ciphertext, user_password)
            if combined is None:
                return "❌ Decryption failed"

            # Step 2: Extract hex data and length indicator
            if '|' in combined:
                parts = combined.split('|')
                if len(parts) >= 2:
                    hex_data = parts[0]
                    key_length_encoded = parts[1]
                else:
                    hex_data = parts[0]
                    key_length_encoded = 'g'
            else:
                hex_data = combined
                key_length_encoded = 'g'

            # Step 3: Decode length and get dynamic key
            key_length_str = self._decode_digit(key_length_encoded)
            try:
                key_length = int(key_length_str)
            except ValueError:
                key_length = 16

            dynamic_key = self._get_dynamic_key(key_length)

            # Step 4: Decrypt with dynamic key
            xor_decrypted = self._xor_decrypt_with_key(hex_data, dynamic_key)
            if xor_decrypted is None:
                return "❌ Decryption failed"

            # Step 5: Extract main data and padding count
            if '|' in xor_decrypted:
                main_part, equal_char = xor_decrypted.split('|')
                equal_count = int(self.equal_reverse.get(equal_char, '0'))
            else:
                main_part = xor_decrypted
                equal_count = 0

            # Step 6: Reverse, flip, substitute
            reversed_text = main_part[::-1]
            flipped = self._apply_flip(reversed_text)
            sub = self._substitute_letters(flipped, self.decrypt_map)

            # Step 7: Restore Base64 special chars and padding
            for old, new in self.special_decrypt.items():
                sub = sub.replace(old, new)

            b64_with_equal = sub + '=' * equal_count

            if len(b64_with_equal) % 4 != 0:
                return "❌ Decryption failed"

            # Step 8: Base64 decode
            decoded = base64.b64decode(b64_with_equal.encode('utf-8')).decode('utf-8')
            return decoded

        except Exception:
            return "❌ Decryption failed"

    def print_keys(self):
        """Display information about the currently loaded keys."""
        if not self.keys_loaded:
            print("❌ Keys not loaded")
            return

        print("=" * 60)
        print("Key Information")
        print("=" * 60)
        print(f"Key file: {self.key_file}")
        print(f"Uppercase mapping: {self.upper_mapping}")
        print(f"Lowercase mapping: {self.lower_mapping}")
        print(f"Flip pattern: {self.flip_pattern}")
        print(f"Long key length: {len(self.long_key)} bits")
        print(f"Short key length: {len(self.short_key)} bits")
        print("=" * 60)


# ----------------------------------------------------------------------
# CLI Entry Point
# ----------------------------------------------------------------------

def main():
    """Command-line interface for the encryption system."""
    print("=" * 60)
    print("Welcome to the Encryption System")
    print("=" * 60)

    crypto = None
    key_password = None

    # Try loading default key file
    if os.path.exists("encryption.key"):
        key_password = input("Enter password for default key file (encryption.key): ")
        crypto = EncryptionSystem("encryption.key", key_password)
        if not crypto.keys_loaded:
            print("❌ Loading failed")
            crypto = None
    else:
        print("\nDefault key file not found. Please generate one first.")
        choice = input("Generate default key file? (y/n): ").strip().lower()
        if choice == 'y':
            crypto = EncryptionSystem()
            result = crypto.generate_keys("encryption.key")
            if result is None:
                print("❌ Generation failed")
                crypto = None
        else:
            print("⚠️ Use options 4, 5, or 6 to manage key files manually.")

    # Main interaction loop
    while True:
        print("\nSelect an option:")
        print("1. Encrypt with default key")
        print("2. Decrypt with default key")
        print("3. Generate new key (overwrite default)")
        print("4. Encrypt with custom key file")
        print("5. Decrypt with custom key file")
        print("6. Generate key and save to current folder")
        print("7. View current key info")
        print("8. Exit")

        choice = input("\nEnter choice (1-8): ").strip()

        if choice == '1':
            # Encrypt with default key
            if crypto is None or not crypto.keys_loaded:
                if os.path.exists("encryption.key"):
                    if key_password is None:
                        key_password = input("Enter default key file password: ")
                    crypto = EncryptionSystem("encryption.key", key_password)
                    if not crypto.keys_loaded:
                        print("❌ Loading failed")
                        continue
                else:
                    print("❌ Key file not found")
                    continue

            password = input("Enter encryption password: ")
            text = input("Enter text to encrypt: ")
            if text and password:
                encrypted = crypto.encrypt(text, password)
                print(f"\n✅ Encrypted result: {encrypted}")

        elif choice == '2':
            # Decrypt with default key
            if crypto is None or not crypto.keys_loaded:
                if os.path.exists("encryption.key"):
                    if key_password is None:
                        key_password = input("Enter default key file password: ")
                    crypto = EncryptionSystem("encryption.key", key_password)
                    if not crypto.keys_loaded:
                        print("❌ Loading failed")
                        continue
                else:
                    print("❌ Key file not found")
                    continue

            password = input("Enter encryption password: ")
            text = input("Enter ciphertext to decrypt: ")
            if text and password:
                decrypted = crypto.decrypt(text, password)
                print(f"\n✅ Decrypted result: {decrypted}")

        elif choice == '3':
            # Generate and overwrite default key
            crypto = EncryptionSystem()
            result = crypto.generate_keys("encryption.key")
            if result is None:
                print("❌ Generation failed")
            else:
                print("✅ Default key updated")

        elif choice == '4':
            # Encrypt with custom key file
            key_file = input("Enter key file path: ")
            if not os.path.exists(key_file):
                print(f"❌ File not found: {key_file}")
                continue

            kp = input("Enter key file password: ")
            crypto = EncryptionSystem(key_file, kp)
            if not crypto.keys_loaded:
                print("❌ Loading failed")
                continue

            password = input("Enter encryption password: ")
            text = input("Enter text to encrypt: ")
            if text and password:
                encrypted = crypto.encrypt(text, password)
                print(f"\n✅ Encrypted result: {encrypted}")

        elif choice == '5':
            # Decrypt with custom key file
            key_file = input("Enter key file path: ")
            if not os.path.exists(key_file):
                print(f"❌ File not found: {key_file}")
                continue

            kp = input("Enter key file password: ")
            crypto = EncryptionSystem(key_file, kp)
            if not crypto.keys_loaded:
                print("❌ Loading failed")
                continue

            password = input("Enter encryption password: ")
            text = input("Enter ciphertext to decrypt: ")
            if text and password:
                decrypted = crypto.decrypt(text, password)
                print(f"\n✅ Decrypted result: {decrypted}")

        elif choice == '6':
            # Generate new key with timestamp
            timestamp = int(time.time())
            filename = f"key_{timestamp}.key"
            crypto = EncryptionSystem()
            result = crypto.generate_keys(filename)
            if result is None:
                print("❌ Generation failed")
            else:
                print(f"✅ Key saved to: {filename}")

        elif choice == '7':
            # Show key info
            if crypto is None or not crypto.keys_loaded:
                if os.path.exists("encryption.key"):
                    if key_password is None:
                        key_password = input("Enter default key file password: ")
                    crypto = EncryptionSystem("encryption.key", key_password)
                    if not crypto.keys_loaded:
                        print("❌ Loading failed")
                        continue
                else:
                    print("❌ Keys not loaded")
                    continue
            crypto.print_keys()

        elif choice == '8':
            print("Goodbye!")
            break

        else:
            print("❌ Invalid choice")


if __name__ == "__main__":
    main()