package encryption

import (
	"testing"
	"time"
)

func TestDeriveKey(t *testing.T) {
	password := "test-password"
	key := DeriveKey(password)

	if len(key) != 32 {
		t.Errorf("Expected key length 32, got %d", len(key))
	}

	// Test that same password produces same key
	key2 := DeriveKey(password)
	if string(key) != string(key2) {
		t.Error("Same password should produce same key")
	}

	// Test that different passwords produce different keys
	key3 := DeriveKey("different-password")
	if string(key) == string(key3) {
		t.Error("Different passwords should produce different keys")
	}
}

func TestEncryptDecrypt(t *testing.T) {
	key := DeriveKey("test-key")
	originalData := []byte("Hello, World! This is a test message.")

	// Test encryption
	encryptedData, err := EncryptData(originalData, key)
	if err != nil {
		t.Fatalf("Encryption failed: %v", err)
	}

	if len(encryptedData) <= len(originalData) {
		t.Error("Encrypted data should be longer than original data (due to nonce)")
	}

	// Test decryption
	decryptedData, err := DecryptData(encryptedData, key)
	if err != nil {
		t.Fatalf("Decryption failed: %v", err)
	}

	if string(decryptedData) != string(originalData) {
		t.Errorf("Decrypted data doesn't match original. Expected: %s, Got: %s",
			string(originalData), string(decryptedData))
	}
}

func TestEncryptDecryptEmptyData(t *testing.T) {
	key := DeriveKey("test-key")
	originalData := []byte("")

	encryptedData, err := EncryptData(originalData, key)
	if err != nil {
		t.Fatalf("Encryption of empty data failed: %v", err)
	}

	decryptedData, err := DecryptData(encryptedData, key)
	if err != nil {
		t.Fatalf("Decryption of empty data failed: %v", err)
	}

	if len(decryptedData) != 0 {
		t.Error("Decrypted empty data should be empty")
	}
}

func TestEncryptDecryptLargeData(t *testing.T) {
	key := DeriveKey("test-key")

	// Create a large data block (1MB)
	originalData := make([]byte, 1024*1024)
	for i := range originalData {
		originalData[i] = byte(i % 256)
	}

	encryptedData, err := EncryptData(originalData, key)
	if err != nil {
		t.Fatalf("Encryption of large data failed: %v", err)
	}

	decryptedData, err := DecryptData(encryptedData, key)
	if err != nil {
		t.Fatalf("Decryption of large data failed: %v", err)
	}

	if len(decryptedData) != len(originalData) {
		t.Errorf("Decrypted data length mismatch. Expected: %d, Got: %d",
			len(originalData), len(decryptedData))
	}

	for i := range originalData {
		if decryptedData[i] != originalData[i] {
			t.Errorf("Data mismatch at position %d", i)
			break
		}
	}
}

func TestWrongKeyDecryption(t *testing.T) {
	key1 := DeriveKey("key1")
	key2 := DeriveKey("key2")
	originalData := []byte("test data")

	encryptedData, err := EncryptData(originalData, key1)
	if err != nil {
		t.Fatalf("Encryption failed: %v", err)
	}

	// Try to decrypt with wrong key
	_, err = DecryptData(encryptedData, key2)
	if err == nil {
		t.Error("Decryption with wrong key should fail")
	}
}

func TestCorruptedDataDecryption(t *testing.T) {
	key := DeriveKey("test-key")
	originalData := []byte("test data")

	encryptedData, err := EncryptData(originalData, key)
	if err != nil {
		t.Fatalf("Encryption failed: %v", err)
	}

	// Corrupt the data
	encryptedData[0] ^= 0xFF

	// Try to decrypt corrupted data
	_, err = DecryptData(encryptedData, key)
	if err == nil {
		t.Error("Decryption of corrupted data should fail")
	}
}

func TestValidateEncryptionKey(t *testing.T) {
	// Test valid key
	validKey := "a0e3dd20a761b118ca234160dd8b87230a001e332a97c9cfe3b8b9c99efaae03"
	if err := ValidateEncryptionKey(validKey); err != nil {
		t.Errorf("Valid key should pass validation: %v", err)
	}

	// Test short key
	shortKey := "short"
	if err := ValidateEncryptionKey(shortKey); err == nil {
		t.Error("Short key should fail validation")
	}

	// Test weak keys
	weakKeys := []string{
		"password", "123456", "admin", "test", "default",
		"your-secure-encryption-key-change-this-to-something-random",
		"test-encryption-key-12345", "teleport-key", "secret",
	}

	for _, weakKey := range weakKeys {
		if err := ValidateEncryptionKey(weakKey); err == nil {
			t.Errorf("Weak key '%s' should fail validation", weakKey)
		}
	}

	// Test low entropy key
	lowEntropyKey := "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
	if err := ValidateEncryptionKey(lowEntropyKey); err == nil {
		t.Error("Low entropy key should fail validation")
	}
}

func TestReplayProtection(t *testing.T) {
	rp := NewReplayProtection()
	nonce := uint64(12345)
	timestamp := time.Now().Unix()

	// First use should be valid
	if !rp.IsValidNonce(nonce, timestamp) {
		t.Error("First use of nonce should be valid")
	}

	// Replay should be invalid
	if rp.IsValidNonce(nonce, timestamp) {
		t.Error("Replay of nonce should be invalid")
	}

	// Different nonce should be valid
	if !rp.IsValidNonce(nonce+1, timestamp) {
		t.Error("Different nonce should be valid")
	}
}

func TestValidatePacketTimestamp(t *testing.T) {
	now := time.Now().Unix()

	// Current timestamp should be valid
	if !ValidatePacketTimestamp(now) {
		t.Error("Current timestamp should be valid")
	}

	// Recent timestamp should be valid
	recent := now - 60 // 1 minute ago
	if !ValidatePacketTimestamp(recent) {
		t.Error("Recent timestamp should be valid")
	}

	// Old timestamp should be invalid
	old := now - int64(MaxPacketAge.Seconds()) - 1
	if ValidatePacketTimestamp(old) {
		t.Error("Old timestamp should be invalid")
	}

	// Future timestamp should be invalid
	future := now + 3600 // 1 hour in future
	if ValidatePacketTimestamp(future) {
		t.Error("Future timestamp should be invalid")
	}
}

func TestConstantTimeCompare(t *testing.T) {
	a := []byte("test")
	b := []byte("test")
	c := []byte("different")

	if !ConstantTimeCompare(a, b) {
		t.Error("Identical byte slices should compare equal")
	}

	if ConstantTimeCompare(a, c) {
		t.Error("Different byte slices should not compare equal")
	}

	// Test with empty slices
	empty1 := []byte{}
	empty2 := []byte{}
	if !ConstantTimeCompare(empty1, empty2) {
		t.Error("Empty slices should compare equal")
	}
}