package crypto

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"testing"
)

func TestEncryptor(t *testing.T) {
	encryptor, err := NewEncryptor()
	if err != nil {
		t.Fatalf("Failed to create Encryptor: %v", err)
	}

	pubPEM, err := encryptor.ExportPublicKeyAsPEM()
	if err != nil {
		t.Fatalf("Failed to export public key as PEM: %v", err)
	}
	if len(pubPEM) == 0 {
		t.Fatal("Exported public key PEM is empty")
	}

	pubKey, err := parsePublicKeyFromPEM(pubPEM)
	if err != nil {
		t.Fatalf("Failed to parse public key from PEM: %v", err)
	}

	plaintext := []byte("Hello, World!")
	encryptedData, err := rsa.EncryptOAEP(sha256.New(), rand.Reader, pubKey, plaintext, nil)
	if err != nil {
		t.Fatalf("Failed to encrypt data: %v", err)
	}

	// decrypt the data and compare with original plaintext
	decryptedData, err := encryptor.Decrypt(encryptedData)
	if err != nil {
		t.Fatalf("Failed to decrypt data: %v", err)
	}
	if !bytes.Equal(decryptedData, plaintext) {
		t.Fatalf("Decrypted data does not match original plaintext. Got: %s, Want: %s", decryptedData, plaintext)
	}

	// try to decrypt with a different private key, should fail
	anotherEncryptor, err := NewEncryptor()
	if err != nil {
		t.Fatalf("Failed to create another Encryptor: %v", err)
	}
	_, err = anotherEncryptor.Decrypt(encryptedData)
	if err == nil {
		t.Fatal("Decryption should have failed with a different private key, but it succeeded")
	}

	// encrypt with this new public key and descrypt with the original private key
	anotherPubPEM, err := anotherEncryptor.ExportPublicKeyAsPEM()
	if err != nil {
		t.Fatalf("Failed to export another public key as PEM: %v", err)
	}
	anotherPubKey, err := parsePublicKeyFromPEM(anotherPubPEM)
	if err != nil {
		t.Fatalf("Failed to parse another public key from PEM: %v", err)
	}
	anotherEncryptedData, err := rsa.EncryptOAEP(sha256.New(), rand.Reader, anotherPubKey, plaintext, nil)
	if err != nil {
		t.Fatalf("Failed to encrypt data with another public key: %v", err)
	}
	_, err = encryptor.Decrypt(anotherEncryptedData)
	if err == nil {
		t.Fatal("Decryption should have failed with a different public key, but it succeeded")
	}
}