package dispatcher

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"math/rand"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/tidbcloud/AgentTeam/desktop/internal/memorysearch"
	"github.com/tidbcloud/AgentTeam/desktop/internal/profile"
	"github.com/tidbcloud/AgentTeam/desktop/internal/vault"
)

// memorySearchRegistry owns the per-binding memorysearch.Reconciler instances
// for the dispatcher. It is daemon-scoped: a reconciler lives across sessions
// for as long as the daemon owns the binding, so triggers independent of any
// interactive session (startup, periodic, tool-time) can drive reconciliation.
//
// The session watcher only holds a ChangeListener reference (Reconciler.OnMemoryChanged
// via memorywatcher.Options.Notify); it never closes a reconciler. Lifecycle
// is controlled here and from the dispatcher shutdown / rebind paths.
//
// Keys are binding fingerprints, not agentIDs: the same agent rebinding to a
// different memoryID/branch produces a different fingerprint, and the old
// reconciler must close before the new one acquires the per-binding owner lock.
//
// Slow-path serialization (Round 10 blocker):
//   - agentCreateMu: a sync.Map of per-agent *sync.Mutex. getOrCreate's slow
//     path acquires the agent's mutex, serializing all concurrent constructions
//     for that agent. Two different-fingerprint getOrCreate calls for the same
//     agent can no longer race on Start and the per-derivative-root owner lock
//     (which is shared across fingerprints for the same agent). The second
//     caller waits for the first to either land its entry or fail; on wake it
//     re-checks the fast path (entry may already be present) and the epoch gate.
//   - inflight: tracks the access being constructed by each agent's in-flight
//     slow path. closeAgentIfBindingChanged inspects BOTH entries and inflight
//     so it can detect identity drift even when the old binding's entry hasn't
//     landed yet. When the in-flight identity disagrees with the authoritative
//     access, the gate bumps agentEpoch — the in-flight getOrCreate observes
//     the bump at its insert gate, closes its freshly-Started reconciler
//     (releasing the owner lock), and returns ErrAgentRevoked.
//
// Lifecycle fencing (Round 9 blockers):
//   - closed: set once by closeAll and never reset. Any subsequent getOrCreate
//     observes it at every lock acquisition and returns ErrRegistryClosed.
//   - agentEpoch[agentID]: bumped by closeAgent / closeAgentIfBindingChanged
//     BEFORE entries are removed. getOrCreate captures the epoch when it
//     enters the slow path and re-checks it before inserting. If a revoke/rebind
//     closes the agent mid-Start, the in-flight getOrCreate observes the epoch
//     bump at the insert gate, closes the freshly-Started reconciler, and
//     returns ErrAgentRevoked — the entry never lands in the map and the
//     per-binding owner lock is released.
//
// Together these prevent both the "in-flight getOrCreate revives a closed
// registry" hazard (closed + epoch fencing) and the "different-fingerprint
// getOrCreate races owner lock with stale identity" hazard (per-agent
// serialization + inflight tracking).
type memorySearchRegistry struct {
	resolver memorysearch.AgentWorkDirResolver
	logger   *slog.Logger
	manager  *memorysearch.Manager
	routeMu  sync.RWMutex
	routes   map[string]string
	eligible map[string]*memorysearch.Binding

	mu       sync.Mutex
	entries  map[string]*searchEntry // keyed by Binding.Fingerprint()
	inflight map[string]inflightInfo // agentID → slow-path identity, see agentCreateMu

	// inflightWg counts getOrCreate slow paths that have committed to
	// acquiring (or having acquired) the per-binding owner lock but have
	// not yet either inserted their entry or self-cancelled at the insert
	// gate. closeAll Waits on this before returning so EnableMemorySearch
	// and disableMemorySearch cannot publish the next registry (or complete
	// shutdown) while a prior registry's in-flight Start still holds the
	// owner lock — without this drain, the next registry's first getOrCreate
	// would race on the owner lock and fail with ErrOwnerLockBusy.
	//
	// Add(1) happens under r.mu in the entry gate AFTER the closed/entry
	// checks pass; Done() runs via defer at function exit. closeAll's Wait
	// therefore observes a stable snapshot of inflight goroutines that
	// already committed (any new getOrCreate after closeAll sets closed
	// observes closed=true and returns without Add'ing).
	inflightWg sync.WaitGroup

	// agentCreateMu lazily allocates a *sync.Mutex per agentID to serialize
	// the getOrCreate slow path. Per-agent (not per-fingerprint) because the
	// per-derivative-root owner lock is shared across all fingerprints for
	// the same agent — concurrent Start calls for different fingerprints
	// would otherwise race on that lock with nondeterministic winner.
	agentCreateMu sync.Map

	// closed is set by closeAll and never cleared. Future getOrCreate calls
	// observe it at lock acquisition and refuse without touching entries.
	closed bool
	// agentEpoch tracks per-agent lifecycle generations. Each closeAgent
	// call for an agentID bumps the counter before deleting entries, so an
	// in-flight getOrCreate for the same agent can detect that its captured
	// epoch is stale and refuse to insert.
	agentEpoch map[string]uint64
	// sourceTrusted is the authoritative access gate. Missing means trusted
	// for backward-compatible startup; false freezes existing and in-flight
	// reconcilers before their registry insert gate.
	sourceTrusted map[string]bool

	periodicStop     chan struct{}
	periodicDone     chan struct{}
	periodicWake     chan struct{}
	periodicInterval time.Duration
	periodicJitter   func(time.Duration) time.Duration

	// testHooks is nil in production. Tests inject callbacks at precise
	// points in the getOrCreate slow path to reproduce races that depend on
	// exact timing (close-while-Start-in-flight). See memorySearchRegistryTestHooks.
	testHooks memorySearchRegistryTestHooks
}

