//go:build !vectors

package memorysearch

import (
	"context"
	"errors"
	"path/filepath"
	"sync"
	"sync/atomic"
	"testing"
	"time"

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

type fakePoolIndex struct {
	closeCalls atomic.Int64
	closeErr   error
}

func (f *fakePoolIndex) Close() error {
	f.closeCalls.Add(1)
	return f.closeErr
}

func testPool(t *testing.T) (*generationPool, *fakePoolIndex) {
	t.Helper()
	index := &fakePoolIndex{}
	backing := &Generation{
		dir: filepath.Join(t.TempDir(), "generations", "gen-a"),
		manifest: &Manifest{
			BindingFingerprint: "fp",
			Files: map[string]FileEntry{
				"system/identity.md": {ChunkIDs: []string{"chunk-a"}},
			},
			Chunks: map[string]ChunkEntry{
				"chunk-a": {ID: "chunk-a", Path: "system/identity.md"},
			},
		},
		closeIndex: index.Close,
	}
	return newGenerationPool("gen-a", backing), index
}

func TestGenerationPoolBorrowManifestAndCloseIsolation(t *testing.T) {
	pool, index := testPool(t)
	first, err := pool.borrow()
	require.NoError(t, err)
	second, err := pool.borrow()
	require.NoError(t, err)

	entry := first.Manifest().Files["system/identity.md"]
	entry.ChunkIDs[0] = "mutated"
	first.Manifest().Files["system/identity.md"] = entry
	first.Manifest().Chunks["extra"] = ChunkEntry{ID: "extra"}

	require.Equal(t, []string{"chunk-a"},
		second.Manifest().Files["system/identity.md"].ChunkIDs)
	require.NotContains(t, second.Manifest().Chunks, "extra")
	require.Equal(t, []string{"chunk-a"},
		pool.backing.Manifest().Files["system/identity.md"].ChunkIDs)

	require.NoError(t, first.Close())
	require.NoError(t, first.Close())
	require.Equal(t, int64(0), index.closeCalls.Load())
	require.NoError(t, second.Close())
	require.Equal(t, int64(0), index.closeCalls.Load())
}

func TestGenerationPoolRetireWaitsForLastBorrower(t *testing.T) {
	pool, index := testPool(t)
	first, err := pool.borrow()
	require.NoError(t, err)
	second, err := pool.borrow()
	require.NoError(t, err)

	pool.markRetiring()
	_, err = pool.borrow()
	require.ErrorIs(t, err, ErrPoolRetiring)

	done := make(chan error, 1)
	go func() { done <- pool.waitClosed() }()
	require.Never(t, func() bool {
		select {
		case <-done:
			return true
		default:
			return false
		}
	}, 100*time.Millisecond, 10*time.Millisecond)

	require.NoError(t, first.Close())
	require.Equal(t, int64(0), index.closeCalls.Load())
	require.NoError(t, second.Close())
	require.NoError(t, <-done)
	require.Equal(t, int64(1), index.closeCalls.Load())
	require.Equal(t, poolClosed, pool.stateSnapshot())
}

func TestGenerationPoolCloseErrorIsCached(t *testing.T) {
	pool, index := testPool(t)
	index.closeErr = errors.New("close failed")
	gen, err := pool.borrow()
	require.NoError(t, err)
	pool.markRetiring()

	err = gen.Close()
	require.ErrorIs(t, err, index.closeErr)
	require.ErrorIs(t, gen.Close(), index.closeErr)
	require.ErrorIs(t, pool.waitClosed(), index.closeErr)
	require.Equal(t, int64(1), index.closeCalls.Load())
}

func TestGenerationDirectCloseConcurrentIdempotent(t *testing.T) {
	index := &fakePoolIndex{closeErr: errors.New("close failed")}
	gen := &Generation{closeIndex: index.Close}
	const callers = 16
	errs := make(chan error, callers)
	var wg sync.WaitGroup
	for range callers {
		wg.Add(1)
		go func() {
			defer wg.Done()
			errs <- gen.Close()
		}()
	}
	wg.Wait()
	close(errs)
	for err := range errs {
		require.ErrorIs(t, err, index.closeErr)
	}
	require.Equal(t, int64(1), index.closeCalls.Load())
}

func TestPoolRegistrySingleflightAndWaiterCancellation(t *testing.T) {
	h := newBuildTestHarness(t, map[string]string{
		"system/identity.md": "# Identity\n\nbody\n",
	})
	genDir := filepath.Join(h.derivativeRoot, generationsDirName, "gen-a")
	started := make(chan struct{})
	release := make(chan struct{})
	var calls atomic.Int64
	var opened []*fakePoolIndex
	var openedMu sync.Mutex
	opener := func(
		ctx context.Context,
		_ *Binding,
		_ AgentWorkDirResolver,
		dir string,
	) (*Generation, error) {
		calls.Add(1)
		select {
		case <-started:
		default:
			close(started)
		}
		select {
		case <-release:
		case <-ctx.Done():
			return nil, ctx.Err()
		}
		index := &fakePoolIndex{}
		openedMu.Lock()
		opened = append(opened, index)
		openedMu.Unlock()
		return &Generation{
			dir:        dir,
			manifest:   &Manifest{Files: map[string]FileEntry{}, Chunks: map[string]ChunkEntry{}},
			closeIndex: index.Close,
		}, nil
	}
	registry := newPoolRegistryWithOpener(h.binding, h.resolver, opener)
	t.Cleanup(func() { _ = registry.Close() })

	const waiters = 8
	results := make(chan *Generation, waiters)
	errs := make(chan error, waiters)
	for range waiters {
		go func() {
			gen, err := registry.BorrowFor(context.Background(), genDir)
			results <- gen
			errs <- err
		}()
	}
	<-started

	cancelCtx, cancel := context.WithCancel(context.Background())
	cancelDone := make(chan error, 1)
	go func() {
		_, err := registry.BorrowFor(cancelCtx, genDir)
		cancelDone <- err
	}()
	cancel()
	require.ErrorIs(t, <-cancelDone, context.Canceled)
	require.Equal(t, int64(1), calls.Load())

	close(release)
	var gens []*Generation
	for range waiters {
		require.NoError(t, <-errs)
		gens = append(gens, <-results)
	}
	require.Equal(t, int64(1), calls.Load())
	for _, gen := range gens {
		require.NoError(t, gen.Close())
	}
}

func TestPoolRegistryAllWaitersCancelClosesUnusedOpen(t *testing.T) {
	h := newBuildTestHarness(t, map[string]string{
		"system/identity.md": "# Identity\n\nbody\n",
	})
	genDir := filepath.Join(h.derivativeRoot, generationsDirName, "gen-a")
	started := make(chan struct{})
	release := make(chan struct{})
	index := &fakePoolIndex{}
	registry := newPoolRegistryWithOpener(
		h.binding,
		h.resolver,
		func(ctx context.Context, _ *Binding, _ AgentWorkDirResolver, dir string) (*Generation, error) {
			close(started)
			select {
			case <-release:
			case <-ctx.Done():
				return nil, ctx.Err()
			}
			return &Generation{
				dir:        dir,
				manifest:   &Manifest{Files: map[string]FileEntry{}, Chunks: map[string]ChunkEntry{}},
				closeIndex: index.Close,
			}, nil
		},
	)
	t.Cleanup(func() { _ = registry.Close() })

	ctx, cancel := context.WithCancel(context.Background())
	result := make(chan error, 1)
	go func() {
		_, err := registry.BorrowFor(ctx, genDir)
		result <- err
	}()
	<-started
	cancel()
	require.ErrorIs(t, <-result, context.Canceled)
	close(release)
	require.Eventually(t, func() bool {
		return index.closeCalls.Load() == 1
	}, time.Second, 10*time.Millisecond)
	pooled, _, opening := registry.RetireForGC("gen-a")
	require.False(t, pooled)
	require.False(t, opening)
}

func TestPoolRegistryCloseFencesOpenAndDrainsBorrow(t *testing.T) {
	h := newBuildTestHarness(t, map[string]string{
		"system/identity.md": "# Identity\n\nbody\n",
	})
	genDir := filepath.Join(h.derivativeRoot, generationsDirName, "gen-a")
	index := &fakePoolIndex{}
	registry := newPoolRegistryWithOpener(
		h.binding,
		h.resolver,
		func(ctx context.Context, _ *Binding, _ AgentWorkDirResolver, dir string) (*Generation, error) {
			return &Generation{
				dir:        dir,
				manifest:   &Manifest{Files: map[string]FileEntry{}, Chunks: map[string]ChunkEntry{}},
				closeIndex: index.Close,
			}, nil
		},
	)
	gen, err := registry.BorrowFor(context.Background(), genDir)
	require.NoError(t, err)

	closed := make(chan error, 1)
	go func() { closed <- registry.Close() }()
	require.Never(t, func() bool {
		select {
		case <-closed:
			return true
		default:
			return false
		}
	}, 100*time.Millisecond, 10*time.Millisecond)

	_, err = registry.BorrowFor(context.Background(), genDir)
	require.ErrorIs(t, err, ErrPoolRegistryClosed)
	require.NoError(t, gen.Close())
	require.NoError(t, <-closed)
	require.Equal(t, int64(1), index.closeCalls.Load())
	require.NoError(t, registry.Close())
}

func TestPoolRegistryConcurrentCloseWaitsForSameDrain(t *testing.T) {
	h := newBuildTestHarness(t, map[string]string{
		"system/identity.md": "# Identity\n\nbody\n",
	})
	genDir := filepath.Join(h.derivativeRoot, generationsDirName, "gen-a")
	index := &fakePoolIndex{closeErr: errors.New("close failed")}
	registry := newPoolRegistryWithOpener(
		h.binding,
		h.resolver,
		func(ctx context.Context, _ *Binding, _ AgentWorkDirResolver, dir string) (*Generation, error) {
			return &Generation{
				dir:        dir,
				manifest:   &Manifest{Files: map[string]FileEntry{}, Chunks: map[string]ChunkEntry{}},
				closeIndex: index.Close,
			}, nil
		},
	)
	gen, err := registry.BorrowFor(context.Background(), genDir)
	require.NoError(t, err)

	const closers = 8
	results := make(chan error, closers)
	for range closers {
		go func() { results <- registry.Close() }()
	}
	require.Never(t, func() bool {
		return len(results) != 0
	}, 100*time.Millisecond, 10*time.Millisecond)
	require.ErrorIs(t, gen.Close(), index.closeErr)
	for range closers {
		require.ErrorIs(t, <-results, index.closeErr)
	}
	require.Equal(t, int64(1), index.closeCalls.Load())
}

func TestPoolRegistryRetireDuringOpenPreservesSelectedWaiter(t *testing.T) {
	h := newBuildTestHarness(t, map[string]string{
		"system/identity.md": "# Identity\n\nbody\n",
	})
	genDir := filepath.Join(h.derivativeRoot, generationsDirName, "gen-a")
	started := make(chan struct{})
	release := make(chan struct{})
	index := &fakePoolIndex{}
	registry := newPoolRegistryWithOpener(
		h.binding,
		h.resolver,
		func(ctx context.Context, _ *Binding, _ AgentWorkDirResolver, dir string) (*Generation, error) {
			close(started)
			select {
			case <-release:
			case <-ctx.Done():
				return nil, ctx.Err()
			}
			return &Generation{
				dir:        dir,
				manifest:   &Manifest{Files: map[string]FileEntry{}, Chunks: map[string]ChunkEntry{}},
				closeIndex: index.Close,
			}, nil
		},
	)
	t.Cleanup(func() { _ = registry.Close() })

	result := make(chan *Generation, 1)
	errResult := make(chan error, 1)
	go func() {
		gen, err := registry.BorrowFor(context.Background(), genDir)
		result <- gen
		errResult <- err
	}()
	<-started
	pooled, _, opening := registry.RetireForGC("gen-a")
	require.False(t, pooled)
	require.True(t, opening)
	close(release)

	require.NoError(t, <-errResult)
	gen := <-result
	require.NotNil(t, gen)
	_, err := registry.BorrowFor(context.Background(), genDir)
	require.ErrorIs(t, err, ErrPoolRetiring)
	require.NoError(t, gen.Close())
	require.Eventually(t, func() bool {
		pooled, state, opening := registry.RetireForGC("gen-a")
		return pooled && !opening && state == poolClosed
	}, time.Second, 10*time.Millisecond)
}