package crypto

import (
	"context"
	"crypto/rand"
	"sync"
	"testing"

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

// mutableStore is a KeyStore whose active version can change at runtime,
// modeling a rotation.
type mutableStore struct {
	mu sync.Mutex
	v  int
}

func (m *mutableStore) ActiveVersion(context.Context) (int, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.v, nil
}

func (m *mutableStore) set(v int) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.v = v
}

func rawKey(t *testing.T) []byte {
	t.Helper()
	k := make([]byte, 32)
	_, err := rand.Read(k)
	require.NoError(t, err)
	return k
}

func TestKeyedEncryptor_RoundTripV1AndV2(t *testing.T) {
	ctx := context.Background()
	prov := NewEnvProvider(nil)
	prov.SetKey(1, rawKey(t))
	prov.SetKey(2, rawKey(t))
	store := &mutableStore{v: 1}
	ke := NewKeyedEncryptor(prov, store)

	pt := []byte("super secret value")

	ctV1, v, err := ke.EncryptVersion(ctx, pt)
	require.NoError(t, err)
	require.Equal(t, 1, v)
	got, err := ke.DecryptVersion(ctx, ctV1, 1)
	require.NoError(t, err)
	require.Equal(t, pt, got)

	store.set(2)
	ctV2, v2, err := ke.EncryptVersion(ctx, pt)
	require.NoError(t, err)
	require.Equal(t, 2, v2)
	require.NotEqual(t, ctV1, ctV2)
	got2, err := ke.DecryptVersion(ctx, ctV2, 2)
	require.NoError(t, err)
	require.Equal(t, pt, got2)
}

func TestKeyedEncryptor_V1BlobDecryptsAfterActiveBumpsToV2(t *testing.T) {
	ctx := context.Background()
	prov := NewEnvProvider(nil)
	prov.SetKey(1, rawKey(t))
	prov.SetKey(2, rawKey(t))
	store := &mutableStore{v: 1}
	ke := NewKeyedEncryptor(prov, store)

	pt := []byte("encrypted under v1")
	ctV1, v, err := ke.EncryptVersion(ctx, pt)
	require.NoError(t, err)
	require.Equal(t, 1, v)

	// Rotate the active write version; the OLD blob must still decrypt under v1.
	store.set(2)
	got, err := ke.DecryptVersion(ctx, ctV1, 1)
	require.NoError(t, err)
	require.Equal(t, pt, got)
}

func TestKeyedEncryptor_WrongVersionDecryptFailsCleanly(t *testing.T) {
	ctx := context.Background()
	prov := NewEnvProvider(nil)
	prov.SetKey(1, rawKey(t))
	prov.SetKey(2, rawKey(t))
	ke := NewKeyedEncryptor(prov, NewStaticKeyStore(1))

	ct, _, err := ke.EncryptVersion(ctx, []byte("x"))
	require.NoError(t, err)
	// Decrypting a v1 blob as v2 must fail authentication, not panic.
	_, err = ke.DecryptVersion(ctx, ct, 2)
	require.Error(t, err)
}

func TestKeyedEncryptor_UnknownVersionErrors(t *testing.T) {
	ctx := context.Background()
	prov := NewEnvProvider(nil)
	prov.SetKey(1, rawKey(t))
	ke := NewKeyedEncryptor(prov, NewStaticKeyStore(1))
	_, err := ke.DecryptVersion(ctx, make([]byte, 16), 9)
	require.Error(t, err)
}

func TestEncryptorShim_BacksOntoKeyedV1(t *testing.T) {
	// The legacy Encryptor shim must remain byte-compatible: a blob it produces
	// decrypts via the underlying KeyedEncryptor at version 1, and vice versa.
	key := rawKey(t)
	enc, err := NewEncryptor(key)
	require.NoError(t, err)

	pt := []byte("legacy callers see no versions")
	ct, err := enc.Encrypt(pt)
	require.NoError(t, err)

	got, err := enc.Decrypt(ct)
	require.NoError(t, err)
	require.Equal(t, pt, got)

	// Same key via KeyedEncryptor v1 opens the shim's ciphertext.
	prov := NewEnvProvider(key)
	ke := NewKeyedEncryptor(prov, NewStaticKeyStore(1))
	got2, err := ke.DecryptVersion(context.Background(), ct, 1)
	require.NoError(t, err)
	require.Equal(t, pt, got2)
}

func TestEnvProvider_ReadsVersionedEnv(t *testing.T) {
	t.Setenv("APP_MASTER_KEY_V2", "00112233445566778899aabbccddeeff00112233445566778899aabbccddeeff")
	prov := NewEnvProvider(nil)
	k, err := prov.KeyMaterial(context.Background(), 2)
	require.NoError(t, err)
	require.Len(t, k, 32)

	_, err = prov.KeyMaterial(context.Background(), 3)
	require.Error(t, err, "unstaged version must error, not silently use a default")
}