// inflightInfo records the identity of an in-flight getOrCreate slow path
// for an agent. closeAgentIfBindingChanged reads this (under r.mu) to detect
// drift even when the entry has not yet landed in r.entries.
type inflightInfo struct {
	fp     string
	access MemoryAccess
}

// memorySearchRegistryTestHooks are timing seams for the getOrCreate slow
// path and the EnableMemorySearch swap path. Production leaves every field
// nil. Tests inject channel-parking callbacks to block the path at the named
// point so a racing close can be observed at the right gate.
type memorySearchRegistryTestHooks struct {
	// beforeInsert fires after Start succeeds and the initial Submit is
	// enqueued, but BEFORE the insert gate acquires the registry mutex.
	// The hook receives agentID + fingerprint so the test can identify
	// which getOrCreate is paused. A nil hook is a no-op.
	beforeInsert func(agentID, fp string)

	// beforeCloseAll fires at the start of closeAll, before the closed flag
	// is set and before any entry is closed. Tests use it to park the
	// EnableMemorySearch swap path during the prev.closeAll() window so a
	// concurrent Refresh can be observed. A nil hook is a no-op.
	beforeCloseAll func()
}

const defaultMemorySearchPeriodicInterval = 30 * time.Minute

// searchEntry pairs a binding with its owning reconciler.
type searchEntry struct {
	binding    *memorysearch.Binding
	reconciler *memorysearch.Reconciler
}

// newMemorySearchRegistry constructs a registry. The resolver must agree with
// the one used to build bindings — production wires profile.AgentWorkDir so
// every reconciler's ResolvedRoot is `<profile.AgentWorkDir(agentID)>/memory`
// regardless of which session workDir triggered the getOrCreate.
func newMemorySearchRegistry(
	resolver memorysearch.AgentWorkDirResolver,
	scan func(ctx context.Context, root string) (memorysearch.SourceSnapshot, memorysearch.ScanStatus, error),
	logger *slog.Logger,
) *memorySearchRegistry {
	if resolver == nil {
		resolver = memorysearch.AgentWorkDirResolver(profile.AgentWorkDir)
	}
	if logger == nil {
		logger = slog.Default()
	}
	manager, err := memorysearch.NewManager(memorysearch.ManagerOptions{
		Resolver: resolver,
		Scan:     scan,
		Logger:   logger,
	})
	if err != nil {
		logger.Error("memory search registry: construct manager failed", "error", err)
	}
	registry := &memorySearchRegistry{
		resolver:         resolver,
		logger:           logger,
		manager:          manager,
		routes:           make(map[string]string),
		eligible:         make(map[string]*memorysearch.Binding),
		entries:          make(map[string]*searchEntry),
		inflight:         make(map[string]inflightInfo),
		agentEpoch:       make(map[string]uint64),
		sourceTrusted:    make(map[string]bool),
		periodicStop:     make(chan struct{}),
		periodicDone:     make(chan struct{}),
		periodicWake:     make(chan struct{}, 1),
		periodicInterval: defaultMemorySearchPeriodicInterval,
		periodicJitter: func(base time.Duration) time.Duration {
			if base <= 0 {
				return base
			}
			// Bounded positive jitter spreads daemon wakeups over 10%.
			return base + time.Duration(rand.Int63n(max(1, int64(base/10))))
		},
	}
	go registry.periodicLoop()
	return registry
}

func (r *memorySearchRegistry) cacheEligible(binding *memorysearch.Binding) {
	if binding == nil {
		return
	}
	copy := *binding
	r.routeMu.Lock()
	r.eligible[binding.LocalAgentID] = &copy
	r.routeMu.Unlock()
}

func (r *memorySearchRegistry) clearEligible(agentID string) {
	r.routeMu.Lock()
	delete(r.eligible, strings.TrimSpace(agentID))
	r.routeMu.Unlock()
}

func (r *memorySearchRegistry) eligibleBinding(agentID string) *memorysearch.Binding {
	r.routeMu.RLock()
	binding := r.eligible[strings.TrimSpace(agentID)]
	r.routeMu.RUnlock()
	if binding == nil {
		return nil
	}
	copy := *binding
	return &copy
}

func (r *memorySearchRegistry) bindRoute(routeAgentID, localAgentID string) {
	routeAgentID = strings.TrimSpace(routeAgentID)
	localAgentID = strings.TrimSpace(localAgentID)
	if routeAgentID == "" || localAgentID == "" {
		return
	}
	r.routeMu.Lock()
	r.routes[routeAgentID] = localAgentID
	r.routeMu.Unlock()
}

func (r *memorySearchRegistry) resolvedAgentID(routeAgentID string) string {
	routeAgentID = strings.TrimSpace(routeAgentID)
	r.routeMu.RLock()
	localAgentID := r.routes[routeAgentID]
	r.routeMu.RUnlock()
	if localAgentID != "" {
		return localAgentID
	}
	return routeAgentID
}

