# SPDX-License-Identifier: Apache-2.0
# Copyright 2025 AES Encryption Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""AES Encryption Demo Script.

Comprehensive demonstration of AES encryption capabilities including:
- Password-based encryption
- Key-based encryption
- Multiple encryption modes
- File encryption
- Performance testing
"""

import json
import os
import time
from typing import Dict, Any

from aes_encryption import (
    AESEncryption,
    create_aes_encryptor,
    encrypt_string,
    decrypt_string,
    AESEncryptionError
)


def demo_basic_encryption():
    """Demonstrate basic string encryption with password."""
    print("🔐 Basic Password-Based Encryption Demo")
    print("-" * 50)
    
    password = "MySecurePassword123!"
    message = "This is a confidential message that needs protection."
    
    print(f"Original message: {message}")
    print(f"Password: {password}")
    
    # Encrypt
    start_time = time.time()
    encrypted = encrypt_string(message, password)
    encrypt_time = time.time() - start_time
    
    print(f"Encrypted data: {encrypted[:60]}...")
    print(f"Encryption time: {encrypt_time:.4f} seconds")
    
    # Decrypt
    start_time = time.time()
    decrypted = decrypt_string(encrypted, password)
    decrypt_time = time.time() - start_time
    
    print(f"Decrypted message: {decrypted}")
    print(f"Decryption time: {decrypt_time:.4f} seconds")
    print(f"✅ Encryption/Decryption successful: {message == decrypted}")
    print()


def demo_different_modes():
    """Demonstrate different encryption modes."""
    print("🔄 Different Encryption Modes Demo")
    print("-" * 50)
    
    test_data = "Sensitive data for mode comparison testing."
    password = "TestPassword123"
    
    modes = ['CBC', 'GCM', 'CTR']
    results = {}
    
    for mode in modes:
        print(f"Testing {mode} mode...")
        
        try:
            # Encrypt
            start_time = time.time()
            encrypted = encrypt_string(test_data, password, mode)
            encrypt_time = time.time() - start_time
            
            # Decrypt
            start_time = time.time()
            decrypted = decrypt_string(encrypted, password, mode)
            decrypt_time = time.time() - start_time
            
            success = test_data == decrypted
            
            results[mode] = {
                'success': success,
                'encrypt_time': encrypt_time,
                'decrypt_time': decrypt_time,
                'encrypted_size': len(encrypted)
            }
            
            print(f"  ✅ {mode}: Success={success}, "
                  f"Encrypt={encrypt_time:.4f}s, Decrypt={decrypt_time:.4f}s")
                  
        except Exception as e:
            print(f"  ❌ {mode}: Failed - {str(e)}")
            results[mode] = {'success': False, 'error': str(e)}
    
    print(f"\n📊 Mode Comparison Results:")
    for mode, result in results.items():
        if result['success']:
            print(f"  {mode}: {result['encrypt_time']:.4f}s encrypt, "
                  f"{result['decrypt_time']:.4f}s decrypt, "
                  f"{result['encrypted_size']} bytes")
    print()


def demo_key_generation():
    """Demonstrate key generation and management."""
    print("🔑 Key Generation and Management Demo")
    print("-" * 50)
    
    # Generate different key sizes
    key_sizes = [16, 24, 32]  # AES-128, AES-192, AES-256
    
    for key_size in key_sizes:
        encryptor = AESEncryption()
        key = encryptor.generate_key(key_size)
        
        key_info = encryptor.get_key_info()
        print(f"Generated {key_info['key_size_bits']}-bit key: "
              f"{key.hex()[:32]}...")
        
        # Test encryption with generated key
        test_message = f"Test message for {key_info['key_size_bits']}-bit AES"
        ciphertext, iv = encryptor.encrypt(test_message)
        decrypted = encryptor.decrypt(ciphertext, iv)
        
        print(f"  Test result: {'✅ Success' if test_message == decrypted.decode() else '❌ Failed'}")
    
    print()


def demo_password_derivation():
    """Demonstrate password-based key derivation."""
    print("🔐 Password-Based Key Derivation Demo")
    print("-" * 50)
    
    password = "MyComplexPassword!@#123"
    encryptor = AESEncryption()
    
    # Derive key with custom parameters
    key, salt = encryptor.derive_key_from_password(
        password=password,
        iterations=150000,  # Higher iteration count for better security
        key_size=32  # 256-bit key
    )
    
    print(f"Password: {password}")
    print(f"Salt: {salt.hex()}")
    print(f"Derived key: {key.hex()[:32]}...")
    print(f"Key info: {encryptor.get_key_info()}")
    
    # Test with same password and salt
    encryptor2 = AESEncryption()
    key2, _ = encryptor2.derive_key_from_password(password, salt, 150000, 32)
    
    print(f"Key derivation consistency: {'✅ Consistent' if key == key2 else '❌ Inconsistent'}")
    print()


def demo_file_encryption():
    """Demonstrate file encryption and decryption."""
    print("📁 File Encryption Demo")
    print("-" * 50)
    
    # Create test file
    test_filename = "test_file.txt"
    encrypted_filename = "test_file.txt.encrypted"
    
    test_content = """This is a test file for encryption demonstration.
