package workflows

import (
	"context"
	"fmt"
	"sort"

	"github.com/samber/lo"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"sigs.k8s.io/yaml"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/gitops"
	"github.com/tidbcloud/infra-provider/pkg/namer"
)

type scaleStep struct {
	step     cloudv1.WorkflowStep
	priority int
}

// scaling order. Read more on https://pingcap.feishu.cn/wiki/ZhFvwtiMsiTMQdkKSBOcHJILnFb#share-CSShd3rD3o2ZfQxuisNcsftJnBg
// Note: the priority number doesn't matter as long as the order is maintained
// 1. storage changes
// 2. scaling up
// 3. scaling out
// 4. scaling in
// 5. scaling down
const (
	scaleDelBalanceRegionSchedulerPriority = iota
	scaleStoragePriority
	scaleAddBalanceRegionSchedulerPriority
	scaleUpPriority
	scaleOutPriority
	scaleInPriority
	scaleDownPriority
)

// Step group names for workflow steps
const (
	stepGroupInstanceTypeOrStorageChange = "InstanceTypeOrStorageChangeGroup"
	stepGroupReplicaChange               = "ReplicaChangeGroup"
	defaultStepGroupName                 = ""
)

func NewScaleCluster(ctx context.Context, instanceTypeProvider gitops.InstanceTypeProvider, cluster *cloudv1.Cluster, changedComp cloudv1.ClusterComponentSpecWrapped, desiredConfigRevisionName, nonce string, disabledExistingGKEHeterogeneous bool) (*cloudv1.Workflow, error) {
	originalComp, err := cluster.GetComponentSpecInGroup(changedComp.GetGroupId(), changedComp.GetComponentType())
	if err != nil { // should not happen
		return nil, fmt.Errorf("failed to get original component spec in group: %w", err)
	}

	prefix := fmt.Sprintf("%s-%s", consts.WorkflowTypeScaleCluster, string(changedComp.GetComponentType()))
	name := namer.NewNameHashOnlyExceed(nonce, prefix, "", 63)
	// Serialize the changed component spec to YAML for the action inputs
	componentSpecBytes, err := yaml.Marshal(changedComp.GetClusterComponent())
	if err != nil {
		return nil, fmt.Errorf("failed to marshal desired component spec: %s", err)
	}
	currentComponentSpec := lo.Must(cluster.GetComponentSpecInGroup(changedComp.GetGroupId(), changedComp.GetComponentType()))
	currentCompSpecBytes, err := yaml.Marshal(currentComponentSpec.GetClusterComponent())
	if err != nil {
		return nil, fmt.Errorf("failed to marshal current component spec: %s", err)
	}
	steps, err := buildScaleClusterSteps(ctx, instanceTypeProvider, cluster, originalComp, changedComp, desiredConfigRevisionName)
	if err != nil {
		return nil, fmt.Errorf("failed to build scale cluster steps: %w", err)
	}
	inputs := map[string]string{
		consts.WfInputKeyNamespace:     cluster.Namespace,
		consts.WfInputKeyClusterID:     cluster.Name,
		consts.WfInputKeyGroupID:       changedComp.GetGroupId(),
		consts.WfInputKeyComponentType: string(changedComp.GetComponentType()),

		consts.WfInputKeyScaleClusterOriginalComponentSpec:     string(currentCompSpecBytes),
		consts.WfInputKeyScaleClusterDesiredComponentSpec:      string(componentSpecBytes),
		consts.WfInputKeyScaleClusterCurrentConfigRevisionName: cluster.Spec.ConfigRevisionRef.Name,
		consts.WfInputKeyScaleClusterDesiredConfigRevisionName: desiredConfigRevisionName,
	}
	if disabledExistingGKEHeterogeneous {
		inputs[consts.WfInputKeyScaleClusterDisableGKEHeterogeneous] = "true"
	}
	wf := &cloudv1.Workflow{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: cluster.Namespace,
			Name:      name,
			Labels:    buildLabels(consts.WorkflowTypeScaleCluster, cluster, changedComp.GetGroupId()),
		},
		Spec: cloudv1.WorkflowSpec{
			Type:   consts.WorkflowTypeScaleCluster,
			TTL:    defaultTTL,
			Inputs: inputs,
			Steps:  steps,
		},
	}

	return wf, nil
}