// getOrCreate returns the Reconciler for the binding identified by (agentID,
// access), constructing one if none exists. The clone-side retention metadata
// at the resolved root's parent must agree with the access before a binding
// is constructed; revoked clones or memory/branch/agent mismatches fail closed.
//
// Serialization (Round 10 blocker 2): the slow path is serialized per-agent
// via agentCreateMu so two different-fingerprint getOrCreate calls for the
// same agent cannot race on Start and the per-derivative-root owner lock
// (which is shared across fingerprints for the same agent). The second caller
// waits for the first to either land its entry or fail; on wake it re-checks
// the fast path (entry may already be present), the epoch gate, and the
// entry-drift gate (an entry for the same agent with a different fingerprint
// means closeAgentIfBindingChanged was skipped — refuse with ErrAgentRevoked
// rather than letting the caller race the owner lock).
//
// Lifecycle fencing (Round 9 blockers): the slow path runs Start OUTSIDE the
// registry mutex, so a concurrent closeAgent or closeAll could otherwise
// complete before the slow path inserts its freshly-Started reconciler —
// leaving an entry in the map despite the close, with the reconciler holding
// the per-binding owner lock. To prevent that, the slow path:
//  1. Captures agentEpoch[agentID] BEFORE acquiring the per-agent mutex.
//  2. Re-checks closed + epoch + entry drift AFTER acquiring the mutex.
//  3. Records inflight[agentID] so closeAgentIfBindingChanged can detect
//     drift even before the entry lands.
//  4. Re-checks closed + epoch + entries + drift BEFORE insert (post-Start).
//
// Any failure at these gates closes the freshly-Started reconciler (releasing
// the lock) and returns ErrRegistryClosed / ErrAgentRevoked. The caller
// (lookupSearchReconciler, preCloneMemory) treats these as "no reconciler
// available" and degrades gracefully.
func (r *memorySearchRegistry) getOrCreate(
	ctx context.Context,
	agentID string,
	access MemoryAccess,
) (*memorysearch.Reconciler, *memorysearch.Binding, error) {
	// Authoritative deny precedes retention/filesystem validation. A stale
	// access object arriving after Unbound/Unknown must fail from the in-memory
	// gate without touching the clone or reviving a writer.
	r.mu.Lock()
	trusted, known := r.sourceTrusted[agentID]
	r.mu.Unlock()
	if known && !trusted {
		return nil, nil, ErrAgentRevoked
	}
	binding, err := r.constructBinding(agentID, access)
	if err != nil {
		return nil, nil, err
	}
	if r.manager == nil {
		return nil, nil, fmt.Errorf("memory search registry: manager unavailable")
	}
	fp := binding.Fingerprint()

	// (1) Entry gate: fast path + epoch capture under the lock.
	r.mu.Lock()
	if r.closed {
		r.mu.Unlock()
		return nil, nil, ErrRegistryClosed
	}
	if e, ok := r.entries[fp]; ok {
		r.mu.Unlock()
		r.cacheEligible(e.binding)
		return e.reconciler, e.binding, nil
	}
	epochAtStart := r.agentEpoch[agentID]
	r.mu.Unlock()

	// (2) Per-agent serialization. Two concurrent getOrCreate calls for the
	// same agentID serialize here; the second waits for the first to either
	// land its entry (fast path on re-check) or fail (re-enters its own
	// slow path with a fresh epoch capture if needed).
	agentMu := r.agentMutexFor(agentID)
	agentMu.Lock()
	defer agentMu.Unlock()

	// (3) Post-mutex re-check. A peer may have just inserted; a close may
	// have marked the registry closed or bumped this agent's epoch; or a
	// peer may have inserted a DIFFERENT fingerprint for the same agent
	// (programmer error: closeAgentIfBindingChanged was skipped).
	r.mu.Lock()
	if r.closed {
		r.mu.Unlock()
		return nil, nil, ErrRegistryClosed
	}
	if e, ok := r.entries[fp]; ok {
		r.mu.Unlock()
		r.cacheEligible(e.binding)
		return e.reconciler, e.binding, nil
	}
	if r.agentEpoch[agentID] != epochAtStart {
		r.mu.Unlock()
		return nil, nil, ErrAgentRevoked
	}
	trusted, known = r.sourceTrusted[agentID]
	if known && !trusted {
		r.mu.Unlock()
		return nil, nil, ErrAgentRevoked
	}
	if drift := r.entryForAgentLocked(agentID); drift != nil && drift.binding.Fingerprint() != fp {
		r.mu.Unlock()
		return nil, nil, fmt.Errorf(
			"memory search registry: agent %s has active fingerprint %s, requested %s: %w",
			agentID, drift.binding.Fingerprint(), fp, ErrAgentRevoked)
	}
	// Record inflight identity so closeAgentIfBindingChanged can detect
	// drift even when this slow path has not yet inserted its entry.
	// Add to inflightWg BEFORE releasing r.mu so closeAll that observes
	// the inflight entry is guaranteed to see a non-zero WaitGroup and
	// block on its Wait(); without this ordering closeAll could clear
	// entries and call Wait() before Add(1) lands, returning prematurely.
	r.inflight[agentID] = inflightInfo{fp: fp, access: access}
	r.inflightWg.Add(1)
	r.mu.Unlock()
	defer func() {
		r.mu.Lock()
		// Only clear our own inflight entry; a later getOrCreate that
		// already started under a new identity must not be clobbered.
		if cur, ok := r.inflight[agentID]; ok && cur.fp == fp {
			delete(r.inflight, agentID)
		}
		r.mu.Unlock()
		// Release the closeAll drain token AFTER inflight cleanup so
		// closeAll's Wait() cannot return while this entry is still
		// observable in r.inflight.
		r.inflightWg.Done()
	}()

	rec, err := r.manager.GetOrCreate(ctx, binding)
	if errors.Is(err, memorysearch.ErrBindingTransitionRequired) {
		active := r.manager.ActiveFingerprint(agentID)
		if active == "" {
			return nil, nil, fmt.Errorf("memory search registry: transition without active binding: %w", err)
		}
		if fenceErr := r.manager.FenceBinding(ctx, active); fenceErr != nil {
			return nil, nil, fmt.Errorf("memory search registry: fence transition: %w", fenceErr)
		}
		// One retry only. A second transition or any other error fails closed;
		// callers must wait for the next authoritative refresh.
		rec, err = r.manager.GetOrCreate(ctx, binding)
	}
	if err != nil {
		return nil, nil, fmt.Errorf("memory search registry: get reconciler: %w", err)
	}
	// Apply the latest authoritative gate even before the adapter entry lands.
	r.mu.Lock()
	trusted, known = r.sourceTrusted[agentID]
	r.mu.Unlock()
	if known && !trusted {
		rec.SetSourceTrusted(false)
	}
	// Start already enqueues initial convergence. A second known-dirty Submit
	// here races the first successful publish and can leave public state
	// refreshing until an unnecessary drain no-op finishes.

	// (4) Test seam: pause here so a racing close can be observed at the
	// insert gate. Production testHooks.beforeInsert is nil — no-op.
	if r.testHooks.beforeInsert != nil {
		r.testHooks.beforeInsert(agentID, fp)
	}

	entry := &searchEntry{binding: binding, reconciler: rec}
	// (5) Insert gate: full fence re-check inside the lock. Any failure
	// closes the freshly-Started reconciler (releasing the per-binding
	// owner lock) and returns the matching sentinel. Without this gate
	// a closeAgent/closeAll that raced the slow path's Start would be
	// hidden — the entry would land and the reconciler would keep
	// holding the lock despite the close.
	r.mu.Lock()
	switch {
	case r.closed:
		r.mu.Unlock()
		return nil, nil, ErrRegistryClosed
	case r.agentEpoch[agentID] != epochAtStart:
		r.mu.Unlock()
		return nil, nil, ErrAgentRevoked
	case func() bool {
		trusted, known := r.sourceTrusted[agentID]
		return known && !trusted
	}():
		r.mu.Unlock()
		rec.SetSourceTrusted(false)
		if active := r.manager.ActiveFingerprint(agentID); active != "" {
			_ = r.manager.FenceBinding(context.Background(), active)
		}
		return nil, nil, ErrAgentRevoked
	}
	// Peer race: another getOrCreate for the same fp inserted while we
	// were paused at the test hook or in Start. Close the duplicate.
	if existing, ok := r.entries[fp]; ok {
		r.mu.Unlock()
		return existing.reconciler, existing.binding, nil
	}
	// Drift re-check: another getOrCreate for a DIFFERENT fp inserted
	// for the same agent while we were paused. Same programmer-error
	// response as the pre-Start gate.
	if drift := r.entryForAgentLocked(agentID); drift != nil && drift.binding.Fingerprint() != fp {
		r.mu.Unlock()
		return nil, nil, fmt.Errorf(
			"memory search registry: agent %s drift detected at insert gate: has %s, requested %s: %w",
			agentID, drift.binding.Fingerprint(), fp, ErrAgentRevoked)
	}
	r.entries[fp] = entry
	r.mu.Unlock()
	r.cacheEligible(binding)

	r.logger.InfoContext(ctx, "memory search registry: reconciler started",
		"agent_id", agentID, "memory_id", access.MemoryID, "fingerprint", fp)
	return entry.reconciler, entry.binding, nil
}

