// 04_encryption_decryption.rhai
// Tests for asymmetric encryption and decryption in the Keypair module

// Custom assert function
fn assert_true(condition, message) {
    if !condition {
        print(`ASSERTION FAILED: ${message}`);
        throw message;
    }
}

print("=== Testing Asymmetric Encryption and Decryption ===");

// Test creating keypairs for sender and recipient
print("Setting up sender and recipient keypairs...");
let space_name = "encryption_test_space";
let password = "secure_password";
let sender_name = "sender_keypair";
let recipient_name = "recipient_keypair";

if create_key_space(space_name, password) {
    print(`✓ Key space "${space_name}" created successfully`);
    
    // Create sender keypair
    if create_keypair(sender_name, password) {
        print(`✓ Sender keypair "${sender_name}" created successfully`);
    } else {
        print(`✗ Failed to create sender keypair "${sender_name}"`);
        throw `Failed to create sender keypair "${sender_name}"`;
    }
    
    // Create recipient keypair
    if create_keypair(recipient_name, password) {
        print(`✓ Recipient keypair "${recipient_name}" created successfully`);
    } else {
        print(`✗ Failed to create recipient keypair "${recipient_name}"`);
        throw `Failed to create recipient keypair "${recipient_name}"`;
    }
    
    // Get recipient's public key
    if select_keypair(recipient_name) {
        print(`✓ Selected recipient keypair "${recipient_name}" successfully`);
        let recipient_pub_key = keypair_pub_key();
        assert_true(recipient_pub_key.len() > 0, "Recipient public key should not be empty");
        print(`✓ Retrieved recipient public key: ${recipient_pub_key.len()} bytes`);
        
        // Switch to sender keypair
        if select_keypair(sender_name) {
            print(`✓ Selected sender keypair "${sender_name}" successfully`);
            
            // Test encrypting a message with recipient's public key
            print("\nTesting encrypting a message...");
            let message = "This is a secret message for the recipient";
            let ciphertext = encrypt_asymmetric(recipient_pub_key, message);
            assert_true(ciphertext.len() > 0, "Ciphertext should not be empty");
            print(`✓ Message encrypted successfully: ${ciphertext.len()} bytes`);
            
            // Switch back to recipient keypair to decrypt
            if select_keypair(recipient_name) {
                print(`✓ Switched back to recipient keypair "${recipient_name}" successfully`);
                
                // Test decrypting the message
                print("Testing decrypting the message...");
                let decrypted = decrypt_asymmetric(ciphertext);
                assert_true(decrypted == message, "Decrypted message should match original");
                print(`✓ Message decrypted successfully: "${decrypted}"`);
                
                // Edge case: Test with empty message
                print("\nTesting with empty message...");
                let empty_message = "";
                let empty_ciphertext = encrypt_asymmetric(recipient_pub_key, empty_message);
                assert_true(empty_ciphertext.len() > 0, "Ciphertext for empty message should not be empty");
                
                let empty_decrypted = decrypt_asymmetric(empty_ciphertext);
                assert_true(empty_decrypted == empty_message, "Decrypted empty message should be empty");
                print("✓ Empty message encrypted and decrypted successfully");
                
                // Edge case: Test with large message
                print("\nTesting with large message...");
                let large_message = "A" * 10000; // 10KB message
                let large_ciphertext = encrypt_asymmetric(recipient_pub_key, large_message);
                assert_true(large_ciphertext.len() > 0, "Ciphertext for large message should not be empty");
                
                let large_decrypted = decrypt_asymmetric(large_ciphertext);
                assert_true(large_decrypted == large_message, "Decrypted large message should match original");
                print("✓ Large message encrypted and decrypted successfully");
                
                // Error case: Attempt to decrypt with the wrong keypair
                print("\nTesting decryption with wrong keypair...");
                if select_keypair(sender_name) {
                    print(`✓ Switched to sender keypair "${sender_name}" successfully`);
                    
                    let wrong_keypair_success = true;
                    try {
                        let wrong_decrypted = decrypt_asymmetric(ciphertext);
                        // If we get here, the decryption didn't fail as expected
                        assert_true(wrong_decrypted != message, "Decryption with wrong keypair should not match original message");
                    } catch(err) {
                        wrong_keypair_success = false;
                        print(`✓ Caught expected error for decryption with wrong keypair: ${err}`);
                    }
                    
                    // Note: Some implementations might not throw an error but return garbage data
                    // So we don't assert on wrong_keypair_success
                    
                    // Switch back to recipient for further tests
                    if select_keypair(recipient_name) {
                        print(`✓ Switched back to recipient keypair "${recipient_name}" successfully`);
                    } else {
                        print(`✗ Failed to switch back to recipient keypair "${recipient_name}"`);
                        throw `Failed to switch back to recipient keypair "${recipient_name}"`;
                    }
                } else {
                    print(`✗ Failed to switch to sender keypair "${sender_name}"`);
                    throw `Failed to switch to sender keypair "${sender_name}"`;
                }
                
                // Error case: Test with malformed ciphertext
                print("\nTesting with malformed ciphertext...");
                let malformed_ciphertext = [0, 1, 2, 3]; // Invalid ciphertext bytes
                let malformed_success = false;
                try {
                    decrypt_asymmetric(malformed_ciphertext);
                    malformed_success = true;
                } catch(err) {
                    print(`✓ Caught expected error for malformed ciphertext: ${err}`);
                }
                assert_true(!malformed_success, "Decrypting malformed ciphertext should fail");
                
                // Error case: Test with invalid public key for encryption
                print("\nTesting encryption with invalid public key...");
                if select_keypair(sender_name) {
                    print(`✓ Switched to sender keypair "${sender_name}" successfully`);
                    
                    let invalid_pub_key = [0, 1, 2, 3]; // Invalid public key bytes
                    let invalid_key_success = false;
                    try {
                        encrypt_asymmetric(invalid_pub_key, message);
                        invalid_key_success = true;
                    } catch(err) {
                        print(`✓ Caught expected error for invalid public key: ${err}`);
                    }
                    assert_true(!invalid_key_success, "Encrypting with invalid public key should fail");
                } else {
                    print(`✗ Failed to switch to sender keypair "${sender_name}"`);
                    throw `Failed to switch to sender keypair "${sender_name}"`;
                }
                
                // Error case: Test with tampered ciphertext
                print("\nTesting with tampered ciphertext...");
                if select_keypair(recipient_name) {
                    print(`✓ Switched to recipient keypair "${recipient_name}" successfully`);
                    
                    // Tamper with the ciphertext (change a byte in the middle)
                    let tampered_ciphertext = ciphertext.clone();
                    if tampered_ciphertext.len() > 100 {
                        tampered_ciphertext[100] = (tampered_ciphertext[100] + 1) % 256;
                        
                        let tampered_success = false;
                        try {
                            let tampered_decrypted = decrypt_asymmetric(tampered_ciphertext);
                            tampered_success = tampered_decrypted == message;
                        } catch(err) {
                            print(`✓ Caught expected error for tampered ciphertext: ${err}`);
                        }
                        assert_true(!tampered_success, "Decrypting tampered ciphertext should fail or produce incorrect result");
                    } else {
                        print("Note: Ciphertext too short to test tampering");
                    }
                } else {
                    print(`✗ Failed to switch to recipient keypair "${recipient_name}"`);
                    throw `Failed to switch to recipient keypair "${recipient_name}"`;
                }
                
            } else {
                print(`✗ Failed to switch back to recipient keypair "${recipient_name}"`);
                throw `Failed to switch back to recipient keypair "${recipient_name}"`;
            }
        } else {
            print(`✗ Failed to select sender keypair "${sender_name}"`);
            throw `Failed to select sender keypair "${sender_name}"`;
        }
    } else {
        print(`✗ Failed to select recipient keypair "${recipient_name}"`);
        throw `Failed to select recipient keypair "${recipient_name}"`;
    }
} else {
    print(`✗ Failed to create key space "${space_name}"`);
    throw `Failed to create key space "${space_name}"`;
}

print("All asymmetric encryption and decryption tests completed successfully!");