package dag

import (
	"cmp"
	"context"
	"encoding/json"
	"fmt"
	"slices"

	"github.com/samber/lo"
	"github.com/tidbcloud/fastflow/pkg/entity"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"

	globalcommonv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/common/v1"
	globaldedicatedv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/dedicated/v1"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/service/workflow/action"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/service/workflow/common"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
	"github.com/tidbcloud/tidb-management-service/utilities/clusterevent"
	"github.com/tidbcloud/tidb-management-service/utilities/grpc/metadata"
	"github.com/tidbcloud/tidb-management-service/utilities/workflow"
)

// the order is the same as the upgrade operations' order in TiDB Operator: PD --> TiFlash --> TiKV --> TiDB
var sequentialScalePriority = map[string]int{
	globaldedicatedv1.ComponentType_PD.String():      1,
	globaldedicatedv1.ComponentType_TIFLASH.String(): 2,
	globaldedicatedv1.ComponentType_TIKV.String():    3,
	globaldedicatedv1.ComponentType_TIDB.String():    4,
	globaldedicatedv1.ComponentType_TIPROXY.String(): 5,
}

// scale down PD before scale in TiKV may become a problem if too many TiKV nodes exist, so the order is: TiFlash --> TiKV --> PD --> TiDB
var reverseScalePriority = map[string]int{
	globaldedicatedv1.ComponentType_TIFLASH.String(): 1,
	globaldedicatedv1.ComponentType_TIKV.String():    2,
	globaldedicatedv1.ComponentType_PD.String():      3,
	globaldedicatedv1.ComponentType_TIDB.String():    4,
	globaldedicatedv1.ComponentType_TIPROXY.String(): 5,
}

type ScaleDagTasks struct {
	TenantId                         uint64
	ProjectId                        uint64
	ClusterId                        uint64
	CloudProviderId                  globalcommonv1.CloudProvider_Name
	ComponentsAsConfigOrigin         []*action.ScaleComponentParameter
	ComponentsAsConfigFollower       []*action.ScaleComponentParameter
	ComponentTiProxyAsConfigOrigin   []*action.ScaleComponentParameter
	ComponentTiProxyAsConfigFollower []*action.ScaleComponentParameter

	ScaleComponents  map[globaldedicatedv1.ComponentType][]*action.ScaleComponentParameter
	ScaleOperations  map[string][]string // map[ComponentType][]OperationType
	OperationContext *clusterevent.OperationContext
}

const (
	ScaleDagPrefix                = "global_scale"
	OptionDisableGKEHeterogeneous = "disable_gke_heterogeneous"
)

func CanRunScaleDag(ctx context.Context, wflow workflow.Workflow, clusterId uint64) bool {
	instances, err := wflow.ListRunningDagInstancesByNamePrefix(
		ctx,
		fmt.Sprintf("%s_%s_%d", workflow.DynamicDagPrefix, ScaleDagPrefix, clusterId),
		fmt.Sprintf("%d", clusterId),
	)
	if err != nil {
		tilog.Error(ctx, "list running dag instances failed", zap.Uint64("cluster", clusterId), zap.Error(err))
		return false
	}

	if len(instances) > 0 {
		tilog.Error(ctx, "cluster has running dag instances", zap.Uint64("cluster", clusterId))
		return false
	}

	return true
}

func RunScaleDag(ctx context.Context, wflow workflow.Workflow, param *ScaleDagTasks, force bool) error {
	metadataJson, merr := metadata.MarshalRequestMetaFromCtx(ctx)
	if merr != nil {
		tilog.Warn(ctx, "failed to marshal request meta", zap.Error(merr))
	}

	operationsJson, merr := json.Marshal(param.ScaleOperations)
	if merr != nil {
		tilog.Warn(ctx, "failed to marshal scale operations", zap.Error(merr))
	}
	tasks, err := BuildScaleTopology(param, force)
	if err != nil {
		return err
	}

	dagVars := map[string]string{
		workflow.KeyRequestMeta:         string(metadataJson),
		common.KeyClusterId:             fmt.Sprintf("%d", param.ClusterId),
		workflow.KeyTraceId:             workflow.GetTraceId(ctx),
		common.KeyTargetStatus:          globaldedicatedv1.Cluster_AVAILABLE.String(),
		workflow.KeyComponentOperations: string(operationsJson),
	}
	if param.OperationContext != nil {
		operationContextJson, merr := json.Marshal(param.OperationContext)
		if merr != nil {
			tilog.Warn(ctx, "failed to marshal operation context", zap.Error(merr))
		} else {
			dagVars[common.KeyOperationContext] = string(operationContextJson)
		}
	}

	return wflow.RegisterDynamicDagThenRun(ctx, ScaleDagPrefix, &workflow.CreateDynamicDagOption{
		OrgId:     fmt.Sprintf("%d", param.TenantId),
		ProjectId: fmt.Sprintf("%d", param.ProjectId),
		BizId:     fmt.Sprintf("%d", param.ClusterId),
		DagVars:   dagVars,
		DagTasks:  tasks,
	})
}

