// Package kb resolves per-agent MeshX KB repository bindings and
// prepares them for Desktop's managed KB worktree sync.
package kb

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

	gittransport "github.com/go-git/go-git/v5/plumbing/transport/http"
	"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"
)

type ServerConnectionResolver interface {
	ResolveServerForAgent(ctx context.Context, agentID string) (asset.ConnectionInfo, bool)
}

type KnowledgeBaseAccess struct {
	KnowledgeBaseID          string
	KnowledgeBaseName        string
	KnowledgeBaseDescription string
	SourceID                 string
	SourceName               string
	SourceDescription        string
	SourceType               string
	RepoRole                 string
	Branch                   string
	CloneURL                 string
	RepoID                   string
	AccessProfile            string
	Writable                 bool
	Auth                     gittransport.AuthMethod
}

type LookupState = dispatcher.SkillLookupState

const (
	LookupUnknown = dispatcher.SkillLookupUnknown
	LookupUnbound = dispatcher.SkillLookupUnbound
	LookupBound   = dispatcher.SkillLookupBound
)

type Provider struct {
	resolver   ServerConnectionResolver
	httpClient *http.Client
}

func New(resolver ServerConnectionResolver) *Provider {
	return &Provider{resolver: resolver, httpClient: &http.Client{Timeout: 30 * time.Second}}
}

type batchResponse struct {
	Agents           map[string][]*accessibleKnowledgeBase `json:"agents"`
	SyncToken        string                                `json:"syncToken,omitempty"`
	SyncTokenExpires string                                `json:"syncTokenExpires,omitempty"`
}

type accessibleKnowledgeBase struct {
	ID          string                `json:"id"`
	Name        string                `json:"name"`
	Description string                `json:"description,omitempty"`
	Sources     []*accessibleKBSource `json:"sources,omitempty"`
}

type accessibleKBSource struct {
	ID           string                    `json:"id"`
	Name         string                    `json:"name"`
	Description  string                    `json:"description,omitempty"`
	SourceType   string                    `json:"sourceType,omitempty"`
	Repositories []*accessibleKBRepository `json:"repositories,omitempty"`
}

type accessibleKBRepository struct {
	RepoRole      string         `json:"repoRole"`
	Repository    *repositoryRef `json:"repository,omitempty"`
	Branch        string         `json:"branch"`
	AccessProfile string         `json:"accessProfile"`
	AccessStatus  string         `json:"accessStatus"`
	Writable      bool           `json:"writable,omitempty"`
}

type repositoryRef struct {
	RepoID       string `json:"repoId,omitempty"`
	RepoFullName string `json:"repoFullName"`
}

func (p *Provider) Lookup(ctx context.Context, agentID string) ([]KnowledgeBaseAccess, LookupState, error) {
	agentID = strings.TrimSpace(agentID)
	if agentID == "" {
		return nil, LookupUnknown, nil
	}
	lookup, err := p.lookupBatch(ctx, agentID)
	if err != nil {
		slog.WarnContext(ctx, "kbprovider: batch access request failed", "agent_id", agentID, "error", err)
		return nil, LookupUnknown, err
	}
	rawBases, ok := lookup.response.Agents[agentID]
	if !ok {
		return nil, LookupUnknown, fmt.Errorf("kbprovider: batch response omitted requested agent %s", agentID)
	}
	access := flattenKnowledgeBaseAccess(lookup.serverURL, lookup.workspaceID, rawBases)
	if len(access) == 0 {
		slog.InfoContext(ctx, "kbprovider: agent has no accessible KB repositories", "agent_id", agentID)
		return nil, LookupUnbound, nil
	}
	lease, err := gitcred.ParseTokenLease(lookup.response.SyncToken, lookup.response.SyncTokenExpires)
	if err != nil {
		return nil, LookupUnknown, fmt.Errorf("kbprovider: sync lease for agent %s: %w", agentID, err)
	}
	auth := gitcred.NewRefreshableLeaseAuth(lease, func(refreshCtx context.Context) (gitcred.TokenLease, error) {
		return p.refreshSyncLease(refreshCtx, agentID)
	})
	for i := range access {
		access[i].Auth = auth
	}
	return access, LookupBound, nil
}

type batchLookup struct {
	serverURL   string
	workspaceID string
	endpoint    string
	response    batchResponse
}

func (p *Provider) lookupBatch(ctx context.Context, agentID string) (batchLookup, error) {
	var batchResp batchResponse
	conn, ok := p.resolver.ResolveServerForAgent(ctx, agentID)
	if !ok {
		return batchLookup{}, fmt.Errorf("kbprovider: resolve channel for agent %s: unavailable", agentID)
	}
	endpoint, err := asset.PostAgentsBatchAccess(ctx, p.httpClient, conn, "/api/v1/kb/agents/access:batch", []string{agentID}, &batchResp)
	if err != nil {
		return batchLookup{}, fmt.Errorf("kbprovider: %w", err)
	}
	return batchLookup{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.lookupBatch(ctx, strings.TrimSpace(agentID))
	if err != nil {
		return gitcred.TokenLease{}, err
	}
	if len(lookup.response.Agents[agentID]) == 0 {
		return gitcred.TokenLease{}, fmt.Errorf("kbprovider: KB binding for agent %s is no longer bound", agentID)
	}
	return gitcred.ParseTokenLease(lookup.response.SyncToken, lookup.response.SyncTokenExpires)
}

func flattenKnowledgeBaseAccess(serverURL, workspaceID string, bases []*accessibleKnowledgeBase) []KnowledgeBaseAccess {
	var out []KnowledgeBaseAccess
	for _, base := range bases {
		if base == nil {
			continue
		}
		for _, source := range base.Sources {
			if source == nil {
				continue
			}
			for _, repo := range source.Repositories {
				if repo == nil {
					continue
				}
				cloneURL := composeCloneURL(serverURL, workspaceID, repo.Repository)
				branch := strings.TrimSpace(repo.Branch)
				repoID := strings.TrimSpace(repoIDOf(repo.Repository))
				if cloneURL == "" || branch == "" || repoID == "" {
					continue
				}
				out = append(out, KnowledgeBaseAccess{
					KnowledgeBaseID:          strings.TrimSpace(base.ID),
					KnowledgeBaseName:        strings.TrimSpace(base.Name),
					KnowledgeBaseDescription: strings.TrimSpace(base.Description),
					SourceID:                 strings.TrimSpace(source.ID),
					SourceName:               strings.TrimSpace(source.Name),
					SourceDescription:        strings.TrimSpace(source.Description),
					SourceType:               strings.TrimSpace(source.SourceType),
					RepoRole:                 strings.TrimSpace(repo.RepoRole),
					Branch:                   branch,
					CloneURL:                 cloneURL,
					RepoID:                   repoID,
					AccessProfile:            strings.TrimSpace(repo.AccessProfile),
					Writable:                 repo.Writable,
				})
			}
		}
	}
	return out
}

func composeCloneURL(serverURL, workspaceID string, repo *repositoryRef) string {
	if repo == nil || strings.TrimSpace(repo.RepoFullName) == "" {
		return ""
	}
	return asset.ComposeCloneURL(serverURL, workspaceID, repo.RepoFullName)
}

func repoIDOf(repo *repositoryRef) string {
	if repo == nil {
		return ""
	}
	return repo.RepoID
}

type AgentServerConnectionResolver struct {
	serverConnections ServerConnectionAccessor
}

type ServerConnectionAccessor interface {
	ServerConnection(id string) (types.ServerConnection, bool)
}

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, "kbprovider: 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(), TemplateLocale: profile.EffectiveLanguage(profile.Default().Config())}, true
}