package clusteragent

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"os"
	"path/filepath"
	"testing"

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

func generatePKIXPublicKeyPEM(t *testing.T) ([]byte, *rsa.PrivateKey) {
	t.Helper()
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)

	pubASN1, err := x509.MarshalPKIXPublicKey(&privateKey.PublicKey)
	require.NoError(t, err)

	pubPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: pubASN1,
	})
	return pubPEM, privateKey
}

func generatePKCS1PublicKeyPEM(t *testing.T) ([]byte, *rsa.PrivateKey) {
	t.Helper()
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)

	pubASN1 := x509.MarshalPKCS1PublicKey(&privateKey.PublicKey)
	pubPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PUBLIC KEY",
		Bytes: pubASN1,
	})
	return pubPEM, privateKey
}

func generateMismatchedHeaderPublicKeyPEM(t *testing.T) ([]byte, *rsa.PrivateKey) {
	t.Helper()
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)

	// Intentionally use PKIX bytes with an "RSA PUBLIC KEY" header to match real-world inconsistent outputs.
	pubASN1, err := x509.MarshalPKIXPublicKey(&privateKey.PublicKey)
	require.NoError(t, err)

	pubPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PUBLIC KEY",
		Bytes: pubASN1,
	})
	return pubPEM, privateKey
}

func TestEncryptWithPublicKey(t *testing.T) {
	tests := []struct {
		name      string
		setup     func(t *testing.T) (pubPEM []byte, privKey *rsa.PrivateKey)
		plaintext []byte
		wantErr   bool
	}{
		{
			name: "successful encrypt and decrypt (PKIX PUBLIC KEY)",
			setup: func(t *testing.T) ([]byte, *rsa.PrivateKey) {
				return generatePKIXPublicKeyPEM(t)
			},
			plaintext: []byte("hello world"),
		},
		{
			name: "empty plaintext (PKCS#1 RSA PUBLIC KEY)",
			setup: func(t *testing.T) ([]byte, *rsa.PrivateKey) {
				return generatePKCS1PublicKeyPEM(t)
			},
			plaintext: []byte(""),
		},
		{
			name: "mismatched PEM header is tolerated",
			setup: func(t *testing.T) ([]byte, *rsa.PrivateKey) {
				return generateMismatchedHeaderPublicKeyPEM(t)
			},
			plaintext: []byte("hello world"),
		},
		{
			name: "invalid PEM data",
			setup: func(t *testing.T) ([]byte, *rsa.PrivateKey) {
				return []byte("not a pem"), nil
			},
			plaintext: []byte("test"),
			wantErr:   true,
		},
		{
			name: "non-RSA public key",
			setup: func(t *testing.T) ([]byte, *rsa.PrivateKey) {
				privKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
				require.NoError(t, err)
				pubASN1, err := x509.MarshalPKIXPublicKey(&privKey.PublicKey)
				require.NoError(t, err)
				ecPEM := pem.EncodeToMemory(&pem.Block{
					Type:  "PUBLIC KEY",
					Bytes: pubASN1,
				})
				return ecPEM, nil
			},
			plaintext: []byte("test"),
			wantErr:   true,
		},
		{
			name: "nil PEM input",
			setup: func(t *testing.T) ([]byte, *rsa.PrivateKey) {
				return nil, nil
			},
			plaintext: []byte("test"),
			wantErr:   true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			pubPEM, privKey := tt.setup(t)
			ciphertext, err := EncryptWithPublicKey(pubPEM, tt.plaintext)
			if tt.wantErr {
				assert.Error(t, err)
				return
			}

			require.NoError(t, err)
			require.NotEmpty(t, ciphertext)

			// Verify decryption matches cluster-agent's decrypt logic.
			decrypted, err := rsa.DecryptOAEP(sha256.New(), rand.Reader, privKey, ciphertext, nil)
			require.NoError(t, err)
			assert.Equal(t, tt.plaintext, decrypted)
		})
	}
}

func TestEncryptWithPublicKey_FixturePublicKeyOutput(t *testing.T) {
	// The fixture contains a real /public_key response output from cluster-agent.
	fixturePath := filepath.Join("testdata", "public_key.json")
	data, err := os.ReadFile(fixturePath)
	require.NoError(t, err)

	var resp getPublicKeyResp
	require.NoError(t, json.Unmarshal(data, &resp))
	require.NotEmpty(t, resp.Pem)

	ciphertext, err := EncryptWithPublicKey(resp.Pem, []byte("hello world"))
	require.NoError(t, err)
	require.NotEmpty(t, ciphertext)
}