// agentMutexFor returns the per-agent serialization mutex, allocating one on
// first use. Per-agent (not per-fingerprint) because the per-derivative-root
// owner lock is shared across all fingerprints for the same agent.
//
// Mutexes are intentionally not garbage-collected on agent revoke: agentIDs
// are workspace-scoped stable identifiers whose population is bounded by the
// workspace's agent count (typically dozens). Each *sync.Mutex is ~24 bytes,
// so the worst-case leak under agent churn is O(workspace agents × 24B),
// which is negligible for a desktop daemon. Active cleanup would require
// reference counting to avoid racing an in-flight getOrCreate that captured
// the pointer — the complexity is not worth the savings.
func (r *memorySearchRegistry) agentMutexFor(agentID string) *sync.Mutex {
	v, _ := r.agentCreateMu.LoadOrStore(agentID, &sync.Mutex{})
	return v.(*sync.Mutex)
}

// constructBinding builds a validated Binding and verifies the clone-side
// retention marker matches. Returns ErrRetentionMismatch (wrapped) if the
// clone is revoked or its memory/branch/agent identity disagrees with the
// access. A missing retention marker means the source clone has not been
// initialized by the watcher's clone flow — the reconciler would index a
// non-existent or unrelated directory, so we fail closed.
func (r *memorySearchRegistry) constructBinding(
	agentID string,
	access MemoryAccess,
) (*memorysearch.Binding, error) {
	binding, err := memorysearch.NewBinding(
		r.resolver,
		agentID,
		access.MemoryID,
		access.Branch,
		memorysearch.AnalyzerVersion,
		memorysearch.ChunkerVersion,
	)
	if err != nil {
		return nil, fmt.Errorf("memory search registry: construct binding: %w", err)
	}
	state, err := vault.ReadRetentionState(binding.ResolvedRoot)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return nil, fmt.Errorf("memory search registry: retention marker missing under %s: %w",
				binding.ResolvedRoot, errRetentionMissing)
		}
		return nil, fmt.Errorf("memory search registry: read retention state: %w", err)
	}
	meta := memorysearch.RetentionMetadata{
		SchemaVersion: state.SchemaVersion,
		MemoryID:      state.MemoryID,
		AgentID:       state.AgentID,
		Branch:        state.Branch,
		RevokedAt:     state.RevokedAt,
	}
	if err := binding.MatchesRetention(meta); err != nil {
		return nil, fmt.Errorf("memory search registry: retention check: %w", err)
	}
	return binding, nil
}

