package memorysearch

import (
	"context"
	"errors"
	"fmt"
	"path/filepath"
	"sync"
)

type poolState uint8

const (
	poolActive poolState = iota
	poolRetiring
	poolClosed
)

type generationPool struct {
	basename string
	backing  *Generation

	mu       sync.Mutex
	state    poolState
	refs     int
	closeErr error
	closedCh chan struct{}
}

func newGenerationPool(basename string, backing *Generation) *generationPool {
	return &generationPool{
		basename: basename,
		backing:  backing,
		state:    poolActive,
		closedCh: make(chan struct{}),
	}
}

func (p *generationPool) borrow() (*Generation, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.state != poolActive {
		return nil, p.stateErrorLocked()
	}
	p.refs++
	return &Generation{
		dir:      p.backing.dir,
		manifest: cloneManifest(p.backing.manifest),
		index:    p.backing.index,
		release:  p.release,
	}, nil
}

func (p *generationPool) release() error {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.refs <= 0 {
		return fmt.Errorf("%w: pool %s release with refs=%d",
			ErrPoolRefInvariant, p.basename, p.refs)
	}
	p.refs--
	if p.refs == 0 && p.state == poolRetiring {
		return p.closeLocked()
	}
	return nil
}

func (p *generationPool) markRetiring() {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.state != poolActive {
		return
	}
	p.state = poolRetiring
	if p.refs == 0 {
		_ = p.closeLocked()
	}
}

func (p *generationPool) waitClosed() error {
	p.mu.Lock()
	if p.state == poolClosed {
		err := p.closeErr
		p.mu.Unlock()
		return err
	}
	ch := p.closedCh
	p.mu.Unlock()
	<-ch
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.closeErr
}

// closeUnborrowed is only for a freshly opened handle which never became
// visible to a borrower (for example, registry shutdown won the insert gate).
func (p *generationPool) closeUnborrowed() error {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.refs != 0 {
		return fmt.Errorf("%w: pool %s has %d borrowers",
			ErrPoolRefInvariant, p.basename, p.refs)
	}
	return p.closeLocked()
}

func (p *generationPool) closeLocked() error {
	if p.state == poolClosed {
		return p.closeErr
	}
	if p.refs != 0 {
		return fmt.Errorf("%w: closing pool %s with %d borrowers",
			ErrPoolRefInvariant, p.basename, p.refs)
	}
	p.state = poolClosed
	p.closeErr = p.backing.Close()
	close(p.closedCh)
	return p.closeErr
}

func (p *generationPool) stateSnapshot() poolState {
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.state
}

func (p *generationPool) stateErrorLocked() error {
	if p.state == poolClosed {
		return fmt.Errorf("%w: %s", ErrPoolClosed, p.basename)
	}
	return fmt.Errorf("%w: %s", ErrPoolRetiring, p.basename)
}

type openGenerationFunc func(
	context.Context,
	*Binding,
	AgentWorkDirResolver,
	string,
) (*Generation, error)

// openingEntry remains registered until every Current call which selected
// this basename before/during the open has consumed or cancelled its wait.
// retirePending fences later borrows while preserving those already selected
// by a current.json read before publish/GC moved past the generation.
type openingEntry struct {
	done          chan struct{}
	pool          *generationPool
	err           error
	waiters       int
	completed     bool
	retirePending bool
}

type poolRegistry struct {
	binding  *Binding
	resolver AgentWorkDirResolver
	open     openGenerationFunc

	openCtx    context.Context
	openCancel context.CancelFunc

	mu         sync.Mutex
	closed     bool
	pools      map[string]*generationPool
	opening    map[string]*openingEntry
	openWG     sync.WaitGroup
	closeStart sync.Once

	closeDone chan struct{}
	closeErr  error
}

func newPoolRegistry(binding *Binding, resolver AgentWorkDirResolver) *poolRegistry {
	return newPoolRegistryWithOpener(binding, resolver, OpenGeneration)
}

