package handlers

import (
	"context"
	"crypto/rand"
	"encoding/json"
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/yan1h/agent-coding-workflow/internal/infra/crypto"
	"github.com/yan1h/agent-coding-workflow/internal/jobs"
)

// memRow models one encrypted-column row across versions.
type memRow struct {
	blob       []byte
	keyVersion int
	plaintext  []byte // ground truth for assertions
	null       bool
}

// memStore is an in-memory ReencryptStore over a single logical table, with a
// failNext hook to model a crash mid-run for the resumability test.
type memStore struct {
	rows     []*memRow
	failNext bool
}

func (m *memStore) Tables() []string { return []string{"t"} }

func (m *memStore) StaleCount(_ context.Context, _ string, active int) (int, error) {
	n := 0
	for _, r := range m.rows {
		if !r.null && r.keyVersion < active {
			n++
		}
	}
	return n, nil
}

func (m *memStore) ReencryptTable(ctx context.Context, enc *crypto.KeyedEncryptor, _ string, active, batch int) (int, error) {
	done := 0
	for _, r := range m.rows {
		if done >= batch {
			break
		}
		if r.null || r.keyVersion >= active {
			continue
		}
		if m.failNext {
			m.failNext = false
			return done, context.Canceled // simulate a crash after some rows
		}
		plain, err := enc.DecryptVersion(ctx, r.blob, r.keyVersion)
		if err != nil {
			return done, err
		}
		sealed, err := enc.SealWithVersion(ctx, active, plain)
		if err != nil {
			return done, err
		}
		r.blob = sealed
		r.keyVersion = active
		done++
	}
	return done, nil
}

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

// twoVersionEnc returns a KeyedEncryptor with v1+v2 keys and a store reporting
// active=2, plus a helper to seal under v1.
func twoVersionEnc(t *testing.T) *crypto.KeyedEncryptor {
	prov := crypto.NewEnvProvider(nil)
	prov.SetKey(1, key(t))
	prov.SetKey(2, key(t))
	return crypto.NewKeyedEncryptor(prov, crypto.NewStaticKeyStore(2))
}

func TestSecretReencrypt_AdvancesAllRows(t *testing.T) {
	ctx := context.Background()
	enc := twoVersionEnc(t)

	// Seed 3 rows encrypted under v1 + 1 NULL row.
	store := &memStore{}
	for i := 0; i < 3; i++ {
		pt := []byte{byte(i), byte(i + 1)}
		ct, err := enc.SealWithVersion(ctx, 1, pt)
		require.NoError(t, err)
		store.rows = append(store.rows, &memRow{blob: ct, keyVersion: 1, plaintext: pt})
	}
	store.rows = append(store.rows, &memRow{null: true, keyVersion: 1})

	h := NewSecretReencryptHandler(store, enc, nil)
	payload, _ := json.Marshal(SecretReencryptPayload{ActiveVersion: 2, BatchSize: 2})
	require.NoError(t, h.Handle(ctx, jobs.Job{Payload: payload}))

	// All non-null rows now on v2 and decrypt to identical plaintext.
	for _, r := range store.rows {
		if r.null {
			require.Equal(t, 1, r.keyVersion, "NULL rows are skipped, version unchanged")
			continue
		}
		require.Equal(t, 2, r.keyVersion)
		got, err := enc.DecryptVersion(ctx, r.blob, 2)
		require.NoError(t, err)
		require.Equal(t, r.plaintext, got)
	}
	left, _ := store.StaleCount(ctx, "t", 2)
	require.Equal(t, 0, left)
}

func TestSecretReencrypt_IdempotentRerun(t *testing.T) {
	ctx := context.Background()
	enc := twoVersionEnc(t)
	store := &memStore{}
	pt := []byte("hello")
	ct, err := enc.SealWithVersion(ctx, 1, pt)
	require.NoError(t, err)
	store.rows = append(store.rows, &memRow{blob: ct, keyVersion: 1, plaintext: pt})

	h := NewSecretReencryptHandler(store, enc, nil)
	payload, _ := json.Marshal(SecretReencryptPayload{ActiveVersion: 2})
	require.NoError(t, h.Handle(ctx, jobs.Job{Payload: payload}))
	blobAfterFirst := store.rows[0].blob

	// Second run is a no-op: row already on active version, blob unchanged.
	require.NoError(t, h.Handle(ctx, jobs.Job{Payload: payload}))
	require.Equal(t, blobAfterFirst, store.rows[0].blob)
}

func TestSecretReencrypt_ResumableAfterCrash(t *testing.T) {
	ctx := context.Background()
	enc := twoVersionEnc(t)
	store := &memStore{failNext: true}
	for i := 0; i < 5; i++ {
		pt := []byte{byte(i)}
		ct, err := enc.SealWithVersion(ctx, 1, pt)
		require.NoError(t, err)
		store.rows = append(store.rows, &memRow{blob: ct, keyVersion: 1, plaintext: pt})
	}

	h := NewSecretReencryptHandler(store, enc, nil)
	payload, _ := json.Marshal(SecretReencryptPayload{ActiveVersion: 2, BatchSize: 10})
	// First run crashes (ctx.Canceled bubbles up).
	require.Error(t, h.Handle(ctx, jobs.Job{Payload: payload}))
	left, _ := store.StaleCount(ctx, "t", 2)
	require.Greater(t, left, 0, "crash left some rows on old version")

	// Re-run completes the rest.
	require.NoError(t, h.Handle(ctx, jobs.Job{Payload: payload}))
	left, _ = store.StaleCount(ctx, "t", 2)
	require.Equal(t, 0, left)
	for _, r := range store.rows {
		got, err := enc.DecryptVersion(ctx, r.blob, 2)
		require.NoError(t, err)
		require.Equal(t, r.plaintext, got)
	}
}

func TestSecretReencrypt_BadPayloadIsPermanent(t *testing.T) {
	h := NewSecretReencryptHandler(&memStore{}, twoVersionEnc(t), nil)
	err := h.Handle(context.Background(), jobs.Job{Payload: []byte("not json")})
	require.Error(t, err)
	require.True(t, jobs.IsPermanent(err))
}