package dispatcher

import (
	"context"
	"strings"
	"sync"
	"time"

	gittransport "github.com/go-git/go-git/v5/plumbing/transport/http"
	"github.com/tidbcloud/AgentTeam/desktop/internal/nativememory"
)

const (
	memoryAccessGroupConcurrency = 4
	memoryAccessAgentConcurrency = 4
	memoryAccessGroupTimeout     = 60 * time.Second
)

// MemorySyncProtocol is the server-selected owner for one exact Memory binding.
// Desktop must execute only the selected protocol and never retry through the
// other owner after a protocol failure.
type MemorySyncProtocol string

const (
	MemorySyncProtocolLegacyGitV1     MemorySyncProtocol = "legacy_git_v1"
	MemorySyncProtocolNativeContentV1 MemorySyncProtocol = "native_content_v1"
)

// MemoryAccess describes a single agent's server-authoritative Memory binding.
// Legacy bindings carry a Git remote and credential. Native bindings carry an
// exact Content Store projection and never carry AGS transport state.
type MemoryAccess struct {
	// MemoryID is the mem_xxx identifier this clone backs. Required.
	MemoryID string
	// Protocol is selected by the server from persisted exact-binding authority.
	Protocol MemorySyncProtocol
	// OriginURL is the full HTTPS URL of the memory repo as exposed by the
	// AgentTeam gateway (e.g. https://agentteam.example.com/workspace/mem_xxx). The
	// batch access response is expected to return this fully composed.
	OriginURL string
	// Branch is the explicit memory repository branch Desktop should checkout
	// and sync. Required; the dispatcher does not derive a fallback branch
	// from the agent ID.
	Branch string
	// ImportExistingLocalContent controls whether pre-binding local Memory is
	// imported after clone. A binding-scoped pending root created after a clone
	// outage is always imported because it is already the authoritative bound root.
	ImportExistingLocalContent bool
	// ReplaceLocalContentOnSwitch controls what happens when the agent's
	// binding changes from one Memory to another while a local clone already
	// exists. When true, the local clone is backed up and a fresh clone is
	// created from the new binding. When false, the local clone is retained,
	// its origin remote is repointed to the new binding, and the local branch
	// is force-pushed to the new origin.
	ReplaceLocalContentOnSwitch bool
	// Auth supplies the clone/fetch/push transport auth. May be nil for
	// unauthenticated origins (e.g. file:// in tests). Production AgentTeam memory
	// clones use an AGS repository sync lease, not the user's AgentTeam API token.
	Auth gittransport.AuthMethod
	// Token is the initially issued AGS sync lease plaintext. Auth is the
	// managed git credential and may refresh this lease; this field exists
	// for older call sites that still inspect the resolved access value.
	// Empty in tests.
	Token string
	// Native projection fields are required only for native_content_v1.
	ContentID            string
	ContentVersionID     string
	ContentGeneration    uint64
	AcknowledgedSequence uint64
	ActiveReplicaID      string
	NativeRemote         nativememory.Remote
	// nativeRefreshHooks is a package-private staged-refresh failure seam.
	// Production leaves it nil.
	nativeRefreshHooks *nativeMemoryRefreshHooks
}

// MemoryAccessState distinguishes authoritative binding state from transient
// lookup uncertainty. Unknown must never be treated as revocation because that
// would detach a valid local clone during channel/token/server outages.
type MemoryAccessState int

const (
	MemoryAccessUnknown MemoryAccessState = iota
	MemoryAccessUnbound
	MemoryAccessBound
)

// MemoryAccessProvider resolves memory bindings for agents. The dispatcher
// consults this at session start: when a binding exists, the per-agent memory
// clone is created/refreshed and a memorywatcher is attached; otherwise the
// session starts without memory sync and continues normally.
//
// The contract is intentionally read-only: providers do not mutate state and
// must be safe to call concurrently from session-start paths.
type MemoryAccessProvider interface {
	// Lookup returns the memory binding state for agentID. Unknown represents a
	// transient inability to determine state and must not revoke or clear local
	// cache. Unbound is authoritative "no binding". Bound returns access data.
	Lookup(ctx context.Context, agentID string) (MemoryAccess, MemoryAccessState, error)
}

