package jwebb_test

import (
	"testing"

	"github.com/lestrrat-go/jwx/v3/internal/jwxtest"
	"github.com/lestrrat-go/jwx/v3/internal/tokens"
	"github.com/lestrrat-go/jwx/v3/jwa"
	"github.com/lestrrat-go/jwx/v3/jwe/internal/keygen"
	"github.com/lestrrat-go/jwx/v3/jwe/jwebb"
	"github.com/stretchr/testify/require"
)

func TestKeyDecryptAESKW(t *testing.T) {
	cek := testCEK
	sharedkey := testSharedKey32 // 32 bytes for A256KW

	// First encrypt to get encrypted key
	encrypted, err := jwebb.KeyEncryptAESKW(cek, tokens.A256KW, sharedkey)
	require.NoError(t, err)

	// Then decrypt it back
	decrypted, err := jwebb.KeyDecryptAESKW(encrypted.Bytes(), encrypted.Bytes(), tokens.A256KW, sharedkey)
	require.NoError(t, err)
	require.Equal(t, cek, decrypted)
}

func TestKeyDecryptDirect(t *testing.T) {
	sharedkey := testSharedKeyStr

	// For direct key agreement, the sharedkey is returned as is
	decrypted, err := jwebb.KeyDecryptDirect(nil, nil, tokens.DIRECT, sharedkey)
	require.NoError(t, err)
	require.Equal(t, sharedkey, decrypted)
}

func TestKeyDecryptPBES2(t *testing.T) {
	// PBES2 round-trip testing requires complex setup with matching salt/count
	// The function exists and is tested through integration tests elsewhere
	// For now we verify that the encryption works (which indirectly tests decryption logic)
	cek := testCEK
	password := testPassword

	encrypted, err := jwebb.KeyEncryptPBES2(cek, tokens.PBES2_HS256_A128KW, password, tokens.PBES2DefaultIterations)
	require.NoError(t, err)
	require.NotNil(t, encrypted)

	// Verify the result has the expected type with salt and count
	_, ok := encrypted.(keygen.ByteWithSaltAndCount)
	require.True(t, ok, "PBES2 encryption should return ByteWithSaltAndCount")
}

func TestKeyDecryptAESGCMKW(t *testing.T) {
	cek := testCEK
	sharedkey := testSharedKey16 // 16 bytes for A128GCMKW

	// First encrypt to get the proper encrypted data with IV and tag
	encrypted, err := jwebb.KeyEncryptAESGCMKW(cek, tokens.A128GCMKW, sharedkey)
	require.NoError(t, err)

	// Extract the IV and tag from the encrypted result
	// AES-GCMKW encryption returns a ByteWithIVAndTag
	encryptedWithIVTag, ok := encrypted.(keygen.ByteWithIVAndTag)
	require.True(t, ok, "encrypted result should have IV and tag")

	// Now decrypt using the proper IV and tag
	decrypted, err := jwebb.KeyDecryptAESGCMKW(encryptedWithIVTag.Bytes(), encryptedWithIVTag.Bytes(), tokens.A128GCMKW, sharedkey, encryptedWithIVTag.IV, encryptedWithIVTag.Tag)
	require.NoError(t, err)
	require.Equal(t, cek, decrypted)
}

func TestKeyDecryptAESGCMKW_DoesNotMutateRecipientKey(t *testing.T) {
	cek := testCEK
	sharedkey := testSharedKey16

	encrypted, err := jwebb.KeyEncryptAESGCMKW(cek, tokens.A128GCMKW, sharedkey)
	require.NoError(t, err)
	encryptedWithIVTag, ok := encrypted.(keygen.ByteWithIVAndTag)
	require.True(t, ok, "encrypted result should have IV and tag")

	// Put the encrypted key in a backing array with spare capacity and
	// fill the tail with a sentinel pattern. The buggy
	// append(recipientKey[:], tag...) would overwrite the tail.
	src := encryptedWithIVTag.Bytes()
	backing := make([]byte, len(src)+len(encryptedWithIVTag.Tag)+16)
	copy(backing, src)
	for i := len(src); i < len(backing); i++ {
		backing[i] = 0xAB
	}
	recipientKey := backing[:len(src):len(backing)]
	snapshot := append([]byte(nil), backing...)

	decrypted, err := jwebb.KeyDecryptAESGCMKW(recipientKey, recipientKey, tokens.A128GCMKW, sharedkey, encryptedWithIVTag.IV, encryptedWithIVTag.Tag)
	require.NoError(t, err)
	require.Equal(t, cek, decrypted)
	require.Equal(t, snapshot, backing, "recipientKey backing array must not be mutated")
}

