package crypto

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

// TestNewRSAEncryptor tests the key pair generation function.
func TestNewRSAEncryptor(t *testing.T) {
	// Test with a valid bit size (2048)
	enc, err := NewRSAEncryptor(2048)
	if err != nil {
		t.Fatalf("Failed to create new RSAEncryptor with 2048 bits: %v", err)
	}
	if enc == nil || enc.privateKey == nil {
		t.Fatal("NewRSAEncryptor returned a nil instance or private key")
	}

	// Test with an invalid bit size (too small)
	enc, err = NewRSAEncryptor(1024)
	if err == nil {
		t.Fatalf("Expected an error for a small bit size (1024), but got none")
	}
}

// TestExportAndParsePEM tests the conversion of public key to/from PEM format.
func TestExportAndParsePEM(t *testing.T) {
	// 1. Generate a new key pair
	enc, err := NewRSAEncryptor(2048)
	if err != nil {
		t.Fatalf("Failed to create RSAEncryptor: %v", err)
	}

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

	// 3. Parse the PEM back into a public key object
	parsedPubKey, err := ParsePublicKeyFromPEM(pubPEM)
	if err != nil {
		t.Fatalf("Failed to parse public key from PEM: %v", err)
	}

	// 4. Verify that the parsed key matches the original public key
	// We can compare the public key's exponent and modulus to ensure they are the same.
	originalPubKey := &enc.privateKey.PublicKey
	if parsedPubKey.N.Cmp(originalPubKey.N) != 0 || parsedPubKey.E != originalPubKey.E {
		t.Fatal("Parsed public key does not match the original")
	}
}

// TestVerifyKeyMatch tests the public key and private key matching function.
func TestVerifyKeyMatch(t *testing.T) {
	// 1. Generate a valid key pair for the server
	serverEnc, err := NewRSAEncryptor(2048)
	if err != nil {
		t.Fatalf("Failed to create server RSAEncryptor: %v", err)
	}
	serverPubPEM, _ := serverEnc.ExportPublicKeyAsPEM()

	// 2. Test with a matching key pair: expect (true, nil)
	isMatch, err := serverEnc.VerifyKeyMatch(serverPubPEM)
	if err != nil {
		t.Fatalf("VerifyKeyMatch failed with an unexpected error for a matching key pair: %v", err)
	}
	if !isMatch {
		t.Fatal("VerifyKeyMatch returned false for a matching key pair")
	}

	// 3. Generate an unrelated key pair to simulate a mismatch
	unrelatedKey, _ := rsa.GenerateKey(rand.Reader, 2048)
	unrelatedEnc := &RSAEncryptor{privateKey: unrelatedKey}
	unrelatedPubPEM, _ := unrelatedEnc.ExportPublicKeyAsPEM()

	// 4. Test with a non-matching key pair: expect (false, nil)
	isMatch, err = serverEnc.VerifyKeyMatch(unrelatedPubPEM)
	if err != nil {
		t.Fatalf("VerifyKeyMatch failed with an unexpected error for a non-matching key pair: %v", err)
	}
	if isMatch {
		t.Fatal("VerifyKeyMatch returned true for a non-matching key pair, which is an error")
	}

	// 5. Test with an invalid PEM: expect (false, error)
	invalidPEM := []byte("-----BEGIN RSA PUBLIC KEY-----\nINVALID DATA\n-----END RSA PUBLIC KEY-----")
	isMatch, err = serverEnc.VerifyKeyMatch(invalidPEM)
	if err == nil {
		t.Fatal("Expected an error for an invalid PEM, but got nil")
	}
	if isMatch {
		t.Fatal("VerifyKeyMatch returned true for an invalid PEM, which is an error")
	}
}

// TestEncryptionDecryptionFlow tests the complete end-to-end flow.
func TestEncryptionDecryptionFlow(t *testing.T) {
	// 1. Server-side: Generate key pair
	serverEnc, err := NewRSAEncryptor(2048)
	if err != nil {
		t.Fatalf("Failed to create server RSAEncryptor: %v", err)
	}
	serverPubPEM, _ := serverEnc.ExportPublicKeyAsPEM()

	// 2. Client-side: Encrypt data using the public key PEM
	originalData := []byte("This is a test message to be encrypted.")
	encryptedData, err := EncryptWithPEM(serverPubPEM, originalData)
	if err != nil {
		t.Fatalf("Failed to encrypt data with PEM: %v", err)
	}

	// 3. Server-side: Decrypt the data using the private key
	decryptedData, err := serverEnc.Decrypt(encryptedData)
	if err != nil {
		t.Fatalf("Failed to decrypt data: %v", err)
	}

	// 4. Verification: Check if the decrypted data matches the original
	if !bytes.Equal(originalData, decryptedData) {
		t.Fatalf("Decrypted data does not match original data. Original: %s, Decrypted: %s", originalData, decryptedData)
	}
}