package crypto

import (
	"bytes"
	"testing"

	"filippo.io/age"
)

func TestEncryptor(t *testing.T) {
	// Generate a test key pair
	identity, err := age.GenerateX25519Identity()
	if err != nil {
		t.Fatalf("failed to generate identity: %v", err)
	}

	publicKey := identity.Recipient().String()

	// Create encryptor
	enc, err := NewEncryptor([]string{publicKey})
	if err != nil {
		t.Fatalf("failed to create encryptor: %v", err)
	}

	// Test data
	plaintext := []byte("Hello, World! This is a test message.")

	// Encrypt
	ciphertext, err := enc.Encrypt(plaintext)
	if err != nil {
		t.Fatalf("failed to encrypt: %v", err)
	}

	// Verify it's actually encrypted (should be larger and different)
	if bytes.Equal(plaintext, ciphertext) {
		t.Error("ciphertext equals plaintext")
	}

	// Decrypt to verify
	r, err := age.Decrypt(bytes.NewReader(ciphertext), identity)
	if err != nil {
		t.Fatalf("failed to decrypt: %v", err)
	}

	var decrypted bytes.Buffer
	if _, err := decrypted.ReadFrom(r); err != nil {
		t.Fatalf("failed to read decrypted data: %v", err)
	}

	if !bytes.Equal(plaintext, decrypted.Bytes()) {
		t.Error("decrypted data doesn't match original")
	}
}

func TestEncryptorMultipleRecipients(t *testing.T) {
	// Generate three test key pairs
	identity1, err := age.GenerateX25519Identity()
	if err != nil {
		t.Fatalf("failed to generate identity1: %v", err)
	}
	identity2, err := age.GenerateX25519Identity()
	if err != nil {
		t.Fatalf("failed to generate identity2: %v", err)
	}
	identity3, err := age.GenerateX25519Identity()
	if err != nil {
		t.Fatalf("failed to generate identity3: %v", err)
	}

	publicKeys := []string{
		identity1.Recipient().String(),
		identity2.Recipient().String(),
		identity3.Recipient().String(),
	}

	// Create encryptor with multiple recipients
	enc, err := NewEncryptor(publicKeys)
	if err != nil {
		t.Fatalf("failed to create encryptor: %v", err)
	}

	// Test data
	plaintext := []byte("Secret message for multiple recipients")

	// Encrypt
	ciphertext, err := enc.Encrypt(plaintext)
	if err != nil {
		t.Fatalf("failed to encrypt: %v", err)
	}

	// Verify each recipient can decrypt
	identities := []age.Identity{identity1, identity2, identity3}
	for i, identity := range identities {
		r, err := age.Decrypt(bytes.NewReader(ciphertext), identity)
		if err != nil {
			t.Fatalf("recipient %d failed to decrypt: %v", i+1, err)
		}

		var decrypted bytes.Buffer
		if _, err := decrypted.ReadFrom(r); err != nil {
			t.Fatalf("recipient %d failed to read decrypted data: %v", i+1, err)
		}

		if !bytes.Equal(plaintext, decrypted.Bytes()) {
			t.Errorf("recipient %d: decrypted data doesn't match original", i+1)
		}
	}
}

func TestEncryptorUpdateRecipients(t *testing.T) {
	// Generate two identities
	identity1, _ := age.GenerateX25519Identity()
	identity2, _ := age.GenerateX25519Identity()

	publicKey1 := identity1.Recipient().String()
	publicKey2 := identity2.Recipient().String()

	// Create encryptor with first key
	enc, err := NewEncryptor([]string{publicKey1})
	if err != nil {
		t.Fatalf("failed to create encryptor: %v", err)
	}

	// Encrypt with first key
	plaintext := []byte("test data")
	ciphertext1, err := enc.Encrypt(plaintext)
	if err != nil {
		t.Fatalf("failed to encrypt: %v", err)
	}

	// Update to second key
	if err := enc.UpdateRecipients([]string{publicKey2}); err != nil {
		t.Fatalf("failed to update recipients: %v", err)
	}

	// Encrypt with second key
	ciphertext2, err := enc.Encrypt(plaintext)
	if err != nil {
		t.Fatalf("failed to encrypt: %v", err)
	}

	// First ciphertext should only decrypt with first identity
	if _, err := age.Decrypt(bytes.NewReader(ciphertext1), identity1); err != nil {
		t.Error("failed to decrypt with identity1")
	}
	if _, err := age.Decrypt(bytes.NewReader(ciphertext1), identity2); err == nil {
		t.Error("should not decrypt with identity2")
	}

	// Second ciphertext should only decrypt with second identity
	if _, err := age.Decrypt(bytes.NewReader(ciphertext2), identity2); err != nil {
		t.Error("failed to decrypt with identity2")
	}
	if _, err := age.Decrypt(bytes.NewReader(ciphertext2), identity1); err == nil {
		t.Error("should not decrypt with identity1")
	}
}