// closeAgent closes every entry whose binding belongs to agentID and bumps
// the agent's epoch. Used during rebind and revoke so the old fingerprint
// releases the per-binding owner lock before the new fingerprint is
// constructed, and so an in-flight getOrCreate for the same agent observes
// the bump at its insert gate and refuses to land a stale entry.
//
// The epoch bump MUST happen under the same lock that deletes entries, and
// BEFORE the delete. This way any getOrCreate that captured the old epoch
// and reaches its insert gate after this call returns is guaranteed to see
// both the bump (epoch mismatch → ErrAgentRevoked) and the missing entry
// (it would otherwise re-insert). The caller must have already stopped /
// drained the watcher so no new Notify arrives against the closing reconciler.
func (r *memorySearchRegistry) closeAgent(agentID string) {
	r.clearEligible(agentID)
	r.mu.Lock()
	r.agentEpoch[agentID]++
	for fp, e := range r.entries {
		if e.binding.LocalAgentID != agentID {
			continue
		}
		delete(r.entries, fp)
	}
	r.mu.Unlock()

	// Manager is the authoritative owner. Its handle may already be active
	// while the adapter slow path is paused before inserting searchEntry, so
	// deriving the transition only from entries would leave activeFP exposed
	// until that slow path resumed. Clear it synchronously through Manager.
	if fingerprint := r.manager.ActiveFingerprint(agentID); fingerprint != "" {
		_ = r.manager.MarkUnbound(context.Background(), fingerprint)
		r.logger.InfoContext(context.Background(), "memory search registry: reconciler closed (agent)",
			"agent_id", agentID, "fingerprint", fingerprint)
	}
}

func (r *memorySearchRegistry) dropAgentEntries(agentID string) {
	r.clearEligible(agentID)
	r.mu.Lock()
	defer r.mu.Unlock()
	r.agentEpoch[agentID]++
	for fingerprint, entry := range r.entries {
		if entry.binding.LocalAgentID == agentID {
			delete(r.entries, fingerprint)
		}
	}
}

// closeAll closes every entry and marks the registry permanently closed. Used
// at dispatcher shutdown. After closeAll returns, any subsequent getOrCreate
// returns ErrRegistryClosed without touching entries, and any in-flight
// getOrCreate observes the closed flag at its insert gate, closes its
// freshly-Started reconciler, and returns ErrRegistryClosed.
//
// Swap-then-close: the closed flag and the empty entries map are installed
// under the mutex; reconcilers are closed outside so concurrent Submit/Current
// callers see an empty registry immediately rather than blocking on a slow
// shutdown. The registry is single-use: closed is never cleared.
//
// Inflight drain (Round 10 final blocker): closeAll must not return while a
// prior slow-path getOrCreate is still in flight (Start already done or about
// to run) — otherwise EnableMemorySearch publishes `next` before prev's in-flight
// Start releases the per-binding owner lock, and next's first getOrCreate races
// that lock with ErrOwnerLockBusy. The drain works because:
//   - closed=true is set under r.mu BEFORE any getOrCreate can set r.inflight
//   - inflightWg.Add(1) under r.mu — a getOrCreate that arrives after
//     closeAll observes closed=true at its post-mutex re-check and returns
//     without Add'ing.
//   - Any getOrCreate that already Add(1)'d is committed to either land an
//     entry (impossible once closed=true) or self-cancel at the insert gate
//     (closing its freshly-Started reconciler, releasing the owner lock).
//   - closeAll's Wait() therefore observes a stable set of in-flight
//     goroutines that will all converge on Done() in bounded time.
//
// The Wait() runs OUTSIDE r.mu so the in-flight getOrCreate can acquire r.mu
// at its insert gate (to observe closed=true) and at its deferred inflight
// cleanup. Holding r.mu while waiting would deadlock the drain.
//
// Test seam: testHooks.beforeCloseAll (when non-nil) fires at the start of
// closeAll, BEFORE the closed flag is set and BEFORE any entry is closed.
// Tests use it to park EnableMemorySearch's prev.closeAll() window so a
// concurrent Refresh can be observed mid-replacement. Production leaves it nil.
func (r *memorySearchRegistry) closeAll() {
	if r.testHooks.beforeCloseAll != nil {
		r.testHooks.beforeCloseAll()
	}
	r.mu.Lock()
	if r.closed {
		r.mu.Unlock()
		return
	}
	r.closed = true
	close(r.periodicStop)
	entries := r.entries
	r.entries = make(map[string]*searchEntry)
	r.mu.Unlock()

	for _, e := range entries {
		if e.reconciler != nil {
			r.logger.InfoContext(context.Background(), "memory search registry: reconciler closed (shutdown)",
				"agent_id", e.binding.LocalAgentID, "fingerprint", e.binding.Fingerprint())
		}
	}

	// Wait for any in-flight slow path that already Add(1)'d before
	// closeAll set closed=true. Those goroutines observe closed=true at
	// their insert gate, self-cancel (releasing the per-binding owner
	// lock), and call Done() via their deferred cleanup. Without this
	// Wait, EnableMemorySearch could publish next before prev's in-flight
	// Start released the owner lock, causing next's first getOrCreate to
	// hit ErrOwnerLockBusy.
	r.inflightWg.Wait()
	<-r.periodicDone
	if r.manager != nil {
		_ = r.manager.Close()
	}
}

// Submit triggers a reconcile pass for the given agent. It is the entry point
// the daemon will wire to periodic and tool-time triggers in Task 8. Agents
// without a registered reconciler are silently skipped.
func (r *memorySearchRegistry) Submit(ctx context.Context, agentID string) {
	r.mu.Lock()
	rec := r.reconcilerForAgentLocked(agentID)
	r.mu.Unlock()
	if rec != nil {
		rec.Submit(ctx)
	}
}