func buildScaleClusterSteps(ctx context.Context, instanceTypeProvider gitops.InstanceTypeProvider, cluster *cloudv1.Cluster, originalComp, changedComp cloudv1.ClusterComponentSpecWrapped, desiredConfigRevisionName string) (cloudv1.WorkflowSteps, error) {
	steps, err := buildScaleSteps(ctx, cluster, instanceTypeProvider, originalComp, changedComp, desiredConfigRevisionName)
	if err != nil {
		return nil, err
	}

	steps = append(steps, cloudv1.WorkflowStep{
		Name:      "Finalize",
		StepGroup: defaultStepGroupName,
		Uses:      consts.WorkflowActionMutateAnnotations,
		Inputs: cloudv1.InputsWrapper{
			Object: map[string]interface{}{
				"namespace": cluster.Namespace,
				"name":      cluster.Name,
				"type":      "cluster",
				"annotations": map[string]string{
					consts.AnnoKeyScaleClusterWorkflow: "-",
				},
			},
		},
	})

	return steps, nil
}

func buildScaleSteps(ctx context.Context, cluster *cloudv1.Cluster, instanceTypeProvider gitops.InstanceTypeProvider, originalComp, changedComp cloudv1.ClusterComponentSpecWrapped, desiredConfigRevisionName string) (cloudv1.WorkflowSteps, error) {
	// Handle the case when replicas change from/to zero
	if isReplicaZeroTransition(originalComp.GetReplicas(), changedComp.GetReplicas()) {
		return cloudv1.WorkflowSteps{
			{
				Name:      fmt.Sprintf("Scale %s replicas from %d to %d", changedComp.GetComponentType(), originalComp.GetReplicas(), changedComp.GetReplicas()),
				Uses:      consts.WorkflowActionScaleCluster,
				StepGroup: defaultStepGroupName,
				Inputs: cloudv1.InputsWrapper{
					Object: buildScaleClusterActionInputs([]string{consts.WfScaleClusterScopeReplicas, consts.WfScaleClusterScopeInstanceType, consts.WfScaleClusterScopeStorage, consts.WfScaleClusterScopeConfigRevision}),
				},
			},
		}, nil
	}

	groups, err := buildScaleStepsForInstanceTypeChanged(ctx, instanceTypeProvider, cluster, originalComp, changedComp, desiredConfigRevisionName)
	if err != nil {
		return nil, err
	}
	groups = append(groups, buildScaleStepsForStorageChange(originalComp, changedComp)...)
	groups = append(groups, buildScaleStepsForReplicasChanged(originalComp, changedComp)...)

	sort.SliceStable(groups, func(i, j int) bool {
		return groups[i].priority < groups[j].priority
	})
	return lo.Map(groups, func(each scaleStep, _ int) cloudv1.WorkflowStep {
		return each.step
	}), nil
}

func buildScaleStepsForInstanceTypeChanged(ctx context.Context, instanceTypeProvider gitops.InstanceTypeProvider, cluster *cloudv1.Cluster, originalComp, changedComp cloudv1.ClusterComponentSpecWrapped, desiredConfigRevisionName string) ([]scaleStep, error) {
	instanceTypeChanged := originalComp.GetInstanceType() != changedComp.GetInstanceType()
	configRevisionChanged := desiredConfigRevisionName != cluster.Spec.ConfigRevisionRef.Name
	if !instanceTypeChanged {
		if configRevisionChanged {
			return nil, fmt.Errorf("config revision changed but instance type did not change, new config revision: %s", desiredConfigRevisionName)
		}
		return nil, nil
	}

	instanceTypeScaleUp, err := isInstanceTypeScaleUp(ctx, instanceTypeProvider, originalComp.GetInstanceType(), changedComp.GetInstanceType())
	if err != nil {
		return nil, err
	}

	isScalingUp := instanceTypeChanged && instanceTypeScaleUp
	priority := scaleDownPriority
	tp := "down"
	if isScalingUp {
		priority = scaleUpPriority
		tp = "up"
	}

	scopes := []string{consts.WfScaleClusterScopeInstanceType}
	if configRevisionChanged {
		scopes = append(scopes, consts.WfScaleClusterScopeConfigRevision)
	}

	var groups []scaleStep
	groups = append(groups, scaleStep{
		step: cloudv1.WorkflowStep{
			Name:      fmt.Sprintf("Scale %s %s instance type", tp, changedComp.GetComponentType()),
			StepGroup: stepGroupInstanceTypeOrStorageChange,
			Uses:      consts.WorkflowActionScaleCluster,
			Inputs: cloudv1.InputsWrapper{
				Object: buildScaleClusterActionInputs(scopes),
			},
		},
		priority: priority,
	})

	return groups, nil
}