It contains multiple lines of text.
Some sensitive information might be stored here.
The file will be encrypted and then decrypted to verify integrity."""
    
    try:
        # Write test file
        with open(test_filename, 'w', encoding='utf-8') as f:
            f.write(test_content)
        
        print(f"Created test file: {test_filename}")
        
        # Encrypt file
        password = "FileEncryptionPassword123!"
        encryptor = create_aes_encryptor(password=password, mode='CBC')
        
        with open(test_filename, 'rb') as f:
            file_data = f.read()
        
        encrypted_data, iv = encryptor.encrypt(file_data)
        
        # Save encrypted file with metadata
        file_metadata = {
            'iv': iv.hex(),
            'mode': 'CBC',
            'original_size': len(file_data)
        }
        
        with open(encrypted_filename, 'wb') as f:
            # Write metadata length, metadata, then encrypted data
            metadata_json = json.dumps(file_metadata).encode('utf-8')
            f.write(len(metadata_json).to_bytes(4, 'big'))
            f.write(metadata_json)
            f.write(encrypted_data)
        
        print(f"File encrypted and saved as: {encrypted_filename}")
        
        # Decrypt file
        with open(encrypted_filename, 'rb') as f:
            metadata_length = int.from_bytes(f.read(4), 'big')
            metadata_json = f.read(metadata_length)
            encrypted_content = f.read()
        
        metadata = json.loads(metadata_json.decode('utf-8'))
        iv = bytes.fromhex(metadata['iv'])
        
        decryptor = create_aes_encryptor(password=password, mode=metadata['mode'])
        decrypted_data = decryptor.decrypt(encrypted_content, iv)
        
        # Verify integrity
        success = file_data == decrypted_data
        print(f"File decryption: {'✅ Success' if success else '❌ Failed'}")
        print(f"Original size: {len(file_data)} bytes")
        print(f"Encrypted size: {len(encrypted_content)} bytes")
        print(f"Decrypted size: {len(decrypted_data)} bytes")
        
        # Clean up
        os.remove(test_filename)
        os.remove(encrypted_filename)
        print("Test files cleaned up")
        
    except Exception as e:
        print(f"❌ File encryption demo failed: {str(e)}")
        # Clean up on error
        for filename in [test_filename, encrypted_filename]:
            if os.path.exists(filename):
                os.remove(filename)
    
    print()


def demo_error_handling():
    """Demonstrate error handling and edge cases."""
    print("⚠️  Error Handling Demo")
    print("-" * 50)
    
    # Test invalid key
    try:
        encryptor = AESEncryption()
        encryptor.set_key(b"too_short")  # Invalid key length
    except AESEncryptionError as e:
        print(f"✅ Caught invalid key error: {str(e)}")
    
    # Test decryption with wrong password
    try:
        original_password = "correct_password"
        wrong_password = "wrong_password"
        message = "Secret message"
        
        encrypted = encrypt_string(message, original_password)
        decrypt_string(encrypted, wrong_password)  # Should fail
    except AESEncryptionError as e:
        print(f"✅ Caught wrong password error: {str(e)}")
    
    # Test unsupported mode
    try:
        AESEncryption(mode='INVALID_MODE')
    except ValueError as e:
        print(f"✅ Caught unsupported mode error: {str(e)}")
    
    print()


def performance_benchmark():
    """Run performance benchmarks."""
    print("⚡ Performance Benchmark")
    print("-" * 50)
    
    # Test different data sizes
    data_sizes = [1024, 10240, 102400, 1048576]  # 1KB, 10KB, 100KB, 1MB
    password = "BenchmarkPassword123"
    
    print(f"{'Size':<10} {'Encrypt (ms)':<15} {'Decrypt (ms)':<15} {'Throughput (MB/s)':<20}")
    print("-" * 70)
    
    for size in data_sizes:
        test_data = os.urandom(size)  # Random binary data
        
        encryptor = create_aes_encryptor(password=password, mode='CBC')
        
        # Encryption benchmark
        start_time = time.time()
        ciphertext, iv = encryptor.encrypt(test_data)
        encrypt_time = time.time() - start_time
        
        # Decryption benchmark
        start_time = time.time()
        decrypted = encryptor.decrypt(ciphertext, iv)
        decrypt_time = time.time() - start_time
        
        # Calculate throughput (MB/s)
        total_time = encrypt_time + decrypt_time
        throughput = (size * 2) / (1024 * 1024) / total_time if total_time > 0 else 0
        
        size_str = f"{size//1024}KB" if size < 1048576 else f"{size//1048576}MB"
        
        print(f"{size_str:<10} {encrypt_time*1000:<15.2f} {decrypt_time*1000:<15.2f} {throughput:<20.2f}")
        
        # Verify integrity
        if test_data != decrypted:
            print(f"❌ Data integrity check failed for {size_str}")
    
    print()


def main():
    """Run all demonstrations."""
    print("🚀 AES Encryption Module - Comprehensive Demo")
    print("=" * 60)
    print()
    
    try:
        demo_basic_encryption()
        demo_different_modes()
        demo_key_generation()
        demo_password_derivation()
        demo_file_encryption()
        demo_error_handling()
        performance_benchmark()
        
        print("🎉 All demonstrations completed successfully!")
        
    except Exception as e:
        print(f"❌ Demo failed with error: {str(e)}")
        raise


if __name__ == "__main__":
    main()