func (r *memorySearchRegistry) periodicLoop() {
	defer close(r.periodicDone)
	for {
		delay := r.periodicInterval
		if r.periodicJitter != nil {
			delay = r.periodicJitter(delay)
		}
		if delay <= 0 {
			delay = defaultMemorySearchPeriodicInterval
		}
		timer := time.NewTimer(delay)
		select {
		case <-r.periodicStop:
			if !timer.Stop() {
				<-timer.C
			}
			return
		case <-r.periodicWake:
			if !timer.Stop() {
				<-timer.C
			}
			r.submitAllPeriodic()
		case <-timer.C:
			r.submitAllPeriodic()
		}
	}
}

func (r *memorySearchRegistry) submitAllPeriodic() {
	r.mu.Lock()
	if r.closed {
		r.mu.Unlock()
		return
	}
	reconcilers := make([]*memorysearch.Reconciler, 0, len(r.entries))
	for _, entry := range r.entries {
		if entry.reconciler != nil {
			reconcilers = append(reconcilers, entry.reconciler)
		}
	}
	r.mu.Unlock()
	for _, rec := range reconcilers {
		// Reconciler.Submit is non-blocking and independently drops the
		// trigger while authoritative access is Unknown.
		rec.CheckFreshness(context.Background())
	}
}

// setSourceTrusted freezes/resumes derivative writes for the agent without
// removing the last trusted serving generation. Unknown access uses false;
// authoritative Bound uses true and queues one reconcile on transition.
func (r *memorySearchRegistry) setSourceTrusted(agentID string, trusted bool) {
	if r.manager != nil {
		r.manager.SetSourceTrusted(agentID, trusted)
	}
	r.mu.Lock()
	r.sourceTrusted[agentID] = trusted
	rec := r.reconcilerForAgentLocked(agentID)
	if rec == nil && r.manager != nil {
		rec = r.manager.ActiveReconciler(agentID)
	}
	r.mu.Unlock()
	if rec != nil {
		rec.SetSourceTrusted(trusted)
	}
}

func (r *memorySearchRegistry) denyAgent(agentID string) {
	r.setSourceTrusted(agentID, false)
	r.mu.Lock()
	r.agentEpoch[agentID]++
	r.mu.Unlock()
}

func (r *memorySearchRegistry) rebuildAgent(agentID string) error {
	r.mu.Lock()
	rec := r.reconcilerForAgentLocked(agentID)
	r.mu.Unlock()
	if rec == nil && r.manager != nil {
		rec = r.manager.ActiveReconciler(agentID)
	}
	if rec == nil {
		return memorysearch.ErrNotReady
	}
	rec.ForceRebuild()
	return nil
}

// Current returns the serviceable generation for the agent, or
// memorysearch.ErrNotReady if no reconciler is registered (or its
// generation is not yet built). Callers must Close the returned Generation.
func (r *memorySearchRegistry) Current(ctx context.Context, agentID string) (*memorysearch.Generation, error) {
	r.mu.Lock()
	rec := r.reconcilerForAgentLocked(agentID)
	r.mu.Unlock()
	if rec == nil {
		return nil, memorysearch.ErrNotReady
	}
	return rec.Current(ctx)
}

// fingerprintForAgent returns the active binding fingerprint for the agent,
// or "" if no entry exists. Used by tests and by rebind to identify the entry
// to close before constructing a new one.
func (r *memorySearchRegistry) fingerprintForAgent(agentID string) string {
	r.mu.Lock()
	defer r.mu.Unlock()
	for fp, e := range r.entries {
		if e.binding.LocalAgentID == agentID {
			return fp
		}
	}
	return ""
}

func (r *memorySearchRegistry) reconcilerForAgentLocked(agentID string) *memorysearch.Reconciler {
	for _, e := range r.entries {
		if e.binding.LocalAgentID == agentID {
			return e.reconciler
		}
	}
	return nil
}

// entryForAgentLocked returns the active searchEntry for agentID, or nil if
// none exists. Used by getOrCreate's drift check to detect "an entry for the
// same agent with a different fingerprint already exists" — a programmer
// error (closeAgentIfBindingChanged was skipped) that must not race the
// per-derivative-root owner lock. Caller MUST hold r.mu.
func (r *memorySearchRegistry) entryForAgentLocked(agentID string) *searchEntry {
	for _, e := range r.entries {
		if e.binding.LocalAgentID == agentID {
			return e
		}
	}
	return nil
}