func buildScaleStepsForReplicasChanged(originalComp, changedComp cloudv1.ClusterComponentSpecWrapped) []scaleStep {
	replicasChanged := originalComp.GetReplicas() != changedComp.GetReplicas()
	if !replicasChanged {
		return nil
	}

	isScalingIn := replicasChanged && changedComp.GetReplicas() < originalComp.GetReplicas()

	var groups []scaleStep
	priority := scaleOutPriority
	tp := "out"
	if isScalingIn {
		priority = scaleInPriority
		tp = "in"
	}
	groups = append(groups, scaleStep{
		step: cloudv1.WorkflowStep{
			Name:      fmt.Sprintf("Scale %s %s replicas", tp, changedComp.GetComponentType()),
			StepGroup: stepGroupReplicaChange,
			Uses:      consts.WorkflowActionScaleCluster,
			Inputs: cloudv1.InputsWrapper{
				Object: buildScaleClusterActionInputs([]string{consts.WfScaleClusterScopeReplicas}),
			},
		},
		priority: priority,
	})
	return groups
}

func buildScaleStepsForStorageChange(originalComp, changedComp cloudv1.ClusterComponentSpecWrapped) []scaleStep {
	storageChanged := hasStorageChanged(originalComp, changedComp)
	if !storageChanged {
		return nil
	}

	compType := changedComp.GetComponentType()
	storageSizeChanged := storageSizeChanged(originalComp, changedComp)

	var groups []scaleStep
	// Remove `balance-region-scheduler` before expand storage size for TiKV and TiFlash. Refer to https://tidb.atlassian.net/browse/FRM-2706
	if storageSizeChanged && lo.Contains([]cloudv1.ComponentType{cloudv1.ComponentTiKV, cloudv1.ComponentTiFlash}, compType) {
		groups = append(groups,
			scaleStep{
				step: cloudv1.WorkflowStep{
					Name:      "Remove scheduler balance-region-scheduler",
					StepGroup: stepGroupInstanceTypeOrStorageChange,
					Uses:      consts.WfActionTogglePDScheduler,
					Inputs: cloudv1.InputsWrapper{
						Object: map[string]interface{}{
							"namespace":   fmt.Sprintf("$wf.spec.inputs.%s", consts.WfInputKeyNamespace),
							"clusterName": fmt.Sprintf("$wf.spec.inputs.%s", consts.WfInputKeyClusterID),
							"scheduler":   consts.PDSchedulerBalanceRegion,
							"op":          consts.WfInputKeyTogglePDSchedulerOpRemove,
						},
					},
				},
				priority: scaleDelBalanceRegionSchedulerPriority,
			},
			scaleStep{
				step: cloudv1.WorkflowStep{
					Name:      "Add scheduler balance-region-scheduler",
					StepGroup: stepGroupInstanceTypeOrStorageChange,
					Uses:      consts.WfActionTogglePDScheduler,
					Inputs: cloudv1.InputsWrapper{
						Object: map[string]interface{}{
							"namespace":   fmt.Sprintf("$wf.spec.inputs.%s", consts.WfInputKeyNamespace),
							"clusterName": fmt.Sprintf("$wf.spec.inputs.%s", consts.WfInputKeyClusterID),
							"scheduler":   consts.PDSchedulerBalanceRegion,
							"op":          consts.WfInputKeyTogglePDSchedulerOpAdd,
						},
					},
				},
				priority: scaleAddBalanceRegionSchedulerPriority,
			},
		)
	}

	groups = append(groups, scaleStep{
		step: cloudv1.WorkflowStep{
			Name:      fmt.Sprintf("Scale %s storage", changedComp.GetComponentType()),
			StepGroup: stepGroupInstanceTypeOrStorageChange,
			Uses:      consts.WorkflowActionScaleCluster,
			Inputs: cloudv1.InputsWrapper{
				Object: buildScaleClusterActionInputs([]string{consts.WfScaleClusterScopeStorage}),
			},
		},
		priority: scaleStoragePriority,
	})

	return groups
}

