package storagewrappers

import (
	"context"
	"fmt"
	"strconv"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"go.uber.org/goleak"
	"go.uber.org/mock/gomock"
	"golang.org/x/sync/singleflight"
	"google.golang.org/protobuf/types/known/structpb"

	openfgav1 "github.com/openfga/api/proto/openfga/v1"

	"github.com/openfga/openfga/internal/mocks"
	"github.com/openfga/openfga/pkg/storage"
	"github.com/openfga/openfga/pkg/storage/cache/keys"
	"github.com/openfga/openfga/pkg/tuple"
)

// ─────────────────────────────────────────────────────────────────────────────
// Test Helpers
// ─────────────────────────────────────────────────────────────────────────────

// blockingIterator is a test helper that returns pre-loaded tuples, then
// blocks on Next()/Head() until the context is cancelled or Stop() is called.
// Callers must ensure Stop() is eventually called (or the context cancelled)
// to avoid hanging the test.
type blockingIterator struct {
	tuples  []*openfgav1.Tuple
	index   int
	blockCh chan struct{}
}

// Next returns the next pre-loaded tuple. Once all tuples are consumed, it
// blocks until ctx is cancelled or Stop() closes blockCh. Tests that use a
// short drainTimeout rely on this blocking to trigger context expiry.
func (b *blockingIterator) Next(ctx context.Context) (*openfgav1.Tuple, error) {
	if b.index < len(b.tuples) {
		t := b.tuples[b.index]
		b.index++
		return t, nil
	}
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case <-b.blockCh:
		return nil, storage.ErrIteratorDone
	}
}

func (b *blockingIterator) Head(ctx context.Context) (*openfgav1.Tuple, error) {
	if b.index < len(b.tuples) {
		return b.tuples[b.index], nil
	}
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case <-b.blockCh:
		return nil, storage.ErrIteratorDone
	}
}

func (b *blockingIterator) Stop() {
	select {
	case <-b.blockCh:
	default:
		close(b.blockCh)
	}
}

func (b *blockingIterator) IsOrdered() bool { return true }

// ─────────────────────────────────────────────────────────────────────────────
// LockFreeCachedIterator Tests
// ─────────────────────────────────────────────────────────────────────────────

func TestLockFreeCachedIterator_IsOrdered(t *testing.T) {
	iter := NewLockFreeCachedIterator([]MinimalCacheEntry{}, "document", "viewer", true)
	defer iter.Stop()
	require.True(t, iter.IsOrdered())

	iter2 := NewLockFreeCachedIterator([]MinimalCacheEntry{}, "document", "viewer", false)
	defer iter2.Stop()
	require.False(t, iter2.IsOrdered())
}

func TestLockFreeCachedIterator_Next_Basic(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	entries := []MinimalCacheEntry{
		{ObjectID: "1", User: "user:alice"},
		{ObjectID: "2", User: "user:bob"},
		{ObjectID: "3", User: "user:charlie"},
	}

	iter := NewLockFreeCachedIterator(entries, "document", "viewer", false)

	ctx := context.Background()

	// Verify all entries can be retrieved
	t1, err := iter.Next(ctx)
	require.NoError(t, err)
	require.Equal(t, "document:1", t1.GetKey().GetObject())
	require.Equal(t, "viewer", t1.GetKey().GetRelation())
	require.Equal(t, "user:alice", t1.GetKey().GetUser())

	t2, err := iter.Next(ctx)
	require.NoError(t, err)
	require.Equal(t, "document:2", t2.GetKey().GetObject())
	require.Equal(t, "user:bob", t2.GetKey().GetUser())

	t3, err := iter.Next(ctx)
	require.NoError(t, err)
	require.Equal(t, "document:3", t3.GetKey().GetObject())
	require.Equal(t, "user:charlie", t3.GetKey().GetUser())

	// Should return ErrIteratorDone after exhaustion
	_, err = iter.Next(ctx)
	require.ErrorIs(t, err, storage.ErrIteratorDone)
}

func TestLockFreeCachedIterator_Next_Empty(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	iter := NewLockFreeCachedIterator([]MinimalCacheEntry{}, "document", "viewer", false)

	ctx := context.Background()
	_, err := iter.Next(ctx)
	require.ErrorIs(t, err, storage.ErrIteratorDone)
}

func TestLockFreeCachedIterator_Next_WithCondition(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	condCtx, _ := structpb.NewStruct(map[string]interface{}{"region": "us-east"})
	entries := []MinimalCacheEntry{
		{
			ObjectID:         "1",
			User:             "user:alice",
			ConditionName:    "is_allowed",
			ConditionContext: condCtx,
		},
	}

	iter := NewLockFreeCachedIterator(entries, "document", "viewer", false)

	ctx := context.Background()
	t1, err := iter.Next(ctx)
	require.NoError(t, err)
	require.Equal(t, "is_allowed", t1.GetKey().GetCondition().GetName())
	require.Equal(t, condCtx, t1.GetKey().GetCondition().GetContext())
}

func TestLockFreeCachedIterator_Head_Basic(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	entries := []MinimalCacheEntry{
		{ObjectID: "1", User: "user:alice"},
		{ObjectID: "2", User: "user:bob"},
	}

	iter := NewLockFreeCachedIterator(entries, "document", "viewer", false)

	ctx := context.Background()

	// Head should return first element without advancing
	t1, err := iter.Head(ctx)
	require.NoError(t, err)
	require.Equal(t, "document:1", t1.GetKey().GetObject())

	// Second Head call should return the same element
	t2, err := iter.Head(ctx)
	require.NoError(t, err)
	require.Equal(t, "document:1", t2.GetKey().GetObject())

	// Next should also return the same (first) element
	t3, err := iter.Next(ctx)
	require.NoError(t, err)
	require.Equal(t, "document:1", t3.GetKey().GetObject())

	// Now Next should advance to second element
	t4, err := iter.Next(ctx)
	require.NoError(t, err)
	require.Equal(t, "document:2", t4.GetKey().GetObject())
}

