# SPDX-License-Identifier: Apache-2.0
# Copyright 2025 AES Encryption Project

package main

import (
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"testing"
)

func TestNewAESEncryption(t *testing.T) {
	password := "test-password"
	aes := NewAESEncryption(password)
	
	if aes == nil {
		t.Fatal("NewAESEncryption returned nil")
	}
	
	if len(aes.key) != 32 {
		t.Errorf("Expected key length 32, got %d", len(aes.key))
	}
}

func TestNewAESEncryptionWithKey(t *testing.T) {
	tests := []struct {
		name    string
		keySize int
		wantErr bool
	}{
		{"AES-128", 16, false},
		{"AES-192", 24, false},
		{"AES-256", 32, false},
		{"Invalid size", 20, true},
		{"Too small", 8, true},
		{"Too large", 64, true},
	}
	
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			key := make([]byte, tt.keySize)
			rand.Read(key)
			
			aes, err := NewAESEncryptionWithKey(key)
			
			if tt.wantErr {
				if err == nil {
					t.Error("Expected error but got none")
				}
				return
			}
			
			if err != nil {
				t.Errorf("Unexpected error: %v", err)
				return
			}
			
			if aes == nil {
				t.Error("Expected AESEncryption instance but got nil")
			}
		})
	}
}

func TestEncryptDecrypt(t *testing.T) {
	aes := NewAESEncryption("test-password")
	
	testCases := []string{
		"Hello, World!",
		"This is a longer message with special characters: !@#$%^&*()",
		"Short",
		"Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.",
		"🚀 Unicode test with emojis 🔐",
	}
	
	for _, plaintext := range testCases {
		t.Run("Text: "+plaintext[:min(20, len(plaintext))], func(t *testing.T) {
			// Encrypt
			encrypted, err := aes.Encrypt(plaintext)
			if err != nil {
				t.Fatalf("Encryption failed: %v", err)
			}
			
			// Verify it's base64 encoded
			_, err = base64.StdEncoding.DecodeString(encrypted)
			if err != nil {
				t.Errorf("Encrypted text is not valid base64: %v", err)
			}
			
			// Decrypt
			decrypted, err := aes.Decrypt(encrypted)
			if err != nil {
				t.Fatalf("Decryption failed: %v", err)
			}
			
			// Verify
			if plaintext != decrypted {
				t.Errorf("Decrypted text doesn't match original.\nOriginal: %s\nDecrypted: %s", plaintext, decrypted)
			}
		})
	}
}

func TestEncryptDecryptBytes(t *testing.T) {
	key, err := GenerateKey(32)
	if err != nil {
		t.Fatalf("Failed to generate key: %v", err)
	}
	
	aes, err := NewAESEncryptionWithKey(key)
	if err != nil {
		t.Fatalf("Failed to create AES instance: %v", err)
	}
	
	testData := [][]byte{
		[]byte("Hello, World!"),
		[]byte("Binary data with null bytes: \x00\x01\x02\x03"),
		make([]byte, 1000), // Large data
		[]byte{0xFF, 0xFE, 0xFD, 0xFC}, // Binary data
	}
	
	// Fill large data with random bytes
	rand.Read(testData[2])
	
	for i, data := range testData {
		t.Run(fmt.Sprintf("Data_%d", i), func(t *testing.T) {
			// Encrypt
			encrypted, err := aes.EncryptBytes(data)
			if err != nil {
				t.Fatalf("Encryption failed: %v", err)
			}
			
			// Decrypt
			decrypted, err := aes.DecryptBytes(encrypted)
			if err != nil {
				t.Fatalf("Decryption failed: %v", err)
			}
			
			// Verify
			if !bytes.Equal(data, decrypted) {
				t.Errorf("Decrypted data doesn't match original")
			}
		})
	}
}

func TestEncryptionRandomness(t *testing.T) {
	aes := NewAESEncryption("test-password")
	plaintext := "Same message"
	
	// Encrypt the same message multiple times
	encrypted1, err := aes.Encrypt(plaintext)
	if err != nil {
		t.Fatalf("First encryption failed: %v", err)
	}
	
	encrypted2, err := aes.Encrypt(plaintext)
	if err != nil {
		t.Fatalf("Second encryption failed: %v", err)
	}
	
	// Results should be different due to random nonce
	if encrypted1 == encrypted2 {
		t.Error("Encrypted results are identical - nonce randomization may be broken")
	}
	
	// Both should decrypt to the same plaintext
	decrypted1, err := aes.Decrypt(encrypted1)
	if err != nil {
		t.Fatalf("First decryption failed: %v", err)
	}
	
	decrypted2, err := aes.Decrypt(encrypted2)
	if err != nil {
		t.Fatalf("Second decryption failed: %v", err)
	}
	
	if decrypted1 != plaintext || decrypted2 != plaintext {
		t.Error("Decrypted messages don't match original")
	}
}

