package common

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/pem"
	"testing"

	"github.com/stretchr/testify/require"
)

func TestParsePublicKeyFromPEM(t *testing.T) {
	t.Parallel()

	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)

	pkixBytes, err := x509.MarshalPKIXPublicKey(&privateKey.PublicKey)
	require.NoError(t, err)
	pkcs1Bytes := x509.MarshalPKCS1PublicKey(&privateKey.PublicKey)

	validPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PUBLIC KEY", Bytes: pkcs1Bytes})
	validPEM2 := pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: pkixBytes})
	// Edge case: RSA PUBLIC KEY header but PKIX content (fallback scenario)
	mixedPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PUBLIC KEY", Bytes: pkixBytes})
	wrongTypePEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: pkixBytes})

	testCases := []struct {
		name    string
		pemData []byte
		wantErr bool
	}{
		{
			name:    "valid pem (RSA PUBLIC KEY)",
			pemData: validPEM,
			wantErr: false,
		},
		{
			name:    "valid pem (PUBLIC KEY)",
			pemData: validPEM2,
			wantErr: false,
		},
		{
			name:    "invalid pem",
			pemData: []byte("not-a-pem"),
			wantErr: true,
		},
		{
			name:    "RSA PUBLIC KEY header with PKIX content (fallback)",
			pemData: mixedPEM,
			wantErr: false,
		},
		{
			name:    "wrong block type",
			pemData: wrongTypePEM,
			wantErr: true,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			key, err := ParsePublicKeyFromPEM(tc.pemData)
			if tc.wantErr {
				require.Error(t, err)
				return
			}

			require.NoError(t, err)
			require.NotNil(t, key)
		})
	}
}

func TestEncryptWithPublicKey(t *testing.T) {
	t.Parallel()

	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)

	plainText := []byte("platform_admin_password")
	encrypted, err := EncryptWithPublicKey(&privateKey.PublicKey, plainText)
	require.NoError(t, err)
	require.NotEmpty(t, encrypted)

	decrypted, err := rsa.DecryptOAEP(sha256.New(), rand.Reader, privateKey, encrypted, nil)
	require.NoError(t, err)
	require.Equal(t, plainText, decrypted)
}