func TestLockFreeCachedIterator_Head_AfterStop(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	entries := []MinimalCacheEntry{
		{ObjectID: "1", User: "user:alice"},
	}

	iter := NewLockFreeCachedIterator(entries, "document", "viewer", false)

	iter.Stop()

	_, err := iter.Head(context.Background())
	require.ErrorIs(t, err, storage.ErrIteratorDone)
}

func TestLockFreeCachedIterator_Head_PastEnd(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	entries := []MinimalCacheEntry{
		{ObjectID: "1", User: "user:alice"},
	}

	iter := NewLockFreeCachedIterator(entries, "document", "viewer", false)

	ctx := context.Background()

	// Consume the only entry
	_, err := iter.Next(ctx)
	require.NoError(t, err)

	// Head should return ErrIteratorDone since index >= len(entries)
	_, err = iter.Head(ctx)
	require.ErrorIs(t, err, storage.ErrIteratorDone)
}

func TestLockFreeCachedIterator_Concurrent_Next(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	entries := make([]MinimalCacheEntry, 100)
	for i := 0; i < 100; i++ {
		entries[i] = MinimalCacheEntry{
			ObjectID: string(rune('a' + i%26)),
			User:     "user:test",
		}
	}

	iter := NewLockFreeCachedIterator(entries, "document", "viewer", false)

	ctx := context.Background()
	var wg sync.WaitGroup
	results := make(chan *openfgav1.Tuple, 200)

	// Start multiple goroutines calling Next concurrently
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for {
				t, err := iter.Next(ctx)
				if err != nil {
					return
				}
				results <- t
			}
		}()
	}

	wg.Wait()
	close(results)

	// Count results - should be exactly 100 (no duplicates, no misses)
	count := 0
	for range results {
		count++
	}
	require.Equal(t, 100, count)
}

func TestLockFreeCachedIterator_Concurrent_Stop(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	entries := make([]MinimalCacheEntry, 100)
	for i := 0; i < 100; i++ {
		entries[i] = MinimalCacheEntry{ObjectID: "1", User: "user:test"}
	}

	iter := NewLockFreeCachedIterator(entries, "document", "viewer", false)

	ctx := context.Background()
	var wg sync.WaitGroup

	// Start goroutines calling Next
	for i := 0; i < 5; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for {
				_, err := iter.Next(ctx)
				if err != nil {
					return
				}
			}
		}()
	}

	// Stop from another goroutine
	iter.Stop()

	wg.Wait()

	// After Stop, Next should return ErrIteratorDone
	_, err := iter.Next(ctx)
	require.ErrorIs(t, err, storage.ErrIteratorDone)
}

func TestLockFreeCachedIterator_ContextCanceled(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	entries := []MinimalCacheEntry{
		{ObjectID: "1", User: "user:alice"},
	}

	iter := NewLockFreeCachedIterator(entries, "document", "viewer", false)

	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	_, err := iter.Next(ctx)
	require.ErrorIs(t, err, context.Canceled)

	_, err = iter.Head(ctx)
	require.ErrorIs(t, err, context.Canceled)
}

// ─────────────────────────────────────────────────────────────────────────────
// CachingIterator Tests
// ─────────────────────────────────────────────────────────────────────────────