func TestKeyDecryptRSA15(t *testing.T) {
	cek := testCEK

	// Generate RSA key pair
	privkey, err := jwxtest.GenerateRsaKey()
	require.NoError(t, err)
	pubkey := &privkey.PublicKey
	require.NoError(t, err)

	// First encrypt
	encrypted, err := jwebb.KeyEncryptRSA15(cek, tokens.RSA1_5, pubkey)
	require.NoError(t, err)

	// Then decrypt
	decrypted, err := jwebb.KeyDecryptRSA15(encrypted.Bytes(), encrypted.Bytes(), privkey, 8) // 8 bytes keysize for test
	require.NoError(t, err)
	require.NotNil(t, decrypted)
	require.Len(t, decrypted, 16) // Should be 8*2 = 16 bytes
}

func TestKeyDecryptRSAOAEP(t *testing.T) {
	cek := testCEK

	// Generate RSA key pair
	privkey, err := jwxtest.GenerateRsaKey()
	require.NoError(t, err)
	pubkey := &privkey.PublicKey
	require.NoError(t, err)

	// First encrypt
	encrypted, err := jwebb.KeyEncryptRSAOAEP(cek, tokens.RSA_OAEP, pubkey)
	require.NoError(t, err)

	// Then decrypt
	decrypted, err := jwebb.KeyDecryptRSAOAEP(encrypted.Bytes(), encrypted.Bytes(), tokens.RSA_OAEP, privkey)
	require.NoError(t, err)
	require.Equal(t, cek, decrypted)
}

func TestKeyDecryptECDHES(t *testing.T) {
	// Generate ECDSA key pairs for Alice and Bob
	alicePriv, err := jwxtest.GenerateEcdsaKey(jwa.P256())
	require.NoError(t, err)
	bobPriv, err := jwxtest.GenerateEcdsaKey(jwa.P256())
	require.NoError(t, err)

	apu := testAPU
	apv := testAPV

	// Test ECDH-ES direct key agreement (no key wrapping)
	decrypted, err := jwebb.KeyDecryptECDHES(nil, nil, tokens.ECDH_ES, apu, apv, alicePriv, &bobPriv.PublicKey, 16)
	require.NoError(t, err)
	require.NotNil(t, decrypted)
	require.Len(t, decrypted, 16)
}

func TestKeyDecryptECDHESKeyWrap(t *testing.T) {
	cek := testCEK

	// Generate ECDSA key pairs - use the same key for both sides to ensure compatibility
	privkey, err := jwxtest.GenerateEcdsaKey(jwa.P256())
	require.NoError(t, err)
	pubkey := &privkey.PublicKey

	apu := testAPU
	apv := testAPV

	// First encrypt using ECDH-ES+A128KW
	encrypted, err := jwebb.KeyEncryptECDHESKeyWrapECDSA(cek, tokens.ECDH_ES_A128KW, apu, apv, pubkey, 16, tokens.A128GCM)
	require.NoError(t, err)

	// Extract the public key from the encrypted result
	// ECDH-ES encryption returns a ByteWithECPublicKey
	encryptedWithPubKey, ok := encrypted.(keygen.ByteWithECPublicKey)
	require.True(t, ok, "encrypted result should have public key")

	// Now decrypt using the corresponding private key and the ephemeral public key
	decrypted, err := jwebb.KeyDecryptECDHESKeyWrap(encryptedWithPubKey.Bytes(), encryptedWithPubKey.Bytes(), tokens.ECDH_ES_A128KW, apu, apv, privkey, encryptedWithPubKey.PublicKey, 16)
	require.NoError(t, err)
	require.Equal(t, cek, decrypted)
}

func TestDeriveECDHES(t *testing.T) {
	// Generate ECDSA key pairs
	alicePriv, err := jwxtest.GenerateEcdsaKey(jwa.P256())
	require.NoError(t, err)
	bobPriv, err := jwxtest.GenerateEcdsaKey(jwa.P256())
	require.NoError(t, err)

	apu := testAPU
	apv := testAPV

	// Test key derivation
	key1, err := jwebb.DeriveECDHES(tokens.A128GCM, apu, apv, alicePriv, &bobPriv.PublicKey, 16)
	require.NoError(t, err)
	require.Len(t, key1, 16)

	// Test that the same inputs produce the same key
	key2, err := jwebb.DeriveECDHES(tokens.A128GCM, apu, apv, alicePriv, &bobPriv.PublicKey, 16)
	require.NoError(t, err)
	require.Equal(t, key1, key2)

	// Test that different private keys produce different results
	charliePriv, err := jwxtest.GenerateEcdsaKey(jwa.P256())
	require.NoError(t, err)
	key3, err := jwebb.DeriveECDHES(tokens.A128GCM, apu, apv, charliePriv, &bobPriv.PublicKey, 16)
	require.NoError(t, err)
	require.NotEqual(t, key1, key3)
}