package dispatcher

import (
	"context"
	"errors"
	"sync"
)

var errMemoryWatcherMutationAdmissionClosed = errors.New("dispatcher memory watcher mutation admission is closed")

// memoryWatcherMutationOwner spans one complete watcher-producing operation:
// authoritative access lookup, clone/rebind, watcher construction and start,
// and either publication on the session or synchronous rejection cleanup.
//
// generation identifies the single Dispatcher lifecycle epoch. Shutdown closes
// admission and advances that generation while holding the same mutex used by
// publication, so an admitted owner cannot publish after the shutdown fence.
type memoryWatcherMutationOwner struct {
	generation uint64
	done       chan struct{}
	doneOnce   sync.Once
}

func (d *Dispatcher) beginMemoryWatcherMutation() (*memoryWatcherMutationOwner, error) {
	if d == nil {
		return nil, errMemoryWatcherMutationAdmissionClosed
	}
	d.memoryWatcherMutationMu.Lock()
	defer d.memoryWatcherMutationMu.Unlock()
	if d.memoryWatcherMutationClosed {
		return nil, errMemoryWatcherMutationAdmissionClosed
	}
	if d.memoryWatcherMutationGeneration == 0 {
		d.memoryWatcherMutationGeneration = 1
	}
	if d.memoryWatcherMutationOwners == nil {
		d.memoryWatcherMutationOwners = make(map[*memoryWatcherMutationOwner]struct{})
	}
	owner := &memoryWatcherMutationOwner{
		generation: d.memoryWatcherMutationGeneration,
		done:       make(chan struct{}),
	}
	d.memoryWatcherMutationOwners[owner] = struct{}{}
	return owner, nil
}

func (d *Dispatcher) finishMemoryWatcherMutation(owner *memoryWatcherMutationOwner) {
	if owner == nil {
		return
	}
	if d != nil {
		d.memoryWatcherMutationMu.Lock()
		delete(d.memoryWatcherMutationOwners, owner)
		d.memoryWatcherMutationMu.Unlock()
	}
	owner.doneOnce.Do(func() { close(owner.done) })
}

// closeMemoryWatcherMutationAdmission is the first Dispatcher shutdown fence.
// It permanently rejects new attach/refresh/rebind owners and returns a stable
// snapshot of all already-admitted owners. Publication validates its owner while
// holding this same mutex, so once this method returns no new watcher can become
// visible on a session.
func (d *Dispatcher) closeMemoryWatcherMutationAdmission() []<-chan struct{} {
	if d == nil {
		return nil
	}
	d.memoryWatcherMutationMu.Lock()
	if !d.memoryWatcherMutationClosed {
		d.memoryWatcherMutationClosed = true
		d.memoryWatcherMutationGeneration++
	}
	owners := make([]<-chan struct{}, 0, len(d.memoryWatcherMutationOwners))
	for owner := range d.memoryWatcherMutationOwners {
		if owner != nil && owner.done != nil {
			owners = append(owners, owner.done)
		}
	}
	d.memoryWatcherMutationMu.Unlock()
	d.closeIdleMemoryRetries()
	return owners
}

func waitForMemoryWatcherMutationOwners(ctx context.Context, owners []<-chan struct{}) error {
	return waitForOwnerChannels(ctx, owners)
}

func waitForOwnerChannels(ctx context.Context, owners []<-chan struct{}) error {
	if ctx == nil {
		ctx = context.Background()
	}
	for _, done := range owners {
		if done == nil {
			continue
		}
		select {
		case <-done:
		case <-ctx.Done():
			return ctx.Err()
		}
	}
	return nil
}

func (d *Dispatcher) memoryWatcherMutationOwnerCurrentLocked(owner *memoryWatcherMutationOwner) bool {
	if d == nil || owner == nil || d.memoryWatcherMutationClosed || owner.generation != d.memoryWatcherMutationGeneration {
		return false
	}
	_, ok := d.memoryWatcherMutationOwners[owner]
	return ok
}

