package keys_test

import (
	"bytes"
	"context"
	"encoding/hex"
	"testing"

	"github.com/ansg191/ibd-trader-backend/internal/keys"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/mock/gomock"
)

func TestEncrypt(t *testing.T) {
	ctrl := gomock.NewController(t)

	// Replace RNG with a deterministic RNG
	aesKey := []byte("0123456789abcdef0123456789abcdef")
	nonce := []byte("0123456789ab")
	keys.CSRNG = bytes.NewReader(append(aesKey, nonce...))

	// Create a mock KMS
	kms := NewMockKeyManagementService(ctrl)
	keyName := "keyName"

	ctx := context.Background()
	plaintext := []byte("plaintext")

	kms.EXPECT().
		Encrypt(ctx, keyName, aesKey).
		Return([]byte("encryptedKey"), nil)

	ciphertext, encryptedKey, err := keys.Encrypt(ctx, kms, keyName, plaintext)
	require.NoError(t, err)

	encrypted, err := hex.DecodeString("e9c586532dbefd63812293e1c4baf71edb7042a294c49c2020")
	require.NoError(t, err)
	assert.Equal(t, append(nonce, encrypted...), ciphertext)
	assert.Equal(t, []byte("encryptedKey"), encryptedKey)
}

func TestDecrypt(t *testing.T) {
	ctrl := gomock.NewController(t)

	kms := NewMockKeyManagementService(ctrl)
	keyName := "keyName"

	ctx := context.Background()
	encryptedKey := []byte("encryptedKey")
	ciphertext, err := hex.DecodeString("e9c586532dbefd63812293e1c4baf71edb7042a294c49c2020")
	require.NoError(t, err)
	ciphertext = append([]byte("0123456789ab"), ciphertext...)

	aesKey := []byte("0123456789abcdef0123456789abcdef")
	kms.EXPECT().
		Decrypt(ctx, keyName, encryptedKey).
		Return(aesKey, nil)

	plaintext, err := keys.Decrypt(ctx, kms, keyName, ciphertext, encryptedKey)
	require.NoError(t, err)
	assert.Equal(t, []byte("plaintext"), plaintext)
}