// closeAgentIfBindingChanged is the registry identity gate. It closes the
// agent's reconciler when the authoritative access's (memoryID, branch)
// disagrees with the active entry's identity OR with an in-flight slow
// path's recorded identity, so the new binding's Start can acquire the
// per-derivative-root owner lock and publish a fresh fingerprint.
//
// Returns true when it closed at least one entry OR bumped the epoch (caller
// should proceed with fresh clone/watcher construction); false when no entry
// exists and no slow path is in flight for the agent, or when all observed
// identities already match the access (caller may keep the existing generation).
//
// Inspects BOTH r.entries AND r.inflight under r.mu. The inflight scan is
// critical for the Round 10 blocker 2 race: when a slow path has Start-ed
// its reconciler but not yet inserted (paused at test hook, or just running
// Start concurrently), the entry is invisible to an entries-only scan. If
// the in-flight identity differs from the authoritative access, bumping the
// epoch makes the in-flight getOrCreate self-cancel at its insert gate
// rather than landing a stale entry that holds the owner lock.
//
// This gate MUST run before any source clone mutation (preCloneMemory,
// attachMemoryWatcher, rebindAgentMemory) — otherwise the old reconciler
// keeps holding the per-derivative-root owner lock and the new binding's Start
// returns ErrOwnerLockBusy; worse, the old reconciler may index new-binding
// content against the old fingerprint while the clone is being mutated.
//
// This gate does NOT take agentCreateMu: doing so would block until the
// in-flight slow path releases the mutex, defeating the purpose of the
// inflight scan (the gate must run unconditionally so the in-flight getOrCreate
// observes the epoch bump at its insert gate before it can land).
func (r *memorySearchRegistry) closeAgentIfBindingChanged(agentID string, access MemoryAccess) bool {
	r.mu.Lock()
	// Inspect active entries for this agent and decide whether identity changed.
	var (
		fromMemoryID, fromBranch string
		seen                     bool
	)
	for _, e := range r.entries {
		if e.binding.LocalAgentID != agentID {
			continue
		}
		fromMemoryID = e.binding.MemoryID
		fromBranch = e.binding.Branch
		seen = true
		break
	}
	// Also inspect the in-flight slow path's recorded identity. Even when no
	// entry has landed yet, an in-flight getOrCreate with a different identity
	// must trigger the epoch bump so it self-cancels.
	var (
		inflightMemoryID, inflightBranch string
		inflightSeen                     bool
	)
	if info, ok := r.inflight[agentID]; ok {
		inflightMemoryID = info.access.MemoryID
		inflightBranch = info.access.Branch
		inflightSeen = true
	}

	entryMatch := seen &&
		fromMemoryID == access.MemoryID &&
		strings.TrimSpace(fromBranch) == strings.TrimSpace(access.Branch)
	inflightMatch := !inflightSeen ||
		(inflightMemoryID == access.MemoryID &&
			strings.TrimSpace(inflightBranch) == strings.TrimSpace(access.Branch))
	if (!seen || entryMatch) && inflightMatch {
		r.mu.Unlock()
		return false
	}

	// Identity differs (in entries, inflight, or both) — bump epoch under the
	// same lock that deletes (same invariant as closeAgent) so any in-flight
	// getOrCreate for the same agent observes the bump at its insert gate and
	// refuses to land a stale entry.
	r.agentEpoch[agentID]++
	var toClose []*searchEntry
	for fp, e := range r.entries {
		if e.binding.LocalAgentID == agentID {
			delete(r.entries, fp)
			toClose = append(toClose, e)
		}
	}
	r.mu.Unlock()
	// Keep routeMu and r.mu independent. In particular, do not acquire
	// routeMu while holding r.mu: MCP eligibility reads routeMu on a hot path,
	// while lifecycle transitions use r.mu.
	r.clearEligible(agentID)

	// As with closeAgent, use Manager.activeFP rather than the adapter entry
	// snapshot. The Manager may have finished GetOrCreate while this adapter
	// is paused at beforeInsert, and the identity fence must still be
	// synchronous and fail closed.
	if fingerprint := r.manager.ActiveFingerprint(agentID); fingerprint != "" {
		_ = r.manager.FenceBinding(context.Background(), fingerprint)
	}
	r.logger.InfoContext(context.Background(), "memory search registry: reconciler closed (identity change)",
		"agent_id", agentID,
		"memory_id_from", fromMemoryID, "memory_id_to", access.MemoryID,
		"branch_from", fromBranch, "branch_to", access.Branch,
		"closed_count", len(toClose),
		"inflight_seen", inflightSeen,
		"inflight_match", inflightMatch)
	return len(toClose) > 0 || (inflightSeen && !inflightMatch)
}

// snapshotMemorySearch returns the current registry pointer, bounded-waiting
// on ctx if EnableMemorySearch is mid-replacement. The wait is required so a
// Refresh that arrives during replacement converges on the new registry in
// the SAME call rather than degrading silently and waiting up to 30min for
// the next periodic trigger.
//
// Returns:
//   - (reg, nil) when a registry is wired and ready
//   - (nil, nil) when no registry is wired and no replacement is in progress
//     (truly pre-Enable, or post-shutdown)
//   - (nil, ctx.Err()) when the wait was cancelled before replacement completed
//
// Callers MUST NOT cache the returned pointer across long-lived operations:
// EnableMemorySearch may swap in a fresh registry, and the previous pointer's
// closeAll may run concurrently. Within a single logical operation (one
// Submit, one getOrCreate, one closeAgent), re-using the snapshot is safe —
// even a swapped-out registry remains valid for read-only method calls until
// its closeAll completes, and registry methods are no-ops on entries that
// closeAll has already removed.
func (d *Dispatcher) snapshotMemorySearch(ctx context.Context) (*memorySearchRegistry, error) {
	for {
		d.memorySearchMu.RLock()
		reg := d.memorySearch
		swap := d.memorySearchSwap
		closed := d.memorySearchClosed
		d.memorySearchMu.RUnlock()
		if reg != nil {
			return reg, nil
		}
		if closed {
			return nil, nil
		}
		if swap == nil {
			return nil, nil
		}
		select {
		case <-swap:
			// Replacement completed (or was superseded by a newer Enable that
			// closed this swap). Re-agentteam to re-read state — the new Enable
			// may have already published, or may itself be in progress with
			// its own fresh swap chan.
		case <-ctx.Done():
			return nil, ctx.Err()
		}
	}
}