func newPoolRegistryWithOpener(
	binding *Binding,
	resolver AgentWorkDirResolver,
	open openGenerationFunc,
) *poolRegistry {
	ctx, cancel := context.WithCancel(context.Background())
	bindingCopy := *binding
	return &poolRegistry{
		binding:    &bindingCopy,
		resolver:   resolver,
		open:       open,
		openCtx:    ctx,
		openCancel: cancel,
		pools:      make(map[string]*generationPool),
		opening:    make(map[string]*openingEntry),
		closeDone:  make(chan struct{}),
	}
}

func (r *poolRegistry) BorrowFor(ctx context.Context, genDir string) (*Generation, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	basename := filepath.Base(genDir)
	if err := validateBasename(basename); err != nil || basename == "" {
		return nil, fmt.Errorf("memory search: pool basename %q: %w",
			basename, ErrInvalidBasename)
	}

	r.mu.Lock()
	if r.closed {
		r.mu.Unlock()
		return nil, ErrPoolRegistryClosed
	}
	if entry := r.opening[basename]; entry != nil {
		entry.waiters++
		r.mu.Unlock()
		return r.waitForOpening(ctx, basename, entry)
	}
	if pool := r.pools[basename]; pool != nil {
		gen, err := pool.borrow()
		r.mu.Unlock()
		return gen, err
	}

	entry := &openingEntry{done: make(chan struct{}), waiters: 1}
	r.opening[basename] = entry
	r.openWG.Add(1)
	r.mu.Unlock()

	go r.runOpener(basename, genDir, entry)
	return r.waitForOpening(ctx, basename, entry)
}

func (r *poolRegistry) waitForOpening(
	ctx context.Context,
	basename string,
	entry *openingEntry,
) (*Generation, error) {
	select {
	case <-entry.done:
	case <-ctx.Done():
		r.mu.Lock()
		r.consumeOpeningWaiterLocked(basename, entry)
		r.mu.Unlock()
		return nil, ctx.Err()
	}

	r.mu.Lock()
	defer r.mu.Unlock()
	defer r.consumeOpeningWaiterLocked(basename, entry)
	if r.closed {
		return nil, ErrPoolRegistryClosed
	}
	if entry.err != nil {
		return nil, entry.err
	}
	if r.pools[basename] != entry.pool || entry.pool == nil {
		return nil, fmt.Errorf("%w: %s", ErrPoolRetired, basename)
	}
	// Existing opening waiters selected this basename before its retirement
	// fence. Keep the pool active until the final waiter has borrowed.
	return entry.pool.borrow()
}

func (r *poolRegistry) runOpener(
	basename string,
	genDir string,
	entry *openingEntry,
) {
	defer r.openWG.Done()
	gen, err := r.open(r.openCtx, r.binding, r.resolver, genDir)

	r.mu.Lock()
	defer r.mu.Unlock()
	entry.completed = true
	if err != nil {
		entry.err = err
		close(entry.done)
		r.finalizeOpeningLocked(basename, entry)
		return
	}
	pool := newGenerationPool(basename, gen)
	if r.closed {
		entry.err = ErrPoolRegistryClosed
		_ = pool.closeUnborrowed()
		close(entry.done)
		r.finalizeOpeningLocked(basename, entry)
		return
	}
	if entry.waiters == 0 {
		// Every caller which selected this generation cancelled while the
		// shared open was in flight. Do not retain an unborrowed reader
		// indefinitely; a later Current may reopen it.
		_ = pool.closeUnborrowed()
		close(entry.done)
		r.finalizeOpeningLocked(basename, entry)
		return
	}
	r.pools[basename] = pool
	entry.pool = pool
	close(entry.done)
	r.finalizeOpeningLocked(basename, entry)
}

func (r *poolRegistry) consumeOpeningWaiterLocked(
	basename string,
	entry *openingEntry,
) {
	if entry.waiters <= 0 {
		return
	}
	entry.waiters--
	r.finalizeOpeningLocked(basename, entry)
}

func (r *poolRegistry) finalizeOpeningLocked(
	basename string,
	entry *openingEntry,
) {
	if !entry.completed || entry.waiters != 0 {
		return
	}
	if r.opening[basename] == entry {
		delete(r.opening, basename)
	}
	if entry.retirePending && entry.pool != nil {
		entry.pool.markRetiring()
	}
}