// fenceSessionMemoryWatcherPublication permanently closes watcher publication
// for a session whose owning Process is terminal. The flag and
// publishMemoryWatcher's second validation share sess.mu: publication that wins
// first is visible to the following drain, while a fence that wins first makes
// the mutation owner reject and clean up its unpublished watcher.
func fenceSessionMemoryWatcherPublication(sess *interactiveSession) {
	if sess == nil {
		return
	}
	sess.mu.Lock()
	sess.memoryWatcherTerminalFenced = true
	sess.mu.Unlock()
}

// fenceSessionMemoryWatcherForFinalRemoval installs the terminal publication
// fence before a slow FlushAndClose/final push without retiring the
// manager-visible session early. RemoveIfCurrent performs whole-session
// retirement only after the drain. The optional callback is only the narrow
// final-removal test seam; other terminal fences must not trigger it.
func (d *Dispatcher) fenceSessionMemoryWatcherForFinalRemoval(sess *interactiveSession) {
	if sess == nil {
		return
	}
	fenceSessionMemoryWatcherPublication(sess)
	if d != nil && d.memoryWatcherFinalRemovalAfterFence != nil {
		d.memoryWatcherFinalRemovalAfterFence(sess)
	}
}

// waitForMemoryWatcherQuiescence joins both forms of owner that can temporarily
// hold an unpublished watcher:
//
//   - memory mutation owners span clone/construct through publish/reject cleanup;
//   - session start owners span an attached-but-not-yet-published session through
//     PublishStart failure and abortSessionLaunch cleanup.
//
// Admission for both registries must already be closed before calling this.
func waitForMemoryWatcherQuiescence(
	ctx context.Context,
	memoryOwners, sessionOwners []<-chan struct{},
) error {
	if err := waitForMemoryWatcherMutationOwners(ctx, memoryOwners); err != nil {
		return err
	}
	return waitForSessionStartOwners(ctx, sessionOwners)
}

// shutdownMemoryWatchersAfterOwnerDrain preserves the bounded shutdown API
// without weakening the teardown order. The cleanup continuation is launched
// unconditionally and survives ctx expiry; the caller only waits for it until
// the shared shutdown deadline. Once all admitted mutation/start owners have
// unwound, SessionManager is closed and its final snapshot is stable, so every
// published watcher is drained before the search registry is disabled.
func (d *Dispatcher) shutdownMemoryWatchersAfterOwnerDrain(
	ctx context.Context,
	memoryOwners, sessionOwners []<-chan struct{},
	shutdownSessions []*interactiveSession,
) {
	if ctx == nil {
		ctx = context.Background()
	}
	done := make(chan struct{})
	go func() {
		defer close(done)
		_ = waitForMemoryWatcherQuiescence(context.Background(), memoryOwners, sessionOwners)
		// shutdownSessions was captured atomically with SessionManager admission
		// close, so it retains a drain handle even if a concurrent teardown removes
		// a published session while owners unwind. Merge the now-stable final
		// snapshot defensively and deduplicate by session pointer.
		sessions := mergeMemoryWatcherShutdownSessions(shutdownSessions, d.sessionManager.Snapshot(true))
		d.drainMemoryWatchersForShutdown(sessions)
		d.disableMemorySearch()
	}()

	select {
	case <-done:
	case <-ctx.Done():
	}
}

func mergeMemoryWatcherShutdownSessions(groups ...[]*interactiveSession) []*interactiveSession {
	seen := make(map[*interactiveSession]struct{})
	var merged []*interactiveSession
	for _, sessions := range groups {
		for _, sess := range sessions {
			if sess == nil {
				continue
			}
			if _, ok := seen[sess]; ok {
				continue
			}
			seen[sess] = struct{}{}
			merged = append(merged, sess)
		}
	}
	return merged
}

// drainMemoryWatchersForShutdown drains a stable session snapshot in parallel.
// All Add calls happen before Wait begins; unlike a dynamically shared
// WaitGroup, this cannot race watcher publication or shutdown admission.
func (d *Dispatcher) drainMemoryWatchersForShutdown(sessions []*interactiveSession) {
	var watchers sync.WaitGroup
	watchers.Add(len(sessions))
	for _, sess := range sessions {
		sess := sess
		go func() {
			defer watchers.Done()
			d.stopMemoryWatcherForShutdown(sess)
		}()
	}
	watchers.Wait()
}