package envelope_test

import (
	"context"
	cryptorand "crypto/rand"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"testing"

	"atlas9.dev/c/core"
	"atlas9.dev/c/core/assert"
	"atlas9.dev/c/core/envelope"
)

// memDekStore is an in-memory envelope.DekStore for tests.
type memDekStore struct {
	byTenant map[core.ID]envelope.DEK
	creates  int
}

func newMemDekStore() *memDekStore {

	return &memDekStore{byTenant: map[core.ID]envelope.DEK{}}
}

func (m *memDekStore) ForTenant(ctx context.Context, tenant core.ID) (envelope.DEK, error) {

	dek, ok := m.byTenant[tenant]
	if !ok {
		return envelope.DEK{}, core.ErrNotFound
	}
	return dek, nil
}

func (m *memDekStore) Create(ctx context.Context, tenant core.ID, wrappedKey []byte) (envelope.DEK, error) {

	m.creates++
	dek := envelope.DEK{ID: core.NewID("dek"), WrappedKey: wrappedKey}
	m.byTenant[tenant] = dek
	return dek, nil
}

// countingWrapper wraps another Wrapper and counts Wrap/Unwrap calls.
type countingWrapper struct {
	inner   envelope.Wrapper
	wraps   int
	unwraps int
}

func (c *countingWrapper) Wrap(ctx context.Context, dek []byte) ([]byte, error) {

	c.wraps++
	return c.inner.Wrap(ctx, dek)
}

func (c *countingWrapper) Unwrap(ctx context.Context, blob []byte) ([]byte, error) {

	c.unwraps++
	return c.inner.Unwrap(ctx, blob)
}

func testWrapper(t *testing.T) envelope.Wrapper {

	t.Helper()
	path := filepath.Join(t.TempDir(), "kek")
	var key [32]byte
	if _, err := cryptorand.Read(key[:]); err != nil {
		t.Fatal(err)
	}
	assert.Ok(t, os.WriteFile(path, []byte(fmt.Sprintf("%x\n", key)), 0600))
	w, err := envelope.LoadKeyFile(path)
	assert.Ok(t, err)
	return w
}

func TestEncryptor_RoundTrip(t *testing.T) {

	ctx := context.Background()
	store := newMemDekStore()
	wrapper := testWrapper(t)
	factory := envelope.NewEncryptorFactory(wrapper)
	tenant := core.NewID("t")
	assert.Ok(t, envelope.Provision(ctx, store, wrapper, tenant))

	enc, err := factory.For(ctx, store, tenant)
	assert.Ok(t, err)

	_, ct, err := enc.Seal([]byte("client-secret"))
	assert.Ok(t, err)

	got, err := enc.Open(ct)
	assert.Ok(t, err)
	assert.Eq(t, string(got), "client-secret")
}

// For fails when the tenant was never provisioned — DEKs are not created lazily.
func TestEncryptor_ForRequiresProvision(t *testing.T) {

	ctx := context.Background()
	factory := envelope.NewEncryptorFactory(testWrapper(t))

	_, err := factory.For(ctx, newMemDekStore(), core.NewID("t"))
	assert.Eq(t, errors.Is(err, core.ErrNotFound), true)
}

// Provision creates one DEK per tenant; distinct tenants get distinct DEKs.
func TestEncryptor_ProvisionPerTenant(t *testing.T) {

	ctx := context.Background()
	store := newMemDekStore()
	wrapper := testWrapper(t)
	acme := core.NewID("t")
	skunk := core.NewID("t")

	assert.Ok(t, envelope.Provision(ctx, store, wrapper, acme))
	assert.Eq(t, store.creates, 1)
	assert.Ok(t, envelope.Provision(ctx, store, wrapper, skunk))
	assert.Eq(t, store.creates, 2)
}

// A tenant encryptor unwraps its DEK once (when built) and reuses it, so
// sealing/opening many values does not unwrap again.
func TestEncryptor_UnwrapsDekOnce(t *testing.T) {

	ctx := context.Background()
	store := newMemDekStore()
	counter := &countingWrapper{inner: testWrapper(t)}
	factory := envelope.NewEncryptorFactory(counter)
	tenant := core.NewID("t")
	assert.Ok(t, envelope.Provision(ctx, store, counter, tenant))

	enc, err := factory.For(ctx, store, tenant)
	assert.Ok(t, err)

	var cts [][]byte
	for _, s := range []string{"one", "two", "three"} {
		_, ct, err := enc.Seal([]byte(s))
		assert.Ok(t, err)
		cts = append(cts, ct)
	}
	for _, ct := range cts {
		_, err := enc.Open(ct)
		assert.Ok(t, err)
	}

	// One Wrap at Provision, one Unwrap at For, none during Seal/Open.
	assert.Eq(t, counter.wraps, 1)
	assert.Eq(t, counter.unwraps, 1)
}

// IsSealed accepts real ciphertext and rejects plaintext / junk — the guard the
// store uses so plaintext can't reach the ciphertext column.
func TestIsSealed(t *testing.T) {

	ctx := context.Background()
	store := newMemDekStore()
	wrapper := testWrapper(t)
	factory := envelope.NewEncryptorFactory(wrapper)
	tenant := core.NewID("t")
	assert.Ok(t, envelope.Provision(ctx, store, wrapper, tenant))
	enc, err := factory.For(ctx, store, tenant)
	assert.Ok(t, err)

	_, ct, err := enc.Seal([]byte("secret"))
	assert.Ok(t, err)

	assert.Eq(t, envelope.IsSealed(ct), true)
	assert.Eq(t, envelope.IsSealed([]byte("plaintext-not-sealed")), false)
	assert.Eq(t, envelope.IsSealed(nil), false)
}

func TestEncryptor_OpenRejectsBadInput(t *testing.T) {

	ctx := context.Background()
	store := newMemDekStore()
	wrapper := testWrapper(t)
	factory := envelope.NewEncryptorFactory(wrapper)
	tenant := core.NewID("t")
	assert.Ok(t, envelope.Provision(ctx, store, wrapper, tenant))

	enc, err := factory.For(ctx, store, tenant)
	assert.Ok(t, err)

	_, ct, err := enc.Seal([]byte("secret"))
	assert.Ok(t, err)

	// An encryptor for a different tenant (different DEK) can't open it — the
	// GCM auth tag rejects the wrong key.
	other := core.NewID("t")
	assert.Ok(t, envelope.Provision(ctx, store, wrapper, other))
	otherEnc, err := factory.For(ctx, store, other)
	assert.Ok(t, err)
	_, err = otherEnc.Open(ct)
	assert.Eq(t, err != nil, true)

	// Plaintext bytes (no magic header).
	_, err = enc.Open([]byte("plaintext-not-sealed"))
	assert.Eq(t, err != nil, true)

	// Corrupted header byte.
	badHeader := append([]byte(nil), ct...)
	badHeader[0] ^= 0xFF
	_, err = enc.Open(badHeader)
	assert.Eq(t, err != nil, true)

	// Tampered ciphertext keeps a valid header but fails the GCM auth tag.
	tampered := append([]byte(nil), ct...)
	tampered[len(tampered)-1] ^= 0xFF
	_, err = enc.Open(tampered)
	assert.Eq(t, err != nil, true)
}