// MemoryAccessLookup is one authoritative or explicitly unknown result from a
// bulk provider snapshot.
type MemoryAccessLookup struct {
	Access MemoryAccess
	State  MemoryAccessState
	Err    error
}

// MemoryAccessBatchProvider is the optional daemon/catalog fast path. Each
// yielded snapshot is complete for one independent connection chunk and may be
// consumed before slower connection groups finish. Session start intentionally
// keeps the narrow single Lookup contract.
type MemoryAccessBatchProvider interface {
	LookupBatch(ctx context.Context, agentIDs []string) <-chan map[string]MemoryAccessLookup
}

// RefreshMemoryAccess re-evaluates watcher state for the requested agents.
// Bulk-capable providers yield independent connection chunks so healthy agents
// can converge without waiting for an unrelated slow server.
func (d *memoryService) RefreshMemoryAccess(ctx context.Context, agentIDs []string) {
	if ctx == nil {
		ctx = context.Background()
	}
	d.memoryAccessMu.RLock()
	provider := d.memoryAccess
	d.memoryAccessMu.RUnlock()
	if batch, ok := provider.(MemoryAccessBatchProvider); ok {
		d.refreshMemoryAccessBatches(ctx, batch.LookupBatch(ctx, agentIDs))
		return
	}
	for _, agentID := range agentIDs {
		agentID = strings.TrimSpace(agentID)
		if agentID == "" {
			continue
		}
		if target := d.refreshAgentMemoryAccessWithSnapshot(ctx, agentID, nil); target != nil {
			d.restartSessionAfterMemoryRefresh(ctx, target)
		}
	}
}

func (d *memoryService) refreshMemoryAccessBatches(ctx context.Context, batches <-chan map[string]MemoryAccessLookup) {
	jobs := make(chan map[string]MemoryAccessLookup, memoryAccessGroupConcurrency)
	var workers sync.WaitGroup
	for range memoryAccessGroupConcurrency {
		workers.Add(1)
		go func() {
			defer workers.Done()
			for batch := range jobs {
				d.refreshMemoryAccessBatch(ctx, batch)
			}
		}()
	}
	for batch := range batches {
		select {
		case jobs <- batch:
		case <-ctx.Done():
			close(jobs)
			workers.Wait()
			return
		}
	}
	close(jobs)
	workers.Wait()
}

func (d *memoryService) refreshMemoryAccessBatch(parent context.Context, batch map[string]MemoryAccessLookup) {
	ctx, cancel := context.WithTimeout(parent, memoryAccessGroupTimeout)
	defer cancel()
	type mutation struct {
		agentID string
		lookup  MemoryAccessLookup
	}
	jobs := make(chan mutation)
	var workers sync.WaitGroup
	for range min(memoryAccessAgentConcurrency, max(1, len(batch))) {
		workers.Add(1)
		go func() {
			defer workers.Done()
			for job := range jobs {
				d.refreshMemoryAccessSnapshot(ctx, job.agentID, job.lookup)
			}
		}()
	}
	for agentID, lookup := range batch {
		select {
		case jobs <- mutation{agentID: agentID, lookup: lookup}:
		case <-ctx.Done():
			close(jobs)
			workers.Wait()
			return
		}
	}
	close(jobs)
	workers.Wait()
}

func (d *memoryService) refreshMemoryAccessSnapshot(ctx context.Context, agentID string, snapshot MemoryAccessLookup) {
	agentID = strings.TrimSpace(agentID)
	if agentID == "" {
		return
	}
	if target := d.refreshAgentMemoryAccessWithSnapshot(ctx, agentID, &snapshot); target != nil {
		d.restartSessionAfterMemoryRefresh(ctx, target)
	}
}

func (d *memoryService) restartSessionAfterMemoryRefresh(ctx context.Context, target *memoryRestartTarget) {
	// Thread-contract instructions are fixed for the lifetime of a live
	// runtime. Retire an unbound process after its first managed Memory attach
	// so the next delivery recompiles against memory/MEMORY.md instead of
	// continuing to write the workspace-root fallback.
	if !d.Dispatcher.interruptSessionProcessForRestart(target.session, target.process) {
		// A newer Process may have won publication while the attach was in
		// flight. Close the migration only if that exact current generation was
		// compiled with managed Memory instructions.
		d.finalizePendingLegacyPreferencesMigration(ctx, target.session, false)
	}
}