package server

import (
	"context"
	"errors"
	"fmt"
	"strings"

	"github.com/tidbcloud/AgentTeam/internal/identity/trusted"
	"github.com/tidbcloud/AgentTeam/internal/service/memory"
	memorystore "github.com/tidbcloud/AgentTeam/internal/store/resource/memory"
	"gorm.io/gorm"
)

const DefaultAgentCloudMemoryBackfillLimit = 1000

type AgentCloudMemoryBackfillOptions struct {
	WorkspaceID  string
	AfterAgentID string
	Limit        int
	DryRun       bool
}

type AgentCloudMemoryBackfillResult struct {
	WorkspaceID       string                          `json:"workspaceId,omitempty"`
	AfterAgentID      string                          `json:"afterAgentId,omitempty"`
	NextAfterAgentID  string                          `json:"nextAfterAgentId,omitempty"`
	Limit             int                             `json:"limit"`
	DryRun            bool                            `json:"dryRun"`
	Scanned           int                             `json:"scanned"`
	AlreadyBound      int                             `json:"alreadyBound"`
	Enabled           int                             `json:"enabled"`
	WouldEnable       int                             `json:"wouldEnable"`
	Failed            int                             `json:"failed"`
	Errors            []AgentCloudMemoryBackfillError `json:"errors,omitempty"`
	ImportLocalMemory bool                            `json:"importLocalMemory"`
}

type AgentCloudMemoryBackfillError struct {
	WorkspaceID string `json:"workspaceId"`
	AgentID     string `json:"agentId"`
	Message     string `json:"message"`
}

type agentCloudMemoryBackfillEnabler interface {
	EnableAgentCloudMemory(ctx context.Context, workspaceID, subject, agentID string, input memory.AgentCloudMemoryEnableMutation) (memory.AgentCloudMemory, error)
}

type agentCloudMemoryBackfillCandidate struct {
	WorkspaceID string `gorm:"column:workspace_id"`
	AgentID     string `gorm:"column:agent_id"`
	AgentName   string `gorm:"column:agent_name"`
	DisplayName string `gorm:"column:display_name"`
	OwnerUserID string `gorm:"column:owner_user_id"`
}

func (s *Server) BackfillAgentCloudMemory(ctx context.Context, opts AgentCloudMemoryBackfillOptions) (AgentCloudMemoryBackfillResult, error) {
	if s == nil {
		return AgentCloudMemoryBackfillResult{}, errors.New("meshx server is not initialized")
	}
	return backfillAgentCloudMemory(ctx, s.agentteamCoreDB, s.meshxDB, s.memoryService, opts)
}

func backfillAgentCloudMemory(ctx context.Context, agentteamCoreDB, meshxDB *gorm.DB, enabler agentCloudMemoryBackfillEnabler, opts AgentCloudMemoryBackfillOptions) (AgentCloudMemoryBackfillResult, error) {
	opts.WorkspaceID = strings.TrimSpace(opts.WorkspaceID)
	opts.AfterAgentID = strings.TrimSpace(opts.AfterAgentID)
	if opts.Limit <= 0 {
		opts.Limit = DefaultAgentCloudMemoryBackfillLimit
	}
	result := AgentCloudMemoryBackfillResult{
		WorkspaceID:       opts.WorkspaceID,
		AfterAgentID:      opts.AfterAgentID,
		Limit:             opts.Limit,
		DryRun:            opts.DryRun,
		ImportLocalMemory: true,
	}
	if agentteamCoreDB == nil {
		return result, errors.New("agentteam core database is not configured")
	}
	if meshxDB == nil {
		return result, errors.New("meshx database is not configured")
	}
	if enabler == nil && !opts.DryRun {
		return result, errors.New("memory service is not configured")
	}

	candidates, err := listAgentCloudMemoryBackfillCandidates(ctx, agentteamCoreDB, opts)
	if err != nil {
		return result, err
	}
	result.Scanned = len(candidates)
	if len(candidates) > 0 {
		result.NextAfterAgentID = candidates[len(candidates)-1].AgentID
	}

	bound, err := listBackfillBoundAgentIDs(ctx, meshxDB, candidateAgentIDs(candidates))
	if err != nil {
		return result, err
	}
	importExisting := true
	for _, candidate := range candidates {
		if _, ok := bound[candidate.AgentID]; ok {
			result.AlreadyBound++
			continue
		}
		if opts.DryRun {
			result.WouldEnable++
			continue
		}
		if err := enableAgentCloudMemoryForBackfill(ctx, enabler, candidate, importExisting); err != nil {
			result.Failed++
			result.Errors = append(result.Errors, AgentCloudMemoryBackfillError{
				WorkspaceID: candidate.WorkspaceID,
				AgentID:     candidate.AgentID,
				Message:     err.Error(),
			})
			continue
		}
		result.Enabled++
	}
	if result.Failed > 0 {
		return result, fmt.Errorf("agent cloud memory backfill failed for %d agent(s)", result.Failed)
	}
	return result, nil
}