func (r *poolRegistry) MarkRetire(basename string) bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	if entry := r.opening[basename]; entry != nil {
		entry.retirePending = true
		r.finalizeOpeningLocked(basename, entry)
		return true
	}
	if pool := r.pools[basename]; pool != nil {
		pool.markRetiring()
		return true
	}
	return false
}

// RetireForGC atomically marks an in-flight opener or live pool for
// retirement and returns the resulting state. The mutating name is
// intentional: GC must fence the handle before deciding whether deletion is
// safe.
func (r *poolRegistry) RetireForGC(
	basename string,
) (pooled bool, state poolState, opening bool) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if entry := r.opening[basename]; entry != nil {
		entry.retirePending = true
		r.finalizeOpeningLocked(basename, entry)
		return entry.pool != nil, poolRetiring, true
	}
	if pool := r.pools[basename]; pool != nil {
		pool.markRetiring()
		return true, pool.stateSnapshot(), false
	}
	return false, poolActive, false
}

func (r *poolRegistry) WaitUntilClosed(ctx context.Context, basename string) error {
	r.mu.Lock()
	pool := r.pools[basename]
	r.mu.Unlock()
	if pool == nil {
		return nil
	}
	pool.mu.Lock()
	if pool.state == poolClosed {
		err := pool.closeErr
		pool.mu.Unlock()
		return err
	}
	closedCh := pool.closedCh
	pool.mu.Unlock()
	select {
	case <-closedCh:
		return pool.waitClosed()
	case <-ctx.Done():
		return ctx.Err()
	}
}

func (r *poolRegistry) Forget(basename string) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.opening[basename] != nil {
		return fmt.Errorf("%w: forgetting opening pool %s",
			ErrPoolRefInvariant, basename)
	}
	pool := r.pools[basename]
	if pool == nil {
		return nil
	}
	if pool.stateSnapshot() != poolClosed {
		return fmt.Errorf("%w: forgetting live pool %s",
			ErrPoolRefInvariant, basename)
	}
	delete(r.pools, basename)
	return nil
}

func (r *poolRegistry) Close() error {
	r.closeStart.Do(func() {
		r.mu.Lock()
		r.closed = true
		r.openCancel()
		r.mu.Unlock()

		go r.closeSlow()
	})
	<-r.closeDone
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.closeErr
}

func (r *poolRegistry) closeSlow() {
	// bleve.Open does not reliably interrupt on context cancellation once
	// the underlying store open has begun. openWG may therefore wait for one
	// bounded in-flight open; singleflight and owner-lock serialization keep
	// it from contending with another opener for the same generation.
	r.openWG.Wait()

	r.mu.Lock()
	pools := make([]*generationPool, 0, len(r.pools))
	for _, pool := range r.pools {
		pools = append(pools, pool)
	}
	r.mu.Unlock()

	for _, pool := range pools {
		pool.markRetiring()
	}
	var errs []error
	for _, pool := range pools {
		if err := pool.waitClosed(); err != nil {
			errs = append(errs, fmt.Errorf("pool %s: %w", pool.basename, err))
		}
	}
	r.mu.Lock()
	r.closeErr = errors.Join(errs...)
	close(r.closeDone)
	r.mu.Unlock()
}

func cloneManifest(manifest *Manifest) *Manifest {
	if manifest == nil {
		return nil
	}
	clone := *manifest
	if manifest.Files != nil {
		clone.Files = make(map[string]FileEntry, len(manifest.Files))
		for path, file := range manifest.Files {
			file.ChunkIDs = append([]string(nil), file.ChunkIDs...)
			clone.Files[path] = file
		}
	}
	if manifest.Chunks != nil {
		clone.Chunks = make(map[string]ChunkEntry, len(manifest.Chunks))
		for id, chunk := range manifest.Chunks {
			clone.Chunks[id] = chunk
		}
	}
	return &clone
}

var (
	ErrPoolRetiring       = errors.New("memory search: generation pool retiring")
	ErrPoolClosed         = errors.New("memory search: generation pool closed")
	ErrPoolRetired        = errors.New("memory search: generation pool retired")
	ErrPoolRefInvariant   = errors.New("memory search: generation pool ref invariant")
	ErrPoolRegistryClosed = errors.New("memory search: generation pool registry closed")
)