// workflowInputRef returns a workflow input reference string
func workflowInputRef(key string) string {
	return fmt.Sprintf("$wf.spec.inputs.%s", key)
}

func buildScaleClusterActionInputs(scopes []string) map[string]interface{} {
	return map[string]interface{}{
		"namespace":                 workflowInputRef(consts.WfInputKeyNamespace),
		"clusterId":                 workflowInputRef(consts.WfInputKeyClusterID),
		"groupId":                   workflowInputRef(consts.WfInputKeyGroupID),
		"componentType":             workflowInputRef(consts.WfInputKeyComponentType),
		"componentSpec":             workflowInputRef(consts.WfInputKeyScaleClusterDesiredComponentSpec),
		"originalComponentSpec":     workflowInputRef(consts.WfInputKeyScaleClusterOriginalComponentSpec),
		"desiredConfigRevisionName": workflowInputRef(consts.WfInputKeyScaleClusterDesiredConfigRevisionName),
		"scopes":                    scopes,
	}
}

// hasStorageChanged checks if storage-related properties have changed
func hasStorageChanged(originalComp, changedComp cloudv1.ClusterComponentSpecWrapped) bool {
	if storageSizeChanged(originalComp, changedComp) {
		return true
	}

	originalClassName, _ := originalComp.GetStorageClassName()
	changedClassName, _ := changedComp.GetStorageClassName()
	if originalClassName != changedClassName {
		return true
	}

	if originalComp.GetUseLocalStorage() != changedComp.GetUseLocalStorage() {
		return true
	}

	if isAdditionalStorageChanged(originalComp, changedComp) {
		return true
	}

	return false
}

func isAdditionalStorageChanged(originalComp, changedComp cloudv1.ClusterComponentSpecWrapped) bool {
	origignalKeyToStorage := make(map[string]cloudv1.ComponentStorageVolume)
	for _, storage := range originalComp.GetAdditionalStorages() {
		origignalKeyToStorage[storage.Name] = storage
	}
	changedKeyToStorage := make(map[string]cloudv1.ComponentStorageVolume)
	for _, storage := range changedComp.GetAdditionalStorages() {
		changedKeyToStorage[storage.Name] = storage
	}

	if len(origignalKeyToStorage) != len(changedKeyToStorage) {
		return true
	}
	for name, originalStorage := range origignalKeyToStorage {
		changedStorage, ok := changedKeyToStorage[name]
		if !ok {
			return true
		}
		if originalStorage.StorageSize != changedStorage.StorageSize {
			return true
		}
		if originalStorage.StorageClassName != changedStorage.StorageClassName {
			return true
		}
	}

	return false
}

func storageSizeChanged(originalComp, changedComp cloudv1.ClusterComponentSpecWrapped) bool {
	originalSize, originalHasSize := originalComp.GetStorageSize()
	changedSize, changedHasSize := changedComp.GetStorageSize()
	if originalHasSize && changedHasSize {
		if !originalSize.Equal(*changedSize) {
			return true
		}
	} else if originalHasSize != changedHasSize {
		return true
	}
	return false
}

func isReplicaZeroTransition(originalReplicas, changedReplicas int) bool {
	return (originalReplicas == 0 && changedReplicas != 0) || (originalReplicas != 0 && changedReplicas == 0)
}

func isInstanceTypeScaleUp(ctx context.Context, instanceTypeProvider gitops.InstanceTypeProvider, oldInstanceType, newInstanceType string) (bool, error) {
	oldObj, err := instanceTypeProvider.FindInstanceType(oldInstanceType)
	if err != nil {
		return false, fmt.Errorf("find instance type: %w", err)
	}
	newObj, err := instanceTypeProvider.FindInstanceType(newInstanceType)
	if err != nil {
		return false, fmt.Errorf("find instance type: %w", err)
	}
	return cloudv1.IsScalingUp(oldObj, newObj), nil
}