// Package memory implements the production MemoryAccessProvider
// that resolves per-agent memory bindings from the AgentTeam server's batch access
// endpoint.
//
// The provider resolves each agent's owning AgentTeam server runtime to obtain the
// correct server URL and auth token, so agents on different AgentTeam servers
// are routed to the right endpoint.
package memory

import (
	"context"
	"fmt"
	"log/slog"
	"net/http"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/tidbcloud/AgentTeam/desktop/internal/dispatcher"
	"github.com/tidbcloud/AgentTeam/desktop/internal/profile"
	"github.com/tidbcloud/AgentTeam/desktop/internal/provisioning/asset"
	"github.com/tidbcloud/AgentTeam/desktop/internal/provisioning/gitcred"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

// ServerConnectionResolver resolves an agent's owning AgentTeam server runtime connection info.
// The dispatcher provides this so the provider can route each agent to the
// correct AgentTeam server.
type ServerConnectionResolver interface {
	// ResolveServerForAgent returns the connection info for the AgentTeam server runtime
	// that owns the given agent, including a freshly-cached user session token.
	// Returns false when the agent has no AgentTeam server runtime or the channel is not running.
	ResolveServerForAgent(ctx context.Context, agentID string) (asset.ConnectionInfo, bool)
}

type resolvedConnectionGroup struct {
	conn asset.ConnectionInfo
	ids  []string
	err  error
}

// serverConnectionBatchResolver is implemented by the production resolver so
// catalog refresh can resolve one token per owning server and publish healthy
// groups without waiting for a slow unrelated connection.
type serverConnectionBatchResolver interface {
	ResolveServerGroups(context.Context, []string) <-chan resolvedConnectionGroup
}

const (
	serverResolutionConcurrency = 8
	serverGroupTimeout          = 30 * time.Second
)

// Provider resolves memory bindings for agents by calling the AgentTeam server's
// batch access endpoint. It implements dispatcher.MemoryAccessProvider.
type Provider struct {
	resolver   ServerConnectionResolver
	httpClient *http.Client
}

// New creates a Provider that resolves each agent's server runtime via resolver.
func New(resolver ServerConnectionResolver) *Provider {
	return &Provider{
		resolver:   resolver,
		httpClient: &http.Client{Timeout: 30 * time.Second},
	}
}

// batchResponse mirrors the relevant fields of the server-side
// AgentAccessBatchResponse DTO.
type batchResponse struct {
	Agents           map[string]*accessibleMemory `json:"agents"`
	SyncToken        string                       `json:"syncToken,omitempty"`
	SyncTokenExpires string                       `json:"syncTokenExpires,omitempty"`
}

// accessibleMemory mirrors the relevant fields of AgentAccessibleMemory.
type accessibleMemory struct {
	ID                          string         `json:"id"`
	AccessLevel                 string         `json:"accessLevel"`
	Repository                  *repositoryRef `json:"repository,omitempty"`
	Branch                      string         `json:"branch"`
	ImportExistingLocalContent  bool           `json:"importExistingLocalContent,omitempty"`
	ReplaceLocalContentOnSwitch bool           `json:"replaceLocalContentOnSwitch,omitempty"`
	SyncProtocol                string         `json:"syncProtocol"`
	ContentID                   string         `json:"contentId,omitempty"`
	ContentVersionID            string         `json:"contentVersionId,omitempty"`
	ContentGeneration           uint64         `json:"contentGeneration,omitempty"`
	AcknowledgedSequence        uint64         `json:"acknowledgedSequence,omitempty"`
	ActiveReplicaID             string         `json:"activeReplicaId,omitempty"`
}

// repositoryRef mirrors the relevant fields of ResourceRepositoryRef.
type repositoryRef struct {
	RepoOwner    string `json:"repoOwner"`
	RepoName     string `json:"repoName"`
	RepoFullName string `json:"repoFullName"`
}

// Lookup resolves the memory binding for a single agent. It returns Bound
// when a complete binding exists, Unbound only when the server authoritatively
// reports null for the requested agent, and Unknown for transient local/server
// failures so Desktop never detaches a valid clone on ambiguity.
func (p *Provider) Lookup(ctx context.Context, agentID string) (dispatcher.MemoryAccess, dispatcher.MemoryAccessState, error) {
	agentID = strings.TrimSpace(agentID)
	if agentID == "" {
		return dispatcher.MemoryAccess{}, dispatcher.MemoryAccessUnknown, fmt.Errorf("memoryprovider: agent id is required")
	}
	if p == nil || p.resolver == nil {
		return dispatcher.MemoryAccess{}, dispatcher.MemoryAccessUnknown, fmt.Errorf("memoryprovider: provider is not configured")
	}

	lookup, err := p.lookupSingleBatch(ctx, agentID)
	if err != nil {
		slog.WarnContext(ctx, "memoryprovider: batch access request failed",
			"agent_id", agentID, "error", err)
		return dispatcher.MemoryAccess{}, dispatcher.MemoryAccessUnknown, err
	}
	return p.resolveLookup(ctx, agentID, lookup)
}

func (p *Provider) resolveLookup(ctx context.Context, agentID string, lookup batchLookup) (dispatcher.MemoryAccess, dispatcher.MemoryAccessState, error) {
	batchResp := lookup.response
	serverURL := lookup.serverURL
	workspaceID := lookup.workspaceID
	endpoint := lookup.endpoint
	slog.InfoContext(ctx, "memoryprovider: querying batch access",
		"agent_id", agentID, "endpoint", endpoint)

	binding, ok := batchResp.Agents[agentID]
	if !ok {
		return dispatcher.MemoryAccess{}, dispatcher.MemoryAccessUnknown, fmt.Errorf("memoryprovider: batch response omitted requested agent %s", agentID)
	}
	if binding == nil {
		return dispatcher.MemoryAccess{}, dispatcher.MemoryAccessUnbound, nil
	}

	access := dispatcher.MemoryAccess{
		MemoryID:                    binding.ID,
		ImportExistingLocalContent:  binding.ImportExistingLocalContent,
		ReplaceLocalContentOnSwitch: binding.ReplaceLocalContentOnSwitch,
	}
	switch dispatcher.MemorySyncProtocol(strings.TrimSpace(binding.SyncProtocol)) {
	case dispatcher.MemorySyncProtocolLegacyGitV1:
		if binding.Repository == nil || strings.TrimSpace(binding.Branch) == "" {
			return dispatcher.MemoryAccess{}, dispatcher.MemoryAccessUnknown, fmt.Errorf("memoryprovider: incomplete legacy memory binding for agent %s", agentID)
		}
		originURL := composeOriginURL(serverURL, workspaceID, binding.Repository)
		if originURL == "" {
			return dispatcher.MemoryAccess{}, dispatcher.MemoryAccessUnknown, fmt.Errorf("memoryprovider: incomplete repository ref for agent %s", agentID)
		}
		lease, err := gitcred.ParseTokenLease(batchResp.SyncToken, batchResp.SyncTokenExpires)
		if err != nil {
			return dispatcher.MemoryAccess{}, dispatcher.MemoryAccessUnknown, fmt.Errorf("memoryprovider: sync lease for agent %s: %w", agentID, err)
		}
		access.Protocol = dispatcher.MemorySyncProtocolLegacyGitV1
		access.OriginURL = originURL
		access.Branch = strings.TrimSpace(binding.Branch)
		access.Auth = gitcred.NewRefreshableLeaseAuth(lease, func(refreshCtx context.Context) (gitcred.TokenLease, error) {
			return p.refreshSyncLease(refreshCtx, agentID)
		})
		access.Token = lease.Token
	case dispatcher.MemorySyncProtocolNativeContentV1:
		if strings.TrimSpace(binding.ID) == "" || strings.TrimSpace(binding.ContentID) == "" || strings.TrimSpace(binding.ContentVersionID) == "" || binding.ContentGeneration == 0 {
			return dispatcher.MemoryAccess{}, dispatcher.MemoryAccessUnknown, fmt.Errorf("memoryprovider: incomplete native memory binding for agent %s", agentID)
		}
		if binding.Repository != nil || strings.TrimSpace(binding.Branch) != "" {
			return dispatcher.MemoryAccess{}, dispatcher.MemoryAccessUnknown, fmt.Errorf("memoryprovider: native memory binding for agent %s contains legacy repository state", agentID)
		}
		access.Protocol = dispatcher.MemorySyncProtocolNativeContentV1
		access.ContentID = strings.TrimSpace(binding.ContentID)
		access.ContentVersionID = strings.TrimSpace(binding.ContentVersionID)
		access.ContentGeneration = binding.ContentGeneration
		access.AcknowledgedSequence = binding.AcknowledgedSequence
		access.ActiveReplicaID = strings.TrimSpace(binding.ActiveReplicaID)
		access.NativeRemote = newNativeRemote(lookup.conn, p.httpClient)
	default:
		return dispatcher.MemoryAccess{}, dispatcher.MemoryAccessUnknown, fmt.Errorf("memoryprovider: unsupported memory sync protocol %q for agent %s", binding.SyncProtocol, agentID)
	}
	slog.InfoContext(ctx, "memoryprovider: resolved memory binding",
		"agent_id", agentID,
		"memory_id", binding.ID,
		"protocol", access.Protocol,
	)
	return access, dispatcher.MemoryAccessBound, nil
}

type batchLookup struct {
	conn        asset.ConnectionInfo
	serverURL   string
	workspaceID string
	endpoint    string
	response    batchResponse
}

// LookupBatch resolves a daemon/catalog snapshot with one HTTP request per
// connection and <=200-agent chunk. Connection groups run independently so a
// slow or failed server cannot prevent unrelated groups from being queried.
func (p *Provider) LookupBatch(ctx context.Context, agentIDs []string) <-chan map[string]dispatcher.MemoryAccessLookup {
	if p == nil || p.resolver == nil {
		unresolved := make(map[string]dispatcher.MemoryAccessLookup)
		err := fmt.Errorf("memoryprovider: provider is not configured")
		for _, id := range agentIDs {
			if id = strings.TrimSpace(id); id != "" {
				unresolved[id] = dispatcher.MemoryAccessLookup{State: dispatcher.MemoryAccessUnknown, Err: err}
			}
		}
		out := make(chan map[string]dispatcher.MemoryAccessLookup, 1)
		out <- unresolved
		close(out)
		return out
	}
	ids := uniqueAgentIDs(agentIDs)
	groups := p.resolveConnectionGroups(ctx, ids)
	out := make(chan map[string]dispatcher.MemoryAccessLookup, len(ids))
	go p.lookupResolvedGroups(ctx, groups, out)
	return out
}

func uniqueAgentIDs(agentIDs []string) []string {
	seen := make(map[string]struct{}, len(agentIDs))
	ids := make([]string, 0, len(agentIDs))
	for _, id := range agentIDs {
		id = strings.TrimSpace(id)
		if id == "" {
			continue
		}
		if _, ok := seen[id]; ok {
			continue
		}
		seen[id] = struct{}{}
		ids = append(ids, id)
	}
	return ids
}

func (p *Provider) resolveConnectionGroups(ctx context.Context, ids []string) <-chan resolvedConnectionGroup {
	if batch, ok := p.resolver.(serverConnectionBatchResolver); ok {
		return batch.ResolveServerGroups(ctx, ids)
	}
	out := make(chan resolvedConnectionGroup, len(ids))
	go func() {
		defer close(out)
		type resolved struct {
			id   string
			conn asset.ConnectionInfo
			ok   bool
		}
		jobs := make(chan string)
		results := make(chan resolved, len(ids))
		var workers sync.WaitGroup
		for range min(serverResolutionConcurrency, max(1, len(ids))) {
			workers.Add(1)
			go func() {
				defer workers.Done()
				for id := range jobs {
					conn, ok := p.resolver.ResolveServerForAgent(ctx, id)
					results <- resolved{id: id, conn: conn, ok: ok}
				}
			}()
		}
		go func() {
			defer close(jobs)
			for _, id := range ids {
				select {
				case jobs <- id:
				case <-ctx.Done():
					return
				}
			}
		}()
		go func() { workers.Wait(); close(results) }()
		grouped := make(map[string]*resolvedConnectionGroup)
		for result := range results {
			if !result.ok {
				out <- resolvedConnectionGroup{ids: []string{result.id}, err: fmt.Errorf("memoryprovider: resolve channel for agent %s: unavailable", result.id)}
				continue
			}
			key := connectionGroupKey(result.conn)
			group := grouped[key]
			if group == nil {
				group = &resolvedConnectionGroup{conn: result.conn}
				grouped[key] = group
			}
			group.ids = append(group.ids, result.id)
		}
		for _, group := range grouped {
			sort.Strings(group.ids)
			out <- *group
		}
	}()
	return out
}

func connectionGroupKey(conn asset.ConnectionInfo) string {
	return strings.TrimRight(strings.TrimSpace(conn.ServerURL), "/") + "\x00" + strings.TrimSpace(conn.WorkspaceID) + "\x00" +
		strings.TrimSpace(conn.UserToken) + "\x00" + strings.TrimSpace(conn.TemplateLocale) + fmt.Sprintf("\x00%p\x00%p", conn.HTTPClient, conn.RefreshUserToken)
}

func (p *Provider) lookupResolvedGroups(ctx context.Context, groups <-chan resolvedConnectionGroup, out chan<- map[string]dispatcher.MemoryAccessLookup) {
	defer close(out)
	jobs := make(chan resolvedConnectionGroup)
	var workers sync.WaitGroup
	for range serverResolutionConcurrency {
		workers.Add(1)
		go func() {
			defer workers.Done()
			for group := range jobs {
				p.lookupResolvedGroup(ctx, group, out)
			}
		}()
	}
	for group := range groups {
		select {
		case jobs <- group:
		case <-ctx.Done():
			close(jobs)
			workers.Wait()
			return
		}
	}
	close(jobs)
	workers.Wait()
}

func (p *Provider) lookupResolvedGroup(parent context.Context, group resolvedConnectionGroup, out chan<- map[string]dispatcher.MemoryAccessLookup) {
	if len(group.ids) == 0 {
		return
	}
	ctx, cancel := context.WithTimeout(parent, serverGroupTimeout)
	defer cancel()
	for start := 0; start < len(group.ids); start += 200 {
		end := min(start+200, len(group.ids))
		chunk := group.ids[start:end]
		var response batchResponse
		endpoint, err := "", group.err
		if err == nil {
			endpoint, err = asset.PostAgentsBatchAccess(ctx, p.httpClient, group.conn, "/api/v1/memory/agents/access:batch", chunk, &response)
		}
		lookup := batchLookup{
			conn: group.conn, serverURL: strings.TrimRight(strings.TrimSpace(group.conn.ServerURL), "/"),
			workspaceID: strings.TrimSpace(group.conn.WorkspaceID), endpoint: endpoint, response: response,
		}
		snapshot := make(map[string]dispatcher.MemoryAccessLookup, len(chunk))
		for _, id := range chunk {
			if err != nil {
				snapshot[id] = dispatcher.MemoryAccessLookup{State: dispatcher.MemoryAccessUnknown, Err: err}
				continue
			}
			access, state, resolveErr := p.resolveLookup(ctx, id, lookup)
			snapshot[id] = dispatcher.MemoryAccessLookup{Access: access, State: state, Err: resolveErr}
		}
		select {
		case out <- snapshot:
		case <-ctx.Done():
			return
		}
	}
}

func (p *Provider) lookupSingleBatch(ctx context.Context, agentID string) (batchLookup, error) {
	var batchResp batchResponse
	conn, ok := p.resolver.ResolveServerForAgent(ctx, agentID)
	if !ok {
		return batchLookup{}, fmt.Errorf("memoryprovider: resolve channel for agent %s: unavailable", agentID)
	}
	endpoint, err := asset.PostAgentsBatchAccess(ctx, p.httpClient, conn, "/api/v1/memory/agents/access:batch", []string{agentID}, &batchResp)
	if err != nil {
		return batchLookup{}, fmt.Errorf("memoryprovider: %w", err)
	}

	return batchLookup{conn: conn, serverURL: strings.TrimRight(strings.TrimSpace(conn.ServerURL), "/"), workspaceID: strings.TrimSpace(conn.WorkspaceID), endpoint: endpoint, response: batchResp}, nil
}

func (p *Provider) refreshSyncLease(ctx context.Context, agentID string) (gitcred.TokenLease, error) {
	lookup, err := p.lookupSingleBatch(ctx, strings.TrimSpace(agentID))
	if err != nil {
		return gitcred.TokenLease{}, err
	}
	binding, ok := lookup.response.Agents[agentID]
	if !ok {
		return gitcred.TokenLease{}, fmt.Errorf("memoryprovider: refresh response omitted requested agent %s", agentID)
	}
	if binding == nil {
		return gitcred.TokenLease{}, fmt.Errorf("memoryprovider: memory binding for agent %s is no longer bound", agentID)
	}
	if dispatcher.MemorySyncProtocol(strings.TrimSpace(binding.SyncProtocol)) != dispatcher.MemorySyncProtocolLegacyGitV1 {
		return gitcred.TokenLease{}, fmt.Errorf("memoryprovider: memory binding for agent %s no longer uses legacy Git sync", agentID)
	}
	return gitcred.ParseTokenLease(lookup.response.SyncToken, lookup.response.SyncTokenExpires)
}

// composeOriginURL builds the git clone URL from the canonical repository full
// name returned by MeshX.
func composeOriginURL(serverURL, workspaceID string, repo *repositoryRef) string {
	if repo == nil {
		return ""
	}
	fullName := strings.TrimSpace(repo.RepoFullName)
	if fullName == "" {
		return ""
	}
	return asset.ComposeCloneURL(serverURL, workspaceID, fullName)
}

// AgentServerConnectionResolver resolves agent→channel using profile agent config
// (Source.ID) and the running server runtime manager. It implements ServerConnectionResolver.
type AgentServerConnectionResolver struct {
	serverConnections ServerConnectionAccessor
}

// ServerConnectionAccessor exposes the channel manager's lookup.
type ServerConnectionAccessor interface {
	ServerConnection(id string) (types.ServerConnection, bool)
}

// ConnectionInfoProvider is satisfied by *wire.Protocol.
type ConnectionInfoProvider interface {
	ServerURL() string
	WorkspaceID() string
	UserToken(ctx context.Context) (string, error)
	RefreshUserToken(ctx context.Context, staleToken string) (string, error)
}

func NewAgentServerConnectionResolver(serverConnections ServerConnectionAccessor) *AgentServerConnectionResolver {
	return &AgentServerConnectionResolver{serverConnections: serverConnections}
}

func (r *AgentServerConnectionResolver) ResolveServerForAgent(ctx context.Context, agentID string) (asset.ConnectionInfo, bool) {
	ag, ok := profile.Default().Agents().Get(agentID)
	if !ok || ag.Source.Type != "agentteam" || strings.TrimSpace(ag.Source.ID) == "" {
		return asset.ConnectionInfo{}, false
	}
	ch, ok := r.serverConnections.ServerConnection(ag.Source.ID)
	if !ok {
		return asset.ConnectionInfo{}, false
	}
	conn, ok := ch.(ConnectionInfoProvider)
	if !ok {
		return asset.ConnectionInfo{}, false
	}
	token, err := conn.UserToken(ctx)
	if err != nil {
		slog.WarnContext(ctx, "memoryprovider: failed to resolve user token",
			"agent_id", agentID, "error", err)
		return asset.ConnectionInfo{}, false
	}
	return asset.ConnectionInfo{
		ServerURL:        conn.ServerURL(),
		UserToken:        token,
		RefreshUserToken: conn.RefreshUserToken,
		WorkspaceID:      conn.WorkspaceID(),
		HTTPClient:       nil,
	}, true
}

// ResolveServerGroups resolves the catalog path by owning server rather than
// repeating UserToken for every agent. Each server has its own deadline and
// runs under a bounded worker pool, so a slow channel cannot delay snapshots
// from healthy channels.
func (r *AgentServerConnectionResolver) ResolveServerGroups(ctx context.Context, agentIDs []string) <-chan resolvedConnectionGroup {
	type serverGroup struct {
		ids  []string
		conn ConnectionInfoProvider
	}
	out := make(chan resolvedConnectionGroup, len(agentIDs))
	go func() {
		defer close(out)
		groups := make(map[string]*serverGroup)
		for _, agentID := range agentIDs {
			ag, ok := profile.Default().Agents().Get(agentID)
			if !ok || ag.Source.Type != "agentteam" || strings.TrimSpace(ag.Source.ID) == "" {
				out <- resolvedConnectionGroup{ids: []string{agentID}, err: fmt.Errorf("memoryprovider: resolve channel for agent %s: unavailable", agentID)}
				continue
			}
			sourceID := strings.TrimSpace(ag.Source.ID)
			group := groups[sourceID]
			if group == nil {
				ch, ok := r.serverConnections.ServerConnection(sourceID)
				if !ok {
					out <- resolvedConnectionGroup{ids: []string{agentID}, err: fmt.Errorf("memoryprovider: resolve channel for agent %s: unavailable", agentID)}
					continue
				}
				conn, ok := ch.(ConnectionInfoProvider)
				if !ok {
					out <- resolvedConnectionGroup{ids: []string{agentID}, err: fmt.Errorf("memoryprovider: resolve channel for agent %s: incompatible runtime", agentID)}
					continue
				}
				group = &serverGroup{conn: conn}
				groups[sourceID] = group
			}
			group.ids = append(group.ids, agentID)
		}
		jobs := make(chan *serverGroup)
		var workers sync.WaitGroup
		for range min(serverResolutionConcurrency, max(1, len(groups))) {
			workers.Add(1)
			go func() {
				defer workers.Done()
				for group := range jobs {
					groupCtx, cancel := context.WithTimeout(ctx, serverGroupTimeout)
					token, err := group.conn.UserToken(groupCtx)
					if err != nil {
						result := resolvedConnectionGroup{ids: group.ids, err: fmt.Errorf("memoryprovider: resolve user token: %w", err)}
						select {
						case out <- result:
						case <-ctx.Done():
						}
						cancel()
						continue
					}
					result := resolvedConnectionGroup{ids: group.ids, conn: asset.ConnectionInfo{
						ServerURL: group.conn.ServerURL(), UserToken: token,
						RefreshUserToken: group.conn.RefreshUserToken, WorkspaceID: group.conn.WorkspaceID(),
					}}
					select {
					case out <- result:
					case <-ctx.Done():
					}
					cancel()
				}
			}()
		}
		for _, group := range groups {
			sort.Strings(group.ids)
			select {
			case jobs <- group:
			case <-ctx.Done():
				close(jobs)
				workers.Wait()
				return
			}
		}
		close(jobs)
		workers.Wait()
	}()
	return out
}