func TestCachingIterator_IsOrdered(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()
	mockCache := mocks.NewMockInMemoryCache[any](ctrl)
	mockCache.EXPECT().Get(gomock.Any()).Return(nil).AnyTimes()
	mockCache.EXPECT().Set(gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes()

	sf := &singleflight.Group{}
	wg := &sync.WaitGroup{}

	inner := storage.NewStaticTupleIterator([]*openfgav1.Tuple{})
	iter := newCachingIterator(inner, mockCache, testCacheKey("key"), 100, time.Hour, 30*time.Second, sf, wg, "document", "viewer", "Read")
	require.True(t, iter.IsOrdered())
	iter.Stop()
	wg.Wait()
}

func TestCachingIterator_Next_Basic(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	mockController := gomock.NewController(t)
	defer mockController.Finish()

	mockCache := mocks.NewMockInMemoryCache[any](mockController)

	ctx := context.Background()
	sf := &singleflight.Group{}
	wg := &sync.WaitGroup{}

	tuples := []*openfgav1.Tuple{
		{Key: tuple.NewTupleKey("document:1", "viewer", "user:alice")},
		{Key: tuple.NewTupleKey("document:2", "viewer", "user:bob")},
	}

	// Use static iterator
	innerIter := storage.NewStaticTupleIterator(tuples)

	// Expect cache.Get to check if already cached (optimization 1), return nil (not cached)
	mockCache.EXPECT().Get(testCacheKey("test-key")).Return(nil).Times(1)
	// Expect cache.Set when Stop() is called after iterator is exhausted
	mockCache.EXPECT().Set(gomock.Any(), gomock.Any(), gomock.Any()).Times(1)

	iter := newCachingIterator(
		innerIter, mockCache, testCacheKey("test-key"), 1000, time.Hour, 30*time.Second,
		sf, wg, "document", "viewer", "ReadUsersetTuples",
	)

	t1, err := iter.Next(ctx)
	require.NoError(t, err)
	require.Equal(t, "document:1", t1.GetKey().GetObject())

	t2, err := iter.Next(ctx)
	require.NoError(t, err)
	require.Equal(t, "document:2", t2.GetKey().GetObject())

	// Iterator done
	_, err = iter.Next(ctx)
	require.ErrorIs(t, err, storage.ErrIteratorDone)

	// Stop triggers flush to cache (now async via goroutine)
	iter.Stop()

	// Wait for background goroutine to complete
	wg.Wait()
}

// TestCachingIterator_Next_NonIteratorDoneError_NilsTuples verifies that a
// non-done/non-cancelled error from the inner iterator causes tuples to be
// set to nil, preventing caching of incomplete results.
func TestCachingIterator_Next_NonIteratorDoneError_NilsTuples(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	mockController := gomock.NewController(t)
	defer mockController.Finish()

	mockCache := mocks.NewMockInMemoryCache[any](mockController)

	ctx := context.Background()
	sf := &singleflight.Group{}
	wg := &sync.WaitGroup{}

	// errorIterator returns first tuple, then ErrSimulatedError on second Next()
	innerIter := mocks.NewErrorTupleIterator([]*openfgav1.Tuple{
		{Key: tuple.NewTupleKey("document:1", "viewer", "user:alice")},
	})

	iter := newCachingIterator(
		innerIter, mockCache, testCacheKey("test-key"), 1000, time.Hour, 30*time.Second,
		sf, wg, "document", "viewer", "ReadUsersetTuples",
	)

	// First call succeeds
	t1, err := iter.Next(ctx)
	require.NoError(t, err)
	require.Equal(t, "document:1", t1.GetKey().GetObject())
	require.NotNil(t, iter.tuples) // tuples still collecting

	// Second call returns simulated error (not done/cancelled)
	_, err = iter.Next(ctx)
	require.ErrorIs(t, err, mocks.ErrSimulatedError)

	// tuples should be nil — incomplete results must not be cached
	require.Nil(t, iter.tuples)

	// Stop should not panic and should call wg.Done
	iter.Stop()
	wg.Wait()
}

func TestCachingIterator_Head_DelegatesToInner(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	mockController := gomock.NewController(t)
	defer mockController.Finish()

	mockCache := mocks.NewMockInMemoryCache[any](mockController)
	mockCache.EXPECT().Get(gomock.Any()).Return(nil).AnyTimes()
	mockCache.EXPECT().Set(gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes()

	ctx := context.Background()
	sf := &singleflight.Group{}
	wg := &sync.WaitGroup{}

	tuples := []*openfgav1.Tuple{
		{Key: tuple.NewTupleKey("document:1", "viewer", "user:alice")},
		{Key: tuple.NewTupleKey("document:2", "viewer", "user:bob")},
	}

	innerIter := storage.NewStaticTupleIterator(tuples)

	iter := newCachingIterator(
		innerIter, mockCache, testCacheKey("test-key"), 1000, time.Hour, 30*time.Second,
		sf, wg, "document", "viewer", "ReadUsersetTuples",
	)

	// Head returns first tuple without advancing
	t1, err := iter.Head(ctx)
	require.NoError(t, err)
	require.Equal(t, "document:1", t1.GetKey().GetObject())

	// Calling Head again returns the same tuple
	t2, err := iter.Head(ctx)
	require.NoError(t, err)
	require.Equal(t, "document:1", t2.GetKey().GetObject())

	// Next also returns the first tuple (Head didn't advance)
	t3, err := iter.Next(ctx)
	require.NoError(t, err)
	require.Equal(t, "document:1", t3.GetKey().GetObject())

	iter.Stop()
	wg.Wait()
}

func TestCachingIterator_Head_AfterStop(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	mockController := gomock.NewController(t)
	defer mockController.Finish()

	mockCache := mocks.NewMockInMemoryCache[any](mockController)
	mockCache.EXPECT().Get(gomock.Any()).Return(nil).AnyTimes()
	mockCache.EXPECT().Set(gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes()

	ctx := context.Background()
	sf := &singleflight.Group{}
	wg := &sync.WaitGroup{}

	tuples := []*openfgav1.Tuple{
		{Key: tuple.NewTupleKey("document:1", "viewer", "user:alice")},
	}

	innerIter := storage.NewStaticTupleIterator(tuples)

	iter := newCachingIterator(
		innerIter, mockCache, testCacheKey("test-key"), 1000, time.Hour, 30*time.Second,
		sf, wg, "document", "viewer", "ReadUsersetTuples",
	)

	iter.Stop()
	wg.Wait()

	// Head after Stop returns ErrIteratorDone
	_, err := iter.Head(ctx)
	require.ErrorIs(t, err, storage.ErrIteratorDone)
}

func TestCachingIterator_ExceedsMaxSize_AbandonsCaching(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	mockController := gomock.NewController(t)
	defer mockController.Finish()

	mockCache := mocks.NewMockInMemoryCache[any](mockController)

	ctx := context.Background()
	sf := &singleflight.Group{}
	wg := &sync.WaitGroup{}

	maxSize := 3 // Small max size to trigger abandonment

	// Create more tuples than maxSize
	tuples := make([]*openfgav1.Tuple, maxSize+1)
	for i := 0; i < maxSize+1; i++ {
		tuples[i] = &openfgav1.Tuple{Key: tuple.NewTupleKey("document:"+string(rune('1'+i)), "viewer", "user:test")}
	}

	innerIter := storage.NewStaticTupleIterator(tuples)

	// No cache.Set expected because we exceed maxSize

	iter := newCachingIterator(
		innerIter, mockCache, testCacheKey("test-key"), maxSize, time.Hour, 30*time.Second,
		sf, wg, "document", "viewer", "ReadUsersetTuples",
	)

	// Consume all tuples
	for i := 0; i < maxSize+1; i++ {
		_, err := iter.Next(ctx)
		require.NoError(t, err)
	}

	// Verify tuples is nil (abandoned due to exceeding maxSize)
	require.Nil(t, iter.tuples)
}

func TestCachingIterator_ConfigurableMaxSize_AbandonsCaching(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	mockController := gomock.NewController(t)
	defer mockController.Finish()

	mockCache := mocks.NewMockInMemoryCache[any](mockController)

	ctx := context.Background()
	sf := &singleflight.Group{}
	wg := &sync.WaitGroup{}

	customMaxSize := 5

	// Create more tuples than customMaxSize
	tuples := make([]*openfgav1.Tuple, customMaxSize+1)
	for i := 0; i < customMaxSize+1; i++ {
		tuples[i] = &openfgav1.Tuple{Key: tuple.NewTupleKey("document:"+string(rune('1'+i)), "viewer", "user:test")}
	}

	innerIter := storage.NewStaticTupleIterator(tuples)

	iter := newCachingIterator(
		innerIter, mockCache, testCacheKey("test-key"), customMaxSize, time.Hour, 30*time.Second,
		sf, wg, "document", "viewer", "ReadUsersetTuples",
	)

	// Consume all tuples
	for i := 0; i < customMaxSize+1; i++ {
		_, err := iter.Next(ctx)
		require.NoError(t, err)
	}

	// Verify tuples is nil (abandoned due to exceeding maxSize)
	require.Nil(t, iter.tuples)
}

func TestCachingIterator_PopulatesCache(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	mockController := gomock.NewController(t)
	defer mockController.Finish()

	mockCache := mocks.NewMockInMemoryCache[any](mockController)

	ctx := context.Background()
	sf := &singleflight.Group{}
	wg := &sync.WaitGroup{}
	cacheKey := testCacheKey("test-key")
	ttl := time.Hour

	tuples := []*openfgav1.Tuple{
		{Key: tuple.NewTupleKey("document:1", "viewer", "user:alice")},
	}

	innerIter := storage.NewStaticTupleIterator(tuples)

	// Expect cache.Get to check if already cached (optimization 1), return nil (not cached)
	mockCache.EXPECT().Get(cacheKey).Return(nil).Times(1)
	// Capture the cache entry
	var capturedEntry *V2IteratorCacheEntry
	mockCache.EXPECT().Set(cacheKey, gomock.Any(), ttl).DoAndReturn(
		func(_ keys.Key, value interface{}, _ time.Duration) {
			capturedEntry = value.(*V2IteratorCacheEntry)
		},
	)

	iter := newCachingIterator(
		innerIter, mockCache, cacheKey, 1000, ttl, 30*time.Second,
		sf, wg, "document", "viewer", "ReadUsersetTuples",
	)

	// Consume all tuples
	_, err := iter.Next(ctx)
	require.NoError(t, err)

	_, err = iter.Next(ctx)
	require.ErrorIs(t, err, storage.ErrIteratorDone)

	// Stop triggers flush to cache (now async via goroutine)
	iter.Stop()

	// Wait for background goroutine to complete
	wg.Wait()

	// Verify cache was populated correctly
	require.NotNil(t, capturedEntry)
	require.Len(t, capturedEntry.Entries, 1)
	require.Equal(t, "1", capturedEntry.Entries[0].ObjectID)
	require.Equal(t, "user:alice", capturedEntry.Entries[0].User)
}

func TestCachingIterator_InnerError(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	mockController := gomock.NewController(t)
	defer mockController.Finish()

	mockCache := mocks.NewMockInMemoryCache[any](mockController)

	ctx := context.Background()
	sf := &singleflight.Group{}
	wg := &sync.WaitGroup{}

	// Empty iterator returns ErrIteratorDone immediately
	innerIter := storage.NewStaticTupleIterator([]*openfgav1.Tuple{})

	// No cache.Set expected on empty iteration

	iter := newCachingIterator(
		innerIter, mockCache, testCacheKey("test-key"), 1000, time.Hour, 30*time.Second,
		sf, wg, "document", "viewer", "ReadUsersetTuples",
	)

	_, err := iter.Next(ctx)
	require.ErrorIs(t, err, storage.ErrIteratorDone)
}

// TestCachingIterator_Stop_Idempotent verifies that Stop() can be called
// multiple times safely.
func TestCachingIterator_Stop_Idempotent(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	mockController := gomock.NewController(t)
	defer mockController.Finish()

	mockCache := mocks.NewMockInMemoryCache[any](mockController)

	ctx := context.Background()
	sf := &singleflight.Group{}
	wg := &sync.WaitGroup{}

	tuples := []*openfgav1.Tuple{
		{Key: tuple.NewTupleKey("document:1", "viewer", "user:alice")},
	}

	innerIter := storage.NewStaticTupleIterator(tuples)

	// Expect cache operations
	mockCache.EXPECT().Get(testCacheKey("test-key")).Return(nil).AnyTimes()
	mockCache.EXPECT().Set(gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes()

	iter := newCachingIterator(
		innerIter, mockCache, testCacheKey("test-key"), 1000, time.Hour, 30*time.Second,
		sf, wg, "document", "viewer", "ReadUsersetTuples",
	)

	// Consume all tuples
	_, err := iter.Next(ctx)
	require.NoError(t, err)
	_, err = iter.Next(ctx)
	require.ErrorIs(t, err, storage.ErrIteratorDone)

	// Call Stop multiple times - should be idempotent and not panic
	iter.Stop()
	iter.Stop()
	iter.Stop()

	// Wait for background goroutine
	wg.Wait()

	// Verify closing flag is set
	require.True(t, iter.closing)
}

// TestCachingIterator_Flush_EmptyAndNil verifies that flush() handles
// nil and empty slices correctly.
func TestCachingIterator_Flush_EmptyAndNil(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	mockController := gomock.NewController(t)
	defer mockController.Finish()

	mockCache := mocks.NewMockInMemoryCache[any](mockController)

	sf := &singleflight.Group{}
	wg := &sync.WaitGroup{}

	t.Run("flush_with_nil_tuples_does_not_panic", func(t *testing.T) {
		iter := newCachingIterator(
			storage.NewStaticTupleIterator([]*openfgav1.Tuple{}),
			mockCache, testCacheKey("test-key"), 1000, time.Hour, 30*time.Second,
			sf, wg, "document", "viewer", "ReadUsersetTuples",
		)

		// Manually set tuples to nil to simulate abandoned state
		iter.tuples = nil

		// flush should not panic and should not call cache.Set
		// (no mockCache.EXPECT().Set() means test fails if Set is called)
		iter.flush()
	})

	t.Run("flush_with_empty_tuples_does_not_cache", func(t *testing.T) {
		iter := newCachingIterator(
			storage.NewStaticTupleIterator([]*openfgav1.Tuple{}),
			mockCache, testCacheKey("test-key"), 1000, time.Hour, 30*time.Second,
			sf, wg, "document", "viewer", "ReadUsersetTuples",
		)

		// tuples is initialized as empty slice
		require.NotNil(t, iter.tuples)
		require.Empty(t, iter.tuples)

		// flush should not panic and should not call cache.Set
		iter.flush()
	})
}

// TestCachingIterator_WaitGroup_AddInConstructor verifies that wg.Add is
// called in the constructor, not in Stop(), to prevent Add-after-Wait panic.
func TestCachingIterator_WaitGroup_AddInConstructor(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	mockController := gomock.NewController(t)
	defer mockController.Finish()

	mockCache := mocks.NewMockInMemoryCache[any](mockController)
	mockCache.EXPECT().Get(gomock.Any()).Return(nil).AnyTimes()
	mockCache.EXPECT().Set(gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes()

	sf := &singleflight.Group{}
	wg := &sync.WaitGroup{}

	tuples := []*openfgav1.Tuple{
		{Key: tuple.NewTupleKey("document:1", "viewer", "user:alice")},
	}

	// Create multiple iterators - wg.Add(1) is called in constructor for each
	// This verifies that even if we create iterators rapidly, the WaitGroup
	// is properly incremented at construction time (not at Stop time)
	iterators := make([]*CachingIterator, 10)
	for i := 0; i < 10; i++ {
		innerIter := storage.NewStaticTupleIterator(tuples)
		iterators[i] = newCachingIterator(
			innerIter, mockCache, testCacheKey("test-key-"+strconv.Itoa(i)), 1000, time.Hour, 30*time.Second,
			sf, wg, "document", "viewer", "ReadUsersetTuples",
		)
	}

	// Start a goroutine that will wait - this should not panic because
	// all Add() calls happened at construction time
	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()

	// Now stop all iterators - this calls wg.Done() for each
	for _, iter := range iterators {
		iter.Stop()
	}

	// Wait should complete without panic
	select {
	case <-done:
		// Success - all iterators properly decremented WaitGroup
	case <-time.After(5 * time.Second):
		t.Fatal("wg.Wait() timed out - wg.Done() was not called for all iterators")
	}
}

// TestCachingIterator_WaitGroup_DoneCalledOnNilTuples verifies that wg.Done()
// is called even when tuples is nil (early return path in Stop()).
func TestCachingIterator_WaitGroup_DoneCalledOnNilTuples(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	mockController := gomock.NewController(t)
	defer mockController.Finish()

	mockCache := mocks.NewMockInMemoryCache[any](mockController)

	sf := &singleflight.Group{}
	wg := &sync.WaitGroup{}

	// Create iterator with tuples that will exceed maxSize
	maxSize := 2
	tuples := make([]*openfgav1.Tuple, maxSize+1)
	for i := 0; i < maxSize+1; i++ {
		tuples[i] = &openfgav1.Tuple{Key: tuple.NewTupleKey("document:"+strconv.Itoa(i), "viewer", "user:test")}
	}

	innerIter := storage.NewStaticTupleIterator(tuples)

	iter := newCachingIterator(
		innerIter, mockCache, testCacheKey("test-key"), maxSize, time.Hour, 30*time.Second,
		sf, wg, "document", "viewer", "ReadUsersetTuples",
	)

	ctx := context.Background()

	// Consume all tuples - this will set tuples to nil due to maxSize exceeded
	for i := 0; i < maxSize+1; i++ {
		_, err := iter.Next(ctx)
		require.NoError(t, err)
	}

	// Verify tuples is nil (abandoned)
	require.Nil(t, iter.tuples)

	// Stop should call wg.Done() even though tuples is nil
	iter.Stop()

	// This should not hang - wg.Done() must have been called
	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()

	select {
	case <-done:
		// Success - wg.Wait() returned
	case <-time.After(1 * time.Second):
		t.Fatal("wg.Wait() timed out - wg.Done() was not called")
	}
}

// TestCachingIterator_ConcurrentNextAndStop verifies that concurrent Next()
// and Stop() calls on a CachingIterator do not race. This test should pass
// with the -race flag.
func TestCachingIterator_ConcurrentNextAndStop(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	mockController := gomock.NewController(t)
	defer mockController.Finish()

	mockCache := mocks.NewMockInMemoryCache[any](mockController)
	// Expect Get calls for optimization 1 (check if already cached)
	mockCache.EXPECT().Get(gomock.Any()).Return(nil).AnyTimes()
	mockCache.EXPECT().Set(gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes()

	tuples := createTestTuples(100)
	ctx := context.Background()
	sf := &singleflight.Group{}
	wg := &sync.WaitGroup{}

	// Run multiple iterations to increase chance of detecting race
	for i := 0; i < 100; i++ {
		innerIter := storage.NewStaticTupleIterator(tuples)
		iter := newCachingIterator(
			innerIter, mockCache, testCacheKey(fmt.Sprintf("test-key-%d", i)), 1000, time.Hour, 30*time.Second,
			sf, wg, "document", "viewer", "ReadUsersetTuples",
		)

		// Concurrent access pattern that previously caused race
		var testWg sync.WaitGroup
		testWg.Add(2)

		go func() {
			defer testWg.Done()
			for {
				_, err := iter.Next(ctx)
				if err != nil {
					return
				}
			}
		}()

		go func() {
			defer testWg.Done()
			time.Sleep(time.Microsecond)
			iter.Stop()
		}()

		testWg.Wait()
	}

	wg.Wait()
	// If the test got here with the -race flag without failing, no race condition detected
}

// ─────────────────────────────────────────────────────────────────────────────
// CachingIterator Background Drain Tests
// ─────────────────────────────────────────────────────────────────────────────

// TestCachingIterator_BackgroundDrainIgnoresRequestContextCancellation verifies
// that the background drain completes and caches even when the original request
// context is canceled. This is important because:
// 1. Consumer calls Stop() after finding a result early
// 2. Original request completes and cancels its context
// 3. Background drain should still complete using context.Background().
func TestCachingIterator_BackgroundDrainIgnoresRequestContextCancellation(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	mockController := gomock.NewController(t)
	defer mockController.Finish()

	mockCache := mocks.NewMockInMemoryCache[any](mockController)

	// Create a request context that we'll cancel to simulate request completion
	requestCtx, cancelRequest := context.WithCancel(context.Background())

	sf := &singleflight.Group{}
	wg := &sync.WaitGroup{}

	// Create tuples - use StaticTupleIterator since background drain uses context.Background()
	tuples := make([]*openfgav1.Tuple, 10)
	for i := 0; i < 10; i++ {
		tuples[i] = &openfgav1.Tuple{Key: tuple.NewTupleKey("document:"+string(rune('1'+i)), "viewer", "user:test")}
	}

	innerIter := storage.NewStaticTupleIterator(tuples)

	// Expect cache.Get to check if already cached (optimization 1), return nil (not cached)
	mockCache.EXPECT().Get(testCacheKey("test-key")).Return(nil).Times(1)
	// Expect cache.Set because background drain uses context.Background() and should complete
	mockCache.EXPECT().Set(gomock.Any(), gomock.Any(), gomock.Any()).Times(1)

	iter := newCachingIterator(
		innerIter, mockCache, testCacheKey("test-key"), 1000, time.Hour, 30*time.Second,
		sf, wg, "document", "viewer", "ReadUsersetTuples",
	)

	// Consume just a few tuples (simulating finding a result early)
	for i := 0; i < 3; i++ {
		_, err := iter.Next(requestCtx)
		require.NoError(t, err)
	}

	// Cancel the request context to simulate request completion
	cancelRequest()

	// Stop triggers background drain which should complete and cache
	// despite the request context being canceled
	iter.Stop()

	// Wait for background drain to complete
	wg.Wait()

	// Verify: cache.Set was called (checked by gomock)
}

func TestCachingIterator_BackgroundDrainCompletes_DoesCache(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	mockController := gomock.NewController(t)
	defer mockController.Finish()

	mockCache := mocks.NewMockInMemoryCache[any](mockController)

	ctx := context.Background()
	sf := &singleflight.Group{}
	wg := &sync.WaitGroup{}

	tuples := make([]*openfgav1.Tuple, 10)
	for i := 0; i < 10; i++ {
		tuples[i] = &openfgav1.Tuple{Key: tuple.NewTupleKey("document:"+string(rune('1'+i)), "viewer", "user:test")}
	}

	innerIter := storage.NewStaticTupleIterator(tuples)

	// Expect cache.Get to check if already cached (optimization 1), return nil (not cached)
	mockCache.EXPECT().Get(testCacheKey("test-key")).Return(nil).Times(1)
	// Expect cache.Set because drain will complete successfully
	mockCache.EXPECT().Set(gomock.Any(), gomock.Any(), gomock.Any()).Times(1)

	iter := newCachingIterator(
		innerIter, mockCache, testCacheKey("test-key"), 1000, time.Hour, 30*time.Second,
		sf, wg, "document", "viewer", "ReadUsersetTuples",
	)

	// Consume just a few tuples (not all)
	for i := 0; i < 3; i++ {
		_, err := iter.Next(ctx)
		require.NoError(t, err)
	}

	// Stop triggers background drain which should complete and cache
	iter.Stop()

	// Wait for background drain to complete
	wg.Wait()

	// Verify: cache.Set was called once (checked by gomock)
}

// TestCachingIterator_DrainTimeout_AbandonsCaching verifies that when the
// drain timeout expires mid-drain, tuples are abandoned.
func TestCachingIterator_DrainTimeout_AbandonsCaching(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	mockController := gomock.NewController(t)
	defer mockController.Finish()

	mockCache := mocks.NewMockInMemoryCache[any](mockController)
	// cache.Get returns nil (not cached) — needed for optimization 1 check
	mockCache.EXPECT().Get(testCacheKey("test-key")).Return(nil).Times(1)
	// No Set expected — drain should abandon due to timeout

	sf := &singleflight.Group{}
	wg := &sync.WaitGroup{}

	initialTuples := []*openfgav1.Tuple{
		{Key: tuple.NewTupleKey("document:1", "viewer", "user:alice")},
		{Key: tuple.NewTupleKey("document:2", "viewer", "user:bob")},
	}

	blockIter := &blockingIterator{
		tuples:  initialTuples,
		blockCh: make(chan struct{}),
	}

	iter := newCachingIterator(
		blockIter, mockCache, testCacheKey("test-key"), 1000,
		time.Hour,
		1*time.Nanosecond, // Very short drain timeout to trigger context expiry
		sf, wg, "document", "viewer", "ReadUsersetTuples",
	)

	ctx := context.Background()

	// Consume first tuple only
	_, err := iter.Next(ctx)
	require.NoError(t, err)

	// Stop triggers background drain. The drain timeout is 1ns so context
	// will expire almost immediately, causing abandonment.
	iter.Stop()

	wg.Wait()
	// If we get here without hanging, drain properly abandoned and called wg.Done
}

// TestCachingIterator_DrainError_AbandonsCaching verifies that a non-done
// error during background drain causes tuples to be abandoned.
func TestCachingIterator_DrainError_AbandonsCaching(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	mockController := gomock.NewController(t)
	defer mockController.Finish()

	mockCache := mocks.NewMockInMemoryCache[any](mockController)
	mockCache.EXPECT().Get(testCacheKey("test-key")).Return(nil).Times(1)
	// No Set expected — drain encounters error and abandons

	sf := &singleflight.Group{}
	wg := &sync.WaitGroup{}

	// errorIterator: first Next() returns tuple, second Next() returns ErrSimulatedError
	innerIter := mocks.NewErrorTupleIterator([]*openfgav1.Tuple{
		{Key: tuple.NewTupleKey("document:1", "viewer", "user:alice")},
		{Key: tuple.NewTupleKey("document:2", "viewer", "user:bob")},
	})

	iter := newCachingIterator(
		innerIter, mockCache, testCacheKey("test-key"), 1000, time.Hour, 30*time.Second,
		sf, wg, "document", "viewer", "ReadUsersetTuples",
	)

	ctx := context.Background()

	// Consume first tuple
	_, err := iter.Next(ctx)
	require.NoError(t, err)

	// Stop triggers background drain. The drain calls inner.Next() which
	// returns ErrSimulatedError (not done/cancelled), causing abandonment.
	iter.Stop()

	wg.Wait()
}

// TestCachingIterator_DrainExceedsMaxSize_AbandonsCaching verifies that when
// the total number of tuples collected during drain exceeds maxSize, caching
// is abandoned.
func TestCachingIterator_DrainExceedsMaxSize_AbandonsCaching(t *testing.T) {
	t.Cleanup(func() {
		goleak.VerifyNone(t)
	})

	mockController := gomock.NewController(t)
	defer mockController.Finish()

	mockCache := mocks.NewMockInMemoryCache[any](mockController)
	mockCache.EXPECT().Get(testCacheKey("test-key")).Return(nil).Times(1)
	// No Set expected — max size exceeded during drain

	sf := &singleflight.Group{}
	wg := &sync.WaitGroup{}

	maxSize := 5
	tuples := make([]*openfgav1.Tuple, 10)
	for i := 0; i < 10; i++ {
		tuples[i] = &openfgav1.Tuple{Key: tuple.NewTupleKey("document:"+strconv.Itoa(i), "viewer", "user:test")}
	}

	innerIter := storage.NewStaticTupleIterator(tuples)

	iter := newCachingIterator(
		innerIter, mockCache, testCacheKey("test-key"), maxSize, time.Hour, 30*time.Second,
		sf, wg, "document", "viewer", "ReadUsersetTuples",
	)

	ctx := context.Background()

	// Consume 3 tuples (under maxSize, so tuples is still collecting)
	for i := 0; i < 3; i++ {
		_, err := iter.Next(ctx)
		require.NoError(t, err)
	}
	require.NotNil(t, iter.tuples)
	require.Len(t, iter.tuples, 3)

	// Stop triggers drain. Drain collects remaining 7 tuples.
	// After 3 more, total=6 > maxSize=5, causing abandonment.
	iter.Stop()

	wg.Wait()
}

// ─────────────────────────────────────────────────────────────────────────────
// extractObjectID Tests
// ─────────────────────────────────────────────────────────────────────────────

func TestExtractObjectID(t *testing.T) {
	tests := []struct {
		input    string
		expected string
	}{
		{"document:123", "123"},
		{"folder:abc-def", "abc-def"},
		{"type:", ""},
		{"nocolon", "nocolon"},
		{"a:b:c", "b:c"},
	}

	for _, tt := range tests {
		t.Run(tt.input, func(t *testing.T) {
			result := extractObjectID(tt.input)
			require.Equal(t, tt.expected, result)
		})
	}
}

// ─────────────────────────────────────────────────────────────────────────────
// Benchmarks
// ─────────────────────────────────────────────────────────────────────────────

// BenchmarkCachingIterator_CacheMiss benchmarks cache miss scenario:
// wrapping a database iterator and collecting tuples for caching.
func BenchmarkCachingIterator_CacheMiss(b *testing.B) {
	mockController := gomock.NewController(b)
	defer mockController.Finish()

	mockCache := mocks.NewMockInMemoryCache[any](mockController)
	mockCache.EXPECT().Set(gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes()

	ctx := context.Background()
	sf := &singleflight.Group{}
	wg := &sync.WaitGroup{}

	// Create 100 tuples for iteration
	tuples := make([]*openfgav1.Tuple, 100)
	for i := 0; i < 100; i++ {
		tuples[i] = &openfgav1.Tuple{
			Key: tuple.NewTupleKey("document:"+string(rune('a'+i%26)), "viewer", "user:test"),
		}
	}

	b.ResetTimer()
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		innerIter := storage.NewStaticTupleIterator(tuples)
		iter := newCachingIterator(
			innerIter, mockCache, testCacheKey("test-key"), 1000, time.Hour, 30*time.Second,
			sf, wg, "document", "viewer", "benchmark",
		)

		// Consume all tuples
		for {
			_, err := iter.Next(ctx)
			if err != nil {
				break
			}
		}
	}
}

// BenchmarkCachingIterator_CacheHit benchmarks cache hit scenario:
// creating a LockFreeCachedIterator and iterating over cached entries.
func BenchmarkCachingIterator_CacheHit(b *testing.B) {
	// Create 100 cached entries
	entries := make([]MinimalCacheEntry, 100)
	for i := 0; i < 100; i++ {
		entries[i] = MinimalCacheEntry{
			ObjectID: string(rune('a' + i%26)),
			User:     "user:test",
		}
	}

	ctx := context.Background()

	b.ResetTimer()
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		iter := NewLockFreeCachedIterator(entries, "document", "viewer", false)

		// Consume all entries
		for {
			_, err := iter.Next(ctx)
			if err != nil {
				break
			}
		}
	}
}

// BenchmarkLockFreeCachedIterator_Next benchmarks the Next() operation
// of the lock-free cached iterator.
func BenchmarkLockFreeCachedIterator_Next(b *testing.B) {
	// Create 100 cached entries
	entries := make([]MinimalCacheEntry, 100)
	for i := 0; i < 100; i++ {
		entries[i] = MinimalCacheEntry{
			ObjectID:      string(rune('a' + i%26)),
			User:          "user:test",
			ConditionName: "test_condition",
		}
	}

	ctx := context.Background()

	b.ResetTimer()
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		iter := NewLockFreeCachedIterator(entries, "document", "viewer", false)

		// Consume all entries
		for {
			_, err := iter.Next(ctx)
			if err != nil {
				break
			}
		}
	}
}

// BenchmarkLockFreeCachedIterator_VsStaticIterator compares the lock-free
// cached iterator against the standard StaticTupleIterator.
func BenchmarkLockFreeCachedIterator_VsStaticIterator(b *testing.B) {
	// Create data for both iterators
	entries := make([]MinimalCacheEntry, 100)
	tuples := make([]*openfgav1.Tuple, 100)
	for i := 0; i < 100; i++ {
		entries[i] = MinimalCacheEntry{
			ObjectID: string(rune('a' + i%26)),
			User:     "user:test",
		}
		tuples[i] = &openfgav1.Tuple{
			Key: tuple.NewTupleKey("document:"+string(rune('a'+i%26)), "viewer", "user:test"),
		}
	}

	ctx := context.Background()

	b.Run("LockFreeCachedIterator", func(b *testing.B) {
		b.ReportAllocs()
		for i := 0; i < b.N; i++ {
			iter := NewLockFreeCachedIterator(entries, "document", "viewer", false)
			for {
				_, err := iter.Next(ctx)
				if err != nil {
					break
				}
			}
		}
	})

	b.Run("StaticTupleIterator", func(b *testing.B) {
		b.ReportAllocs()
		for i := 0; i < b.N; i++ {
			iter := storage.NewStaticTupleIterator(tuples)
			for {
				_, err := iter.Next(ctx)
				if err != nil {
					break
				}
			}
		}
	})
}

// BenchmarkMinimalCacheEntry_Memory benchmarks memory usage of MinimalCacheEntry
// vs full TupleRecord storage.
func BenchmarkMinimalCacheEntry_Memory(b *testing.B) {
	b.Run("MinimalCacheEntry_100", func(b *testing.B) {
		b.ReportAllocs()
		for i := 0; i < b.N; i++ {
			entries := make([]MinimalCacheEntry, 100)
			for j := 0; j < 100; j++ {
				entries[j] = MinimalCacheEntry{
					ObjectID: "object-" + string(rune('a'+j%26)),
					User:     "user:test-user",
				}
			}
			_ = entries
		}
	})

	b.Run("FullTuple_100", func(b *testing.B) {
		b.ReportAllocs()
		for i := 0; i < b.N; i++ {
			tuples := make([]*openfgav1.Tuple, 100)
			for j := 0; j < 100; j++ {
				tuples[j] = &openfgav1.Tuple{
					Key: tuple.NewTupleKey(
						"document:object-"+string(rune('a'+j%26)),
						"viewer",
						"user:test-user",
					),
				}
			}
			_ = tuples
		}
	})
}

// BenchmarkCacheKeyGeneration benchmarks cache key generation performance.
func BenchmarkCacheKeyGeneration(b *testing.B) {
	b.Run("ReadUsersetTuples", func(b *testing.B) {
		filter := storage.ReadUsersetTuplesFilter{
			Object:   "document:1",
			Relation: "viewer",
			AllowedUserTypeRestrictions: []*openfgav1.RelationReference{
				{Type: "group", RelationOrWildcard: &openfgav1.RelationReference_Relation{Relation: "member"}},
				{Type: "user"},
			},
			Conditions: []string{"cond1", "cond2"},
		}

		b.ReportAllocs()
		b.ResetTimer()

		for i := 0; i < b.N; i++ {
			_ = storage.ReadUsersetTuplesKey("store123", filter)
		}
	})

	b.Run("Read", func(b *testing.B) {
		filter := storage.ReadFilter{
			Object:     "document:1",
			Relation:   "parent",
			User:       "folder:",
			Conditions: []string{"cond1"},
		}

		b.ReportAllocs()
		b.ResetTimer()

		for i := 0; i < b.N; i++ {
			_ = storage.ReadKey("store123", filter)
		}
	})

	b.Run("ReadStartingWithUser", func(b *testing.B) {
		filter := storage.ReadStartingWithUserFilter{
			ObjectType: "document",
			Relation:   "viewer",
			UserFilter: []*openfgav1.ObjectRelation{
				{Object: "user:alice"},
				{Object: "user:bob"},
			},
			Conditions: []string{"cond1", "cond2"},
		}

		b.ReportAllocs()
		b.ResetTimer()

		for i := 0; i < b.N; i++ {
			_ = storage.ReadStartingWithUserKey("store123", filter)
		}
	})
}