// EnableMemorySearch wires the daemon-scoped memorysearch.Reconciler owner.
// Production wires profile.AgentWorkDir as the resolver so every reconciler's
// ResolvedRoot is `<profile.AgentWorkDir(agentID)>/memory` regardless of
// which session workDir later triggers a getOrCreate.
//
// Replacement transaction (Round 10 blocker 1): the full sequence
// (take-old → set-swap/nil → closeAll-old → publish-next → close-swap) runs
// under memorySearchReplace so concurrent Enables linearize by mutex
// acquisition order. The second Enable blocks until the first has fully
// released — including closeAll — which guarantees the previous registry's
// per-binding owner locks are released before the next registry's first
// getOrCreate. Without this serialization, E2 could publish next2 while E1's
// closeAll(prev) was still in progress, and Refresh's getOrCreate(next2)
// would hit ErrOwnerLockBusy (owner lock still held by prev).
//
// Refresh observes memorySearch=nil + memorySearchSwap!=nil during the
// replacement window and bounded-waits on the swap chan (via ctx) so the
// same Refresh call converges on the new registry.
//
// Safe to call after Close: returns without publishing once
// memorySearchClosed is set. This prevents a goroutine that was blocked on
// memorySearchReplace during shutdown from publishing a fresh registry onto a
// torn-down dispatcher.
//
// A nil resolver is normalized to profile.AgentWorkDir (production default);
// a nil logger falls back to slog.Default().
func (d *Dispatcher) EnableMemorySearch(
	resolver memorysearch.AgentWorkDirResolver,
	scan func(ctx context.Context, root string) (memorysearch.SourceSnapshot, memorysearch.ScanStatus, error),
	logger *slog.Logger,
) {
	if resolver == nil {
		resolver = memorysearch.AgentWorkDirResolver(profile.AgentWorkDir)
	}
	if logger == nil {
		logger = slog.Default()
	}
	next := newMemorySearchRegistry(resolver, scan, logger)

	// Serialize the full replacement transaction. Concurrent Enables block
	// here until the prior Enable has completed its closeAll + publish.
	d.memorySearchReplace.Lock()
	defer d.memorySearchReplace.Unlock()

	d.memorySearchMu.Lock()
	if d.memorySearchClosed {
		d.memorySearchMu.Unlock()
		logger.InfoContext(context.Background(),
			"memory search registry: EnableMemorySearch skipped (dispatcher closed)")
		return
	}
	prev := d.memorySearch
	swap := make(chan struct{})
	d.memorySearchSwap = swap
	d.memorySearch = nil
	d.memorySearchMu.Unlock()

	if prev != nil {
		// Close outside the memorySearchMu write lock so concurrent snapshot
		// readers don't block on slow reconciler shutdowns (in-flight build,
		// owner-lock release). memorySearchReplace still holds, so a
		// concurrent Enable waits for this closeAll to finish before it can
		// even begin its own replacement.
		prev.closeAll()
		logger.InfoContext(context.Background(),
			"memory search registry: previous registry closed after swap")
	}

	d.memorySearchMu.Lock()
	d.memorySearch = next
	d.memorySearchSwap = nil
	d.memorySearchMu.Unlock()
	close(swap)
}

// disableMemorySearch tears down the memory search registry at dispatcher
// shutdown. Idempotent: a second call is a no-op. Serialized by
// memorySearchReplace so it cannot race with an in-flight EnableMemorySearch
// transaction (Enable acquires the same mutex before publishing).
//
// The sequence:
//  1. Take memorySearchReplace (defer unlock) to serialize against Enable.
//  2. Take memorySearchMu, set memorySearchClosed=true (idempotency flag),
//     extract prev pointer + swap chan, clear memorySearch + memorySearchSwap
//     fields, release memorySearchMu.
//  3. Close the swap chan so any Refresh bounded-waiting on it wakes up and
//     observes (reg=nil, swap=nil, closed=true) → returns (nil, nil).
//  4. Close prev's reconcilers OUTSIDE the locks so a slow reconciler
//     shutdown (in-flight build, owner-lock release) does not block other
//     shutdown work.
//
// Once memorySearchClosed is set, any subsequent EnableMemorySearch call
// observes the flag under memorySearchMu and returns without publishing —
// preventing a goroutine that blocked on memorySearchReplace during shutdown
// from later publishing a fresh registry onto a torn-down dispatcher.
func (d *Dispatcher) disableMemorySearch() {
	d.memorySearchReplace.Lock()
	defer d.memorySearchReplace.Unlock()

	d.memorySearchMu.Lock()
	if d.memorySearchClosed {
		d.memorySearchMu.Unlock()
		return
	}
	d.memorySearchClosed = true
	prev := d.memorySearch
	swap := d.memorySearchSwap
	d.memorySearch = nil
	d.memorySearchSwap = nil
	d.memorySearchMu.Unlock()

	if swap != nil {
		close(swap)
	}
	if prev != nil {
		prev.closeAll()
		slog.InfoContext(context.Background(),
			"memory search registry: closed via dispatcher shutdown")
	}
}

// errRetentionMissing signals that no retention marker exists under the
// resolved root's parent. The source clone has not been initialized by
// the watcher's clone flow, so no reconciler may run.
var errRetentionMissing = errors.New("memory search registry: retention marker missing")

// ErrRegistryClosed is returned by getOrCreate when the registry has been
// permanently closed via closeAll. Callers treat it as "no reconciler
// available" and degrade gracefully. The state is terminal: a closed registry
// never reopens.
var ErrRegistryClosed = errors.New("memory search registry: closed")

// ErrAgentRevoked is returned by getOrCreate when the agent's epoch has been
// bumped by closeAgent between the entry gate and the insert gate — i.e. the
// binding was revoked or rebound while Start was in flight. The freshly-Started
// reconciler is closed (releasing the per-binding owner lock) and the entry
// never lands in the map.
var ErrAgentRevoked = errors.New("memory search registry: agent revoked during start")