func listAgentCloudMemoryBackfillCandidates(ctx context.Context, db *gorm.DB, opts AgentCloudMemoryBackfillOptions) ([]agentCloudMemoryBackfillCandidate, error) {
	conditions := []string{
		"a.status <> 'archived'",
		"mac.status = 'active'",
		"TRIM(COALESCE(mac.created_by_user_id, '')) <> ''",
	}
	args := make([]any, 0, 3)
	if opts.WorkspaceID != "" {
		conditions = append(conditions, "a.workspace_id = ?")
		args = append(args, opts.WorkspaceID)
	}
	if opts.AfterAgentID != "" {
		conditions = append(conditions, "a.id > ?")
		args = append(args, opts.AfterAgentID)
	}
	args = append(args, opts.Limit)

	query := `
SELECT a.workspace_id,
       a.id AS agent_id,
       a.name AS agent_name,
       m.display_name,
       mac.created_by_user_id AS owner_user_id
FROM agents a
JOIN members m ON m.id = a.member_id
JOIN machines mac ON mac.workspace_id = a.workspace_id AND mac.id = a.machine_id
WHERE ` + strings.Join(conditions, "\n  AND ") + `
ORDER BY a.id ASC
LIMIT ?`

	var candidates []agentCloudMemoryBackfillCandidate
	if err := db.WithContext(ctx).Raw(query, args...).Scan(&candidates).Error; err != nil {
		return nil, fmt.Errorf("list agent cloud memory backfill candidates: %w", err)
	}
	for i := range candidates {
		candidates[i].WorkspaceID = strings.TrimSpace(candidates[i].WorkspaceID)
		candidates[i].AgentID = strings.TrimSpace(candidates[i].AgentID)
		candidates[i].AgentName = strings.TrimSpace(candidates[i].AgentName)
		candidates[i].DisplayName = strings.TrimSpace(candidates[i].DisplayName)
		candidates[i].OwnerUserID = strings.TrimSpace(candidates[i].OwnerUserID)
	}
	return candidates, nil
}

func listBackfillBoundAgentIDs(ctx context.Context, db *gorm.DB, agentIDs []string) (map[string]struct{}, error) {
	bound := make(map[string]struct{})
	if len(agentIDs) == 0 {
		return bound, nil
	}
	var rows []struct {
		AgentID string `gorm:"column:agent_id"`
	}
	if err := db.WithContext(ctx).
		Model(&memorystore.Memory{}).
		Select("agent_id").
		Where("agent_id IN ? AND deleted_at IS NULL AND is_active = ?", agentIDs, true).
		Scan(&rows).Error; err != nil {
		return nil, fmt.Errorf("list bound memory agents: %w", err)
	}
	for _, row := range rows {
		if id := strings.TrimSpace(row.AgentID); id != "" {
			bound[id] = struct{}{}
		}
	}
	return bound, nil
}

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

func enableAgentCloudMemoryForBackfill(ctx context.Context, enabler agentCloudMemoryBackfillEnabler, candidate agentCloudMemoryBackfillCandidate, importExisting bool) error {
	identity := trusted.Identity{
		Subject: candidate.OwnerUserID,
		RawClaims: map[string]any{
			"sub":                          candidate.OwnerUserID,
			"agent_ownership_workspace_id": candidate.WorkspaceID,
			"owned_agent_ids":              []any{candidate.AgentID},
			"owned_agents": []any{
				map[string]any{
					"id":           candidate.AgentID,
					"name":         candidate.AgentName,
					"display_name": candidate.DisplayName,
				},
			},
		},
	}
	_, err := enabler.EnableAgentCloudMemory(trusted.WithIdentity(ctx, identity), candidate.WorkspaceID, candidate.OwnerUserID, candidate.AgentID, memory.AgentCloudMemoryEnableMutation{
		ImportExistingLocalContent: &importExisting,
	})
	return err
}