func sortComponents(param []*action.ScaleComponentParameter) []*action.ScaleComponentParameter {
	pdParam, hasPd := lo.Find(param, func(item *action.ScaleComponentParameter) bool {
		return item.ComponentType == globaldedicatedv1.ComponentType_PD.String()
	})

	priority := sequentialScalePriority
	if hasPd && pdParam.NeedToScaleInstanceMu != nil && !pdParam.NeedToScaleInstanceMu.ScaleUp {
		priority = reverseScalePriority
	}

	slices.SortFunc(param, func(a, b *action.ScaleComponentParameter) int {
		return cmp.Compare(priority[a.ComponentType], priority[b.ComponentType])
	})

	return param
}

func BuildScaleTopology(param *ScaleDagTasks, force bool) ([]entity.Task, error) {
	topology := make([]entity.Task, 0)

	buildComponentScaleTasks := func(components []*action.ScaleComponentParameter, actionName string) error {
		for _, component := range components {
			scaleTasks, err := toScaleTasks(
				lo.If(len(topology) == 0, "").ElseF(func() string {
					return topology[len(topology)-1].ID
				}),
				actionName,
				component,
				force,
			)
			if err != nil {
				return err
			}
			topology = append(topology, scaleTasks...)
		}
		return nil
	}

	if err := buildComponentScaleTasks(
		sortComponents(param.ComponentsAsConfigOrigin),
		action.ScaleComponentWithConfigOriginActionName,
	); err != nil {
		return nil, err
	}
	if err := buildComponentScaleTasks(
		sortComponents(param.ComponentsAsConfigFollower),
		action.ScaleComponentWithConfigFollowerActionName,
	); err != nil {
		return nil, err
	}

	// TiProxy scaling special cases:
	// In same tidbgroup, clusterserver require create tidb first, then create proxy.
	// So we need to build tiproxy scaling tasks after tidb scaling tasks.
	if err := buildComponentScaleTasks(
		sortComponents(param.ComponentTiProxyAsConfigOrigin),
		action.ScaleComponentWithConfigOriginActionName,
	); err != nil {
		return nil, err
	}
	if err := buildComponentScaleTasks(
		sortComponents(param.ComponentTiProxyAsConfigFollower),
		action.ScaleComponentWithConfigFollowerActionName,
	); err != nil {
		return nil, err
	}

	if len(topology) == 0 {
		return nil, bizerr.New(bizerr.ErrCodeParams, "scale topology is empty: no components to scale")
	}

	statusDependOn := []string{topology[len(topology)-1].ID}
	if force {
		pollTask := entity.Task{
			ID:         fmt.Sprintf("%s_all_components", action.PollClusterUntilSynced),
			ActionName: action.PollClusterUntilSynced,
			Params: map[string]any{
				common.KeyClusterId: fmt.Sprintf("{{%s}}", common.KeyClusterId),
			},
			DependOn: []string{topology[len(topology)-1].ID},
		}
		topology = append(topology, pollTask)
		statusDependOn = []string{pollTask.ID}
	}

	topology = append(topology, entity.Task{
		ID:         action.SetGlobalClusterStatus,
		ActionName: action.SetGlobalClusterStatus,
		Params: map[string]any{
			common.KeyClusterId:    fmt.Sprintf("{{%s}}", common.KeyClusterId),
			common.KeyTargetStatus: fmt.Sprintf("{{%s}}", common.KeyTargetStatus),
			common.KeyEventName:    string(clusterevent.EventTypeScaled),
		},
		DependOn: statusDependOn,
	})

	return topology, nil
}

func toScaleTasks(
	parentId, actionName string,
	param *action.ScaleComponentParameter,
	force bool,
) ([]entity.Task, error) {
	data, err := json.Marshal(param)
	if err != nil {
		return nil, err
	}

	id := ScaleComponentTaskId(param.ComponentType, param.ID)
	parent := parentId
	tasks := make([]entity.Task, 0, 4)
	scaleTask := entity.Task{
		ID:         id,
		Name:       id,
		ActionName: actionName,
		Params: map[string]any{
			common.ScaleTaskKeyRequestBody: string(data),
		},
	}
	if parent != "" {
		scaleTask.DependOn = []string{parent}
	}
	tasks = append(tasks, scaleTask)

	if !force {
		tasks = append(tasks, entity.Task{
			ID:         WaitScaleComponentSyncedTaskId(scaleTask.ID),
			ActionName: action.PollClusterUntilSynced,
			Params: map[string]any{
				common.KeyClusterId: fmt.Sprintf("{{%s}}", common.KeyClusterId),
			},
			DependOn: []string{scaleTask.ID},
		})
	}

	if param.ComponentType == globaldedicatedv1.ComponentType_TIDB.String() && param.NeedToScaleInstanceMu != nil &&
		param.NeedToScaleInstanceMu.ScaleUpFromZero {
		tasks = append(tasks, entity.Task{
			ID:         fmt.Sprintf("create_private_link_for_%s", scaleTask.ID),
			ActionName: action.CreatePrivateLinkServiceWithExtraTidbGroup,
			Params: map[string]any{
				common.ScaleTaskKeyRequestBody: string(data),
			},
			DependOn: []string{tasks[len(tasks)-1].ID},
		})
	}

	return tasks, nil
}

func ScaleComponentTaskId(componentType string, componentId uint64) string {
	return fmt.Sprintf("scale_%s_%d", componentType, componentId)
}

func WaitScaleComponentSyncedTaskId(componentScaleTaskId string) string {
	return fmt.Sprintf("wait_%s", componentScaleTaskId)
}