func TestErrorHandling(t *testing.T) {
	aes := NewAESEncryption("test-password")
	
	// Test empty plaintext
	_, err := aes.Encrypt("")
	if err == nil {
		t.Error("Expected error for empty plaintext")
	}
	
	// Test empty ciphertext
	_, err = aes.Decrypt("")
	if err == nil {
		t.Error("Expected error for empty ciphertext")
	}
	
	// Test invalid base64
	_, err = aes.Decrypt("invalid-base64!")
	if err == nil {
		t.Error("Expected error for invalid base64")
	}
	
	// Test too short ciphertext
	shortCiphertext := base64.StdEncoding.EncodeToString([]byte("short"))
	_, err = aes.Decrypt(shortCiphertext)
	if err == nil {
		t.Error("Expected error for too short ciphertext")
	}
	
	// Test corrupted ciphertext
	validEncrypted, _ := aes.Encrypt("test message")
	corruptedBytes, _ := base64.StdEncoding.DecodeString(validEncrypted)
	corruptedBytes[len(corruptedBytes)-1] ^= 0xFF // Flip last byte
	corrupted := base64.StdEncoding.EncodeToString(corruptedBytes)
	
	_, err = aes.Decrypt(corrupted)
	if err == nil {
		t.Error("Expected error for corrupted ciphertext")
	}
}

func TestGenerateKey(t *testing.T) {
	validSizes := []int{16, 24, 32}
	
	for _, size := range validSizes {
		t.Run(fmt.Sprintf("Size_%d", size), func(t *testing.T) {
			key, err := GenerateKey(size)
			if err != nil {
				t.Fatalf("GenerateKey failed: %v", err)
			}
			
			if len(key) != size {
				t.Errorf("Expected key size %d, got %d", size, len(key))
			}
			
			// Test that keys are different
			key2, err := GenerateKey(size)
			if err != nil {
				t.Fatalf("Second GenerateKey failed: %v", err)
			}
			
			if bytes.Equal(key, key2) {
				t.Error("Generated keys are identical - randomness may be broken")
			}
		})
	}
	
	// Test invalid sizes
	invalidSizes := []int{0, 8, 15, 17, 20, 64}
	for _, size := range invalidSizes {
		t.Run(fmt.Sprintf("Invalid_Size_%d", size), func(t *testing.T) {
			_, err := GenerateKey(size)
			if err == nil {
				t.Errorf("Expected error for invalid key size %d", size)
			}
		})
	}
}

func TestKeyIsolation(t *testing.T) {
	// Test that external key modification doesn't affect internal key
	originalKey := []byte("1234567890123456") // 16 bytes for AES-128
	keyCopy := make([]byte, len(originalKey))
	copy(keyCopy, originalKey)
	
	aes, err := NewAESEncryptionWithKey(keyCopy)
	if err != nil {
		t.Fatalf("Failed to create AES instance: %v", err)
	}
	
	// Modify the external key
	keyCopy[0] = 0xFF
	
	// Encryption should still work with original key
	plaintext := "test message"
	encrypted, err := aes.Encrypt(plaintext)
	if err != nil {
		t.Fatalf("Encryption failed: %v", err)
	}
	
	decrypted, err := aes.Decrypt(encrypted)
	if err != nil {
		t.Fatalf("Decryption failed: %v", err)
	}
	
	if plaintext != decrypted {
		t.Error("Key isolation failed - external modification affected internal key")
	}
}

// Helper function for min (Go 1.21+ has this built-in)
func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

// Benchmark tests
func BenchmarkEncrypt(b *testing.B) {
	aes := NewAESEncryption("benchmark-password")
	plaintext := "This is a benchmark test message for AES encryption performance."
	
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_, err := aes.Encrypt(plaintext)
		if err != nil {
			b.Fatalf("Encryption failed: %v", err)
		}
	}
}

func BenchmarkDecrypt(b *testing.B) {
	aes := NewAESEncryption("benchmark-password")
	plaintext := "This is a benchmark test message for AES encryption performance."
	encrypted, _ := aes.Encrypt(plaintext)
	
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_, err := aes.Decrypt(encrypted)
		if err != nil {
			b.Fatalf("Decryption failed: %v", err)
		}
	}
}