package server

import (
	"context"
	"fmt"
	"strings"

	"github.com/google/uuid"
	"github.com/samber/lo"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/resource"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/yaml"

	"github.com/tidbcloud/cluster-controller/server/internal/util"
	"github.com/tidbcloud/cluster-controller/server/server/cvt"
	"github.com/tidbcloud/cluster-controller/server/server/workflows"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	apiv1 "github.com/tidbcloud/infra-provider/api/tidbcloud/cluster-server/v1"
	apicommon "github.com/tidbcloud/infra-provider/api/tidbcloud/common"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/gitops"
	pkgutil "github.com/tidbcloud/infra-provider/pkg/util"
)

var (
	ComponentsCanScaleFromZero = []cloudv1.ComponentType{cloudv1.ComponentTiFlash, cloudv1.ComponentTiProxy}
)

const scaleOptionDisableGKEHeterogeneous = "disable_gke_heterogeneous"

// getClusterForUpdating first gets the cluster CR by clusterId, then checks if the cluster is synced so that it can be updated.
func (s *ClusterServer) getClusterForUpdating(ctx context.Context, clusterId string) (*cloudv1.Cluster, error) {
	cluster, err := s.clusterReader.Get(ctx, s.KubeClient, clusterId)
	if err != nil {
		if apierrors.IsNotFound(err) {
			return nil, status.Errorf(codes.NotFound, "Cluster %s not found", clusterId)
		}
		return nil, status.Errorf(codes.Internal, "Failed to get cluster %s: %s", clusterId, err)
	}

	if cvt.CalculateClusterPhase(cluster) != apiv1.Cluster_SYNCED {
		return nil, status.Errorf(codes.Unavailable, "Cluster %s is not synced", clusterId)
	}

	return cluster, nil
}

// ScaleCluster scales up/down or in/out one component of a TiDB cluster.
func (s *ClusterServer) ScaleCluster(ctx context.Context, req *apiv1.ScaleClusterRequest) (*apiv1.Cluster, error) {
	logger := tilog.WithContextL(ctx).With(
		zap.String("cluster", req.ClusterId),
		zap.String("component", req.NodeSpec.GetComponentType().String()),
	)

	cluster, err := s.getClusterForUpdating(ctx, req.ClusterId)
	if err != nil {
		logger.Error("failed to get cluster for updating", zap.Error(err))
		return nil, err
	}

	if err := validateScaleClusterRequest(ctx, req, cluster); err != nil {
		return nil, err
	}

	validation, err := validateScaleClusterForApply(logger, cluster, req)
	if err != nil {
		return nil, err
	}
	comp := validation.currentComp

	desiredConfigRevisionName := cluster.Spec.ConfigRevisionRef.Name
	// Update config revision if provide new configuration
	if len(req.ComponentConfigurations) != 0 {
		var version *apiv1.ClusterVersion
		// If scale tiflash from 0, update the config revision with tiflash's image
		if scaleFromZero(cluster, comp, req) && lo.Contains(ComponentsCanScaleFromZero, comp.GetComponentType()) {
			version = &apiv1.ClusterVersion{
				ClusterVersion: "",
				ComponentVersions: []*apiv1.ClusterVersion_ComponentVersion{
					{
						ComponentType: cvt.ComponentType2PB(comp.GetComponentType()),
						ImageRepo:     req.ImageRepo,
						ImageTag:      req.ImageTag,
					},
				},
			}
		}

		desiredConfigRevisionName, err = updateConfigWithoutConfigChangeCR(ctx, s.KubeClient, cluster, req.ComponentConfigurations, version)
		if err != nil {
			return nil, err
		}
		if desiredConfigRevisionName != cluster.Spec.ConfigRevisionRef.Name {
			logger.Info("modify configuration for component")
		} else {
			logger.Warn("provide component configuration, but no change")
		}
	}

	if err := s.applyScaleClusterValidated(ctx, logger, cluster, desiredConfigRevisionName, req, validation); err != nil {
		logger.Error("failed to scale cluster", zap.Error(err))
		return nil, err
	}
	return cvt.Cluster2PB(cluster), nil
}

func (s *ClusterServer) applyScaleCluster(ctx context.Context, logger *zap.Logger, cluster *cloudv1.Cluster, desiredConfigRevisionName string, req *apiv1.ScaleClusterRequest) error {
	validation, err := validateScaleClusterForApply(logger, cluster, req)
	if err != nil {
		return err
	}

	return s.applyScaleClusterValidated(ctx, logger, cluster, desiredConfigRevisionName, req, validation)
}

func (s *ClusterServer) applyScaleClusterValidated(ctx context.Context, logger *zap.Logger, cluster *cloudv1.Cluster, desiredConfigRevisionName string, req *apiv1.ScaleClusterRequest, validation *scaleClusterValidation) error {
	compType := cloudv1.ComponentType(strings.ToLower(req.NodeSpec.GetComponentType().String()))
	cloned := cluster.DeepCopy()
	toggleType, err := applyClusterScaleChangesToCR(ctx, cloned, desiredConfigRevisionName, req, validation.storagePlan)
	if err != nil {
		logger.Error("failed to apply scale changes to cluster", zap.Error(err))
		return err
	}

	// directly scale the cluster CR if this component is being enabled/disabled
	// Note:
	//   1. when enabling/disabling tiproxy, workflow is mandatory and is created/set by `applyClusterScaleChangesToCR`.
	//   2. to unify, we also use workflow to scale tiflash even if it's being enabled/disabled.
	updateWithWorkflow := (toggleType == componentToggleTypeNone || compType == cloudv1.ComponentTiFlash)
	// enable scale status track if the request has the scale-status-track option
	_, scaleStatusTrackOpt := req.Options[consts.ClusterFeatureScaleStatusTrack]
	if scaleStatusTrackOpt {
		cloned.EnableFeature(consts.ClusterFeatureScaleStatusTrack)
	}

	_, scaleCheckNodeOpt := req.Options[consts.ClusterFeatureNodeScalingCheck]
	if scaleCheckNodeOpt {
		cloned.EnableFeature(consts.ClusterFeatureNodeScalingCheck)
	}

	_, disableClusterAutoscalerOpt := req.Options[consts.ClusterFeatureDisableClusterAutoscaler]
	if disableClusterAutoscalerOpt {
		cloned.EnableFeature(consts.ClusterFeatureDisableClusterAutoscaler)
	}

	disabledExistingGKEHeterogeneous := prepareDisableGKEHeterogeneousScale(cloned, validation.disableGKEHeterogeneous)
	if disabledExistingGKEHeterogeneous {
		logger.Info("gke heterogeneous deploy intent disable is requested with scale")
		updateWithWorkflow = true
	}

	if !updateWithWorkflow {
		if err := s.KubeClient.Update(ctx, cloned); err != nil {
			return status.Errorf(codes.Internal, "Failed to modify cluster: %s", err)
		}
		return nil
	}

	// dry run to validate the changes
	if err := s.KubeClient.Update(ctx, cloned, client.DryRunAll); err != nil {
		return status.Errorf(codes.Internal, "Dry run cluster modification: %s", err)
	}
	// encode changes into scale workflow and set into annotation
	changedComp, err := cloned.GetComponentSpecInGroup(scaleRequestTiDBGroupID(req), compType)
	if err != nil {
		return status.Errorf(codes.Internal, "failed to get component spec in group: %s", err)
	}
	if scaleStatusTrackOpt {
		cluster.EnableFeature(consts.ClusterFeatureScaleStatusTrack)
	}
	if scaleCheckNodeOpt {
		cluster.EnableFeature(consts.ClusterFeatureNodeScalingCheck)
	}
	if disableClusterAutoscalerOpt {
		cluster.EnableFeature(consts.ClusterFeatureDisableClusterAutoscaler)
	}
	if err := setScaleClusterWfIntoAnnoIfNeed(ctx, logger, s.instanceTypeProvider, cluster, changedComp, desiredConfigRevisionName, disabledExistingGKEHeterogeneous); err != nil {
		return err
	}
	if err := s.KubeClient.Update(ctx, cluster); err != nil {
		return status.Errorf(codes.Internal, "Failed to modify cluster: %s", err)
	}

	return nil
}

func disableGKEHeterogeneousRequested(req *apiv1.ScaleClusterRequest) bool {
	if req == nil || req.Options == nil {
		return false
	}
	return strings.EqualFold(strings.TrimSpace(req.Options[scaleOptionDisableGKEHeterogeneous]), "true")
}

type scaleClusterValidation struct {
	currentComp             cloudv1.ClusterComponent
	disableGKEHeterogeneous *disableGKEHeterogeneousScaleValidation
	storagePlan             storageScalePlan
}

type disableGKEHeterogeneousScaleValidation struct {
	intentConfig     gkeHeterogeneousIntentConfig
	tidbGroupID      string
	intentConfigured bool
}

func validateScaleClusterForApply(logger *zap.Logger, cluster *cloudv1.Cluster, req *apiv1.ScaleClusterRequest) (*scaleClusterValidation, error) {
	comp, err := getComponentToBeScaled(logger, req, cluster)
	if err != nil {
		logger.Error("failed to get the component to be scaled", zap.Error(err))
		return nil, err
	}

	disableGKEHeterogeneous, err := validateGKEHeterogeneousScaleRequest(cluster, comp, req)
	if err != nil {
		return nil, err
	}
	storagePlan, err := buildStorageScalePlan(comp, req.NodeSpec)
	if err != nil {
		return nil, err
	}
	if err := validateGKEStorageScale(cluster, comp, req.NodeSpec, storagePlan); err != nil {
		return nil, err
	}
	return &scaleClusterValidation{
		currentComp:             comp,
		disableGKEHeterogeneous: disableGKEHeterogeneous,
		storagePlan:             storagePlan,
	}, nil
}

func validateGKEHeterogeneousScaleRequest(cluster *cloudv1.Cluster, currentComp cloudv1.ClusterComponent, req *apiv1.ScaleClusterRequest) (*disableGKEHeterogeneousScaleValidation, error) {
	if err := validateGKEHeterogeneousEmptyInstanceTypeOnlyCountScale(cluster, currentComp, req); err != nil {
		return nil, err
	}
	if err := validateGKEHeterogeneousInstanceTypeMutationNotBlocked(cluster, currentComp, req); err != nil {
		return nil, err
	}
	return validateDisableGKEHeterogeneousScaleRequest(cluster, currentComp, req)
}

func validateGKEHeterogeneousInstanceTypeMutationNotBlocked(cluster *cloudv1.Cluster, currentComp cloudv1.ClusterComponent, req *apiv1.ScaleClusterRequest) error {
	if cluster == nil || currentComp == nil || req == nil || req.NodeSpec == nil {
		return nil
	}
	if cluster.GetProviderType() != cloudv1.ProviderTypeGKE || disableGKEHeterogeneousRequested(req) {
		return nil
	}
	targetInstanceType := strings.TrimSpace(req.NodeSpec.GetInstanceType())
	if targetInstanceType == "" || strings.TrimSpace(currentComp.GetInstanceType()) == targetInstanceType {
		return nil
	}

	if !gkeHeterogeneousScaleIntentConfigured(cluster, currentComp.GetComponentType(), scaleRequestTiDBGroupID(req)) {
		return nil
	}
	return status.Errorf(codes.FailedPrecondition, "cannot change %s instance type while GKE heterogeneous machine deploy is enabled; set options.%s=true to disable heterogeneous deploy first", currentComp.GetComponentType(), scaleOptionDisableGKEHeterogeneous)
}

func validateGKEHeterogeneousEmptyInstanceTypeOnlyCountScale(cluster *cloudv1.Cluster, currentComp cloudv1.ClusterComponent, req *apiv1.ScaleClusterRequest) error {
	if cluster == nil || currentComp == nil || req == nil || req.NodeSpec == nil {
		return nil
	}
	if cluster.GetProviderType() != cloudv1.ProviderTypeGKE || disableGKEHeterogeneousRequested(req) {
		return nil
	}
	if strings.TrimSpace(req.NodeSpec.GetInstanceType()) != "" {
		return nil
	}
	if !gkeHeterogeneousScaleIntentConfigured(cluster, currentComp.GetComponentType(), scaleRequestTiDBGroupID(req)) {
		return nil
	}
	if isCountOnlyScaleRequest(req) {
		return nil
	}
	return status.Error(codes.InvalidArgument, "empty instanceType is only supported for count-only scale while GKE heterogeneous machine deploy is enabled")
}

func shouldApplyScaleRequestInstanceType(cluster *cloudv1.Cluster, currentComp cloudv1.ClusterComponent, req *apiv1.ScaleClusterRequest) bool {
	if cluster == nil || currentComp == nil || req == nil || req.NodeSpec == nil {
		return true
	}
	if strings.TrimSpace(req.NodeSpec.GetInstanceType()) != "" {
		return true
	}
	if cluster.GetProviderType() != cloudv1.ProviderTypeGKE || disableGKEHeterogeneousRequested(req) {
		return true
	}
	return !gkeHeterogeneousScaleIntentConfigured(cluster, currentComp.GetComponentType(), scaleRequestTiDBGroupID(req)) || !isCountOnlyScaleRequest(req)
}

func isCountOnlyScaleRequest(req *apiv1.ScaleClusterRequest) bool {
	return req != nil && req.NodeSpec != nil &&
		len(req.ComponentConfigurations) == 0 &&
		req.NodeSpec.GetStorage() == nil &&
		len(req.NodeSpec.GetAdditionalStorage()) == 0
}

func gkeHeterogeneousScaleIntentConfigured(cluster *cloudv1.Cluster, componentType cloudv1.ComponentType, tidbGroupID string) bool {
	if cluster == nil {
		return false
	}
	switch componentType {
	case cloudv1.ComponentTiKV:
		return cluster.TiKVGKEComputeClassIntentConfigured()
	case cloudv1.ComponentTiDB:
		return cluster.TiDBGKEComputeClassIntentConfigured(tidbGroupID)
	default:
		return false
	}
}

func validateDisableGKEHeterogeneousScaleRequest(cluster *cloudv1.Cluster, currentComp cloudv1.ClusterComponent, req *apiv1.ScaleClusterRequest) (*disableGKEHeterogeneousScaleValidation, error) {
	if !disableGKEHeterogeneousRequested(req) {
		return nil, nil
	}
	if cluster.GetProviderType() != cloudv1.ProviderTypeGKE {
		return nil, nil
	}
	componentType := req.NodeSpec.GetComponentType()
	if componentType != apicommon.ComponentType_TIKV && componentType != apicommon.ComponentType_TIDB {
		return nil, nil
	}

	intentConfig, tidbGroupID, err := resolveGKEHeterogeneousIntent(cluster, componentType, scaleRequestTiDBGroupID(req))
	if err != nil {
		return nil, err
	}
	validation := &disableGKEHeterogeneousScaleValidation{
		intentConfig:     intentConfig,
		tidbGroupID:      tidbGroupID,
		intentConfigured: gkeHeterogeneousIntentConfigured(cluster, intentConfig, tidbGroupID),
	}
	if !validation.intentConfigured {
		return nil, nil
	}
	if strings.TrimSpace(req.NodeSpec.GetInstanceType()) == "" {
		return nil, status.Error(codes.InvalidArgument, "disable_gke_heterogeneous requires target instanceType")
	}
	if strings.TrimSpace(currentComp.GetInstanceType()) == strings.TrimSpace(req.NodeSpec.GetInstanceType()) {
		return nil, status.Error(codes.InvalidArgument, "disable_gke_heterogeneous requires target instanceType to differ from current instanceType")
	}
	return validation, nil
}

func prepareDisableGKEHeterogeneousScale(cloned *cloudv1.Cluster, validation *disableGKEHeterogeneousScaleValidation) bool {
	if validation == nil {
		return false
	}
	clearGKEHeterogeneousIntentInMemory(cloned, validation.intentConfig, validation.tidbGroupID)
	return validation.intentConfigured
}

func clearGKEHeterogeneousIntentInMemory(cluster *cloudv1.Cluster, intentConfig gkeHeterogeneousIntentConfig, tidbGroupID string) {
	if cluster == nil {
		return
	}
	key := gkeHeterogeneousIntentAnnotationKey(cluster, intentConfig, tidbGroupID)
	cluster.MutateAnnotations(func(annotations map[string]string) {
		delete(annotations, key)
	})
}

func applyClusterScaleChangesToCR(ctx context.Context, cluster *cloudv1.Cluster, desiredConfigRevisionName string, req *apiv1.ScaleClusterRequest, storagePlan storageScalePlan) (compToggleType string, _ error) {
	logger := tilog.WithContextL(ctx).With(
		zap.String("cluster", req.ClusterId),
		zap.String("component", req.NodeSpec.GetComponentType().String()),
	)

	comp, err := getComponentToBeScaled(logger, req, cluster)
	if err != nil {
		logger.Error("failed to get the component to be scaled", zap.Error(err))
		return "", err
	}
	compToggleType = getComponentToggleType(comp, req)
	logger.Info("apply scale changes to component", zap.String("toggleType", compToggleType), zap.String("newCRName", desiredConfigRevisionName))

	cluster.Spec.ConfigRevisionRef.Name = desiredConfigRevisionName

	// If the tiproxy is being disabled, we don't change its replicas here.
	// Instead, we will do remove tiproxy operation in the 'disable-tiproxy' workflow.
	if comp.GetComponentType() == cloudv1.ComponentTiProxy && compToggleType == componentToggleTypeDisable {
		logger.Info("tiproxy is being disabled, don't change its replicas in cluster-server")
	} else {
		comp.SetReplicas(int(req.NodeSpec.GetCount()))
	}
	if shouldApplyScaleRequestInstanceType(cluster, comp, req) {
		comp.SetInstanceType(req.NodeSpec.GetInstanceType())
	}
	if storage := req.NodeSpec.GetStorage(); storage != nil {
		comp.SetStorageClassName(storagePlan.primaryStorageClass)
		comp.SetStorageSize(resource.MustParse(fmt.Sprintf("%dGi", storage.GetSizeGi())))
	}
	if comp.GetComponentType() == cloudv1.ComponentTiKV && len(req.NodeSpec.GetAdditionalStorage()) > 0 {
		comp.SetAdditionalStorages(storagePlan.additionalStorages)
	}

	// Annotate the cluster with the tiproxy toggle workflow if needed.
	if err := setTiproxyToggleWfIntoAnnoIfNeed(logger, cluster, compToggleType, scaleRequestTiDBGroupID(req), comp.GetComponentType()); err != nil {
		return "", status.Errorf(codes.Internal, "Failed to set or get tiproxy toggle workflow into annotation: %s", err)
	}

	return compToggleType, nil
}

type storageScalePlan struct {
	primaryStorageClass string
	additionalStorages  []cloudv1.ComponentStorageVolume
}

func buildStorageScalePlan(comp cloudv1.ClusterComponent, nodeSpec *apiv1.Cluster_NodeSpec) (storageScalePlan, error) {
	plan := storageScalePlan{additionalStorages: comp.GetAdditionalStorages()}
	plan.primaryStorageClass, _ = comp.GetStorageClassName()
	if storage := nodeSpec.GetStorage(); storage != nil {
		storageClassName, err := util.BuildStorageClassName(strings.ToLower(nodeSpec.GetComponentType().String()), storage)
		if err != nil {
			return storageScalePlan{}, status.Errorf(codes.InvalidArgument, "invalid storage spec: %s", err)
		}
		plan.primaryStorageClass = storageClassName
	}
	if comp.GetComponentType() != cloudv1.ComponentTiKV || len(nodeSpec.GetAdditionalStorage()) == 0 {
		return plan, nil
	}

	newVolumes := make(map[string]cloudv1.ComponentStorageVolume)
	for _, volume := range getAdditionalStorageVolumesFromComponentReq(nodeSpec) {
		newVolumes[volume.Name] = volume
	}
	desiredVolumes := make([]cloudv1.ComponentStorageVolume, len(plan.additionalStorages))
	missingVolumes := make([]string, 0)
	for i, volume := range plan.additionalStorages {
		newVolume, ok := newVolumes[volume.Name]
		if !ok {
			missingVolumes = append(missingVolumes, volume.Name)
			continue
		}
		desiredVolumes[i] = newVolume
		delete(newVolumes, volume.Name)
	}
	if len(newVolumes) > 0 {
		return storageScalePlan{}, status.Errorf(codes.InvalidArgument, "do not support adding new storage volumes %v", newVolumes)
	}
	if len(missingVolumes) > 0 {
		return storageScalePlan{}, status.Errorf(codes.InvalidArgument, "do not support deleting existing storage volumes %v", missingVolumes)
	}
	plan.additionalStorages = desiredVolumes
	return plan, nil
}

func validateGKEStorageScale(
	cluster *cloudv1.Cluster,
	comp cloudv1.ClusterComponent,
	nodeSpec *apiv1.Cluster_NodeSpec,
	plan storageScalePlan,
) error {
	if cluster.GetProviderType() != cloudv1.ProviderTypeGKE ||
		(nodeSpec.GetStorage() == nil && len(nodeSpec.GetAdditionalStorage()) == 0) {
		return nil
	}

	currentPrimaryStorageClass, hasPrimaryStorage := comp.GetStorageClassName()
	hasCurrentPrimaryVolume := hasPrimaryStorage && comp.GetReplicas() > 0
	vacSupported := cluster.FeatureEnabled(consts.ClusterFeatureUseVolumeAttributesClass) &&
		!cluster.FeatureEnabled(consts.ClusterFeatureUseTiDBOperatorV2)
	primaryPerformanceChanged := hasCurrentPrimaryVolume &&
		pkgutil.GCPHyperdiskPerformanceChanged(currentPrimaryStorageClass, plan.primaryStorageClass)
	if primaryPerformanceChanged && !vacSupported {
		return status.Error(codes.FailedPrecondition, "GKE Hyperdisk performance modification requires VolumeAttributesClass support")
	}
	if primaryPerformanceChanged && comp.GetComponentType() == cloudv1.ComponentTiKV &&
		cluster.Spec.Cluster.TiKV.DedicatedRaftEngineStorage != nil {
		return status.Error(codes.FailedPrecondition,
			"GKE Hyperdisk performance modification is not supported for TiKV with dedicatedRaftEngineStorage")
	}
	if vacSupported && hasCurrentPrimaryVolume && currentPrimaryStorageClass != plan.primaryStorageClass {
		volumeStatus := getComponentPrimaryVolumeStatus(cluster, comp.GetComponentType())
		if volumeStatus == nil {
			return status.Errorf(codes.Unavailable, "GKE %s volume status is not available", comp.GetComponentType())
		}
		if err := validateStorageClassChangeWithActiveVAC(string(comp.GetComponentType()), plan.primaryStorageClass, volumeStatus); err != nil {
			return err
		}
	}
	for i, currentVolume := range comp.GetAdditionalStorages() {
		currentStorageClass := inheritStorageClass(currentVolume.StorageClassName, currentPrimaryStorageClass)
		desiredStorageClass := inheritStorageClass(plan.additionalStorages[i].StorageClassName, plan.primaryStorageClass)
		if !vacSupported && pkgutil.GCPHyperdiskPerformanceChanged(currentStorageClass, desiredStorageClass) {
			return status.Error(codes.FailedPrecondition, "GKE Hyperdisk performance modification requires VolumeAttributesClass support")
		}
		if primaryPerformanceChanged {
			currentVolumeType, _, _ := pkgutil.GetTypeIOPSThroughputFromClassName(currentStorageClass)
			desiredVolumeType, _, _ := pkgutil.GetTypeIOPSThroughputFromClassName(desiredStorageClass)
			if currentVolumeType != consts.GCPVolumeTypeHyperdiskBalanced ||
				desiredVolumeType != consts.GCPVolumeTypeHyperdiskBalanced {
				return status.Errorf(codes.InvalidArgument,
					"TiKV additional volume %q must use the same storage type as the primary volume", currentVolume.Name)
			}
		}
		if vacSupported && currentStorageClass != desiredStorageClass {
			volumeStatus := getTiKVAdditionalVolumeStatus(cluster, currentVolume.Name)
			if volumeStatus == nil {
				return status.Errorf(codes.Unavailable, "GKE TiKV additional volume %q status is not available", currentVolume.Name)
			}
			if err := validateStorageClassChangeWithActiveVAC(
				fmt.Sprintf("tikv additional volume %q", currentVolume.Name), desiredStorageClass, volumeStatus,
			); err != nil {
				return err
			}
		}
	}
	return nil
}

func getComponentPrimaryVolumeStatus(cluster *cloudv1.Cluster, componentType cloudv1.ComponentType) *cloudv1.VolumeStatus {
	if cluster.Status.DB == nil {
		return nil
	}
	switch componentType {
	case cloudv1.ComponentPD:
		if cluster.Status.DB.PD != nil {
			return cluster.Status.DB.PD.Volume
		}
	case cloudv1.ComponentTiKV:
		if cluster.Status.DB.TiKV != nil {
			return cluster.Status.DB.TiKV.Volume
		}
	case cloudv1.ComponentTiFlash, cloudv1.ComponentTiFlashLearner:
		if cluster.Status.DB.TiFlash != nil {
			return cluster.Status.DB.TiFlash.Volume
		}
	}
	return nil
}

func getTiKVAdditionalVolumeStatus(cluster *cloudv1.Cluster, name string) *cloudv1.VolumeStatus {
	if cluster.Status.DB == nil || cluster.Status.DB.TiKV == nil {
		return nil
	}
	for i := range cluster.Status.DB.TiKV.StorageVolumes {
		volumeStatus := &cluster.Status.DB.TiKV.StorageVolumes[i]
		if volumeStatus.Name == name {
			return &volumeStatus.VolumeStatus
		}
	}
	return nil
}

func inheritStorageClass(storageClass, primaryStorageClass string) string {
	if storageClass == "" {
		return primaryStorageClass
	}
	return storageClass
}

func validateStorageClassChangeWithActiveVAC(subject, desiredStorageClass string, volumeStatus *cloudv1.VolumeStatus) error {
	if volumeStatus.CurrentVolumeAttributesClass == "" ||
		volumeStatus.CurrentStorageClass == desiredStorageClass ||
		pkgutil.GCPHyperdiskPerformanceChanged(volumeStatus.CurrentStorageClass, desiredStorageClass) {
		return nil
	}
	return status.Errorf(codes.FailedPrecondition,
		"changing GKE %s StorageClass from %q to %q while VolumeAttributesClass %q is active is not supported",
		subject, volumeStatus.CurrentStorageClass, desiredStorageClass, volumeStatus.CurrentVolumeAttributesClass)
}

func validateScaleClusterRequest(ctx context.Context, req *apiv1.ScaleClusterRequest, cluster *cloudv1.Cluster) error {
	logger := tilog.WithContextL(ctx)
	componentType := req.NodeSpec.GetComponentType()

	if componentType == apicommon.ComponentType_TIFLASH_LEARNER {
		return status.Errorf(codes.InvalidArgument, "please use TIFLASH to scale Tiflash component")
	}

	if lo.Contains([]apicommon.ComponentType{apicommon.ComponentType_TIDB, apicommon.ComponentType_TIPROXY}, componentType) &&
		req.TidbGroupId == "" {
		return status.Errorf(codes.InvalidArgument, "tidbGroupId is required for scaling TiDB or TiProxy")
	}

	if req.TidbGroupId != "" && componentType == apicommon.ComponentType_TIDB && req.NodeSpec.GetCount() == 0 {
		return status.Errorf(codes.InvalidArgument, "can not scale TiDB group to 0")
	}

	configurationLen := len(req.ComponentConfigurations)
	switch componentType {
	// only Tiflash can have multiple configurations, the tiflash and tiflashLeaner
	case apicommon.ComponentType_TIFLASH:
		if configurationLen > 2 {
			return status.Errorf(codes.InvalidArgument, "too many configurations")
		}

		// If scale tiflash from 0
		if scaleFromZero(cluster, &cluster.Spec.Cluster.TiFlash, req) {
			if configurationLen == 0 {
				return status.Errorf(codes.InvalidArgument, "should provide configuration for Tiflash when scale from 0")
			}
			if req.ImageRepo == "" || req.ImageTag == "" {
				return status.Errorf(codes.InvalidArgument, "should provide image repo and tag for Tiflash when scale from 0")
			}
			if req.NodeSpec.Storage == nil {
				return status.Errorf(codes.InvalidArgument, "should provide storage for Tiflash when scale from 0")
			}
			if req.NodeSpec.InstanceType == "" {
				return status.Errorf(codes.InvalidArgument, "should provide instance type for Tiflash when scale from 0")
			}
		}
	case apicommon.ComponentType_TIPROXY:
		// If scale tiproxy from 0
		comp, err := getComponentToBeScaled(logger, req, cluster)
		if err != nil {
			return fmt.Errorf("failed to get the corresponding component to be scaled: %w", err)
		}
		if scaleFromZero(cluster, comp, req) {
			if configurationLen == 0 {
				return status.Errorf(codes.InvalidArgument, "should provide configuration for Tiproxy when scale from 0")
			}
			if req.ImageRepo == "" || req.ImageTag == "" {
				return status.Errorf(codes.InvalidArgument, "should provide image repo and tag for Tiproxy when scale from 0")
			}
			if req.NodeSpec.InstanceType == "" {
				return status.Errorf(codes.InvalidArgument, "should provide instance type for Tiproxy when scale from 0")
			}
		}
	default:
		if configurationLen > 1 {
			return status.Errorf(codes.InvalidArgument, "too many configurations")
		}
	}

	return nil
}

func getComponentToBeScaled(logger *zap.Logger, req *apiv1.ScaleClusterRequest, cluster *cloudv1.Cluster) (cloudv1.ClusterComponent, error) {
	tp := cloudv1.ComponentType(strings.ToLower(req.NodeSpec.GetComponentType().String()))
	tidbGroupID := scaleRequestTiDBGroupID(req)
	if tidbGroupID != "" {
		if !isDefaultLoadBalancer(tidbGroupID, cluster) {
			logger.Info("scaling the extra TiDB group", zap.String("tidbGroupId", tidbGroupID))
			for i := range cluster.Spec.ExtraGroups {
				tidbGroup := &cluster.Spec.ExtraGroups[i]
				if tidbGroup.Name == tidbGroupID {
					return tidbGroup.GetComponent(tp), nil
				}
			}
			return nil, status.Errorf(codes.NotFound, "TiDB group %s not found", tidbGroupID)
		} else {
			logger.Info("scaling the default TiDB group", zap.String("tidbGroupId", tidbGroupID))
		}
	}

	return cluster.GetComponent(tp), nil
}

func scaleRequestTiDBGroupID(req *apiv1.ScaleClusterRequest) string {
	if req == nil || req.NodeSpec == nil {
		return ""
	}
	switch req.NodeSpec.GetComponentType() {
	case apicommon.ComponentType_TIDB, apicommon.ComponentType_TIPROXY:
		return req.TidbGroupId
	default:
		return ""
	}
}

func scaleFromZero(cluster *cloudv1.Cluster, comp cloudv1.ClusterComponent, req *apiv1.ScaleClusterRequest) bool {
	switch compType := comp.GetComponentType(); compType {
	case cloudv1.ComponentTiFlash:
		// TiFlash is only supported in the default elastic group, so we don't need to check the extra groups
		return comp.GetReplicas() == 0 && req.NodeSpec.GetCount() > 0
	case cloudv1.ComponentTiProxy:
		// For TiProxy, we sum up the replicas in the default group and extra groups
		countInCluster := cluster.GetComponent(compType).GetReplicas()
		for _, tidbGroup := range cluster.Spec.ExtraGroups {
			countInCluster += tidbGroup.GetComponent(compType).GetReplicas()
		}
		return countInCluster == 0 && req.NodeSpec.GetCount() > 0
	default:
		return false
	}
}

func setTiproxyToggleWfIntoAnnoIfNeed(logger *zap.Logger, cluster *cloudv1.Cluster, tiproxyToggleType, groupId string, compType cloudv1.ComponentType) error {
	if tiproxyToggleType == componentToggleTypeNone || compType != cloudv1.ComponentTiProxy {
		return nil
	}

	if anno, ok := cluster.Annotations[consts.AnnoKeyTiProxyToggleWorkflow]; ok {
		var wf cloudv1.Workflow
		if err := yaml.Unmarshal([]byte(anno), &wf); err != nil {
			return fmt.Errorf("failed to unmarshal tiproxy toggle workflow: %w", err)
		}
		logger.Info("tiproxy toggle workflow is already running", zap.String("workflow", wf.Name))
		return nil
	}

	suffix := fmt.Sprintf("%s-%s", groupId, strings.ReplaceAll(uuid.New().String(), "-", ""))
	var wf *cloudv1.Workflow
	if tiproxyToggleType == componentToggleTypeEnable {
		wf = workflows.NewEnableTiProxy(cluster, groupId, suffix)
	} else if tiproxyToggleType == componentToggleTypeDisable {
		wf = workflows.NewDisableTiProxy(cluster, groupId, suffix)
	} else { // should not happen
		panic(fmt.Sprintf("invalid tiproxy toggle type: %s", tiproxyToggleType))
	}
	body, err := yaml.Marshal(wf)
	if err != nil {
		return fmt.Errorf("failed to marshal tiproxy toggle workflow: %w", err)
	}
	cluster.UpdateAnnotation(consts.AnnoKeyTiProxyToggleWorkflow, string(body))
	logger.Info("tiproxy toggle workflow is annotated in cluster", zap.String("workflow", wf.Name))
	return nil
}

func setScaleClusterWfIntoAnnoIfNeed(ctx context.Context, logger *zap.Logger, instanceTypeProvider gitops.InstanceTypeProvider, cluster *cloudv1.Cluster, changedComp cloudv1.ClusterComponentSpecWrapped, desiredConfigRevisionName string, disabledExistingGKEHeterogeneous bool) error {
	if anno, ok := cluster.Annotations[consts.AnnoKeyScaleClusterWorkflow]; ok {
		var wf cloudv1.Workflow
		if err := yaml.Unmarshal([]byte(anno), &wf); err != nil {
			return fmt.Errorf("failed to unmarshal scale cluster workflow: %w", err)
		}
		logger.Info("scale cluster workflow is already running", zap.String("workflow", wf.Name))
		return nil
	}

	nonce := strings.ReplaceAll(uuid.New().String(), "-", "")
	wf, err := workflows.NewScaleCluster(ctx, instanceTypeProvider, cluster, changedComp, desiredConfigRevisionName, nonce, disabledExistingGKEHeterogeneous)
	if err != nil {
		return fmt.Errorf("failed to create scale cluster workflow: %w", err)
	}
	body, err := yaml.Marshal(wf)
	if err != nil {
		return fmt.Errorf("failed to marshal scale cluster workflow: %w", err)
	}
	cluster.UpdateAnnotation(consts.AnnoKeyScaleClusterWorkflow, string(body))
	logger.Info("scale cluster workflow is annotated in cluster", zap.String("workflow", wf.Name))
	return nil
}

const (
	componentToggleTypeNone    = ""
	componentToggleTypeEnable  = "Enable"
	componentToggleTypeDisable = "Disable"
)

func getComponentToggleType(comp cloudv1.ClusterComponent, req *apiv1.ScaleClusterRequest) string {
	toggleType := componentToggleTypeNone
	if comp.GetReplicas() == 0 && req.NodeSpec.GetCount() > 0 {
		toggleType = componentToggleTypeEnable
	} else if comp.GetReplicas() > 0 && req.NodeSpec.GetCount() == 0 {
		toggleType = componentToggleTypeDisable
	}
	return toggleType
}

// CancelScaleCluster cancels an ongoing scale cluster operation.
func (s *ClusterServer) CancelScaleCluster(ctx context.Context, req *apiv1.CancelScaleClusterRequest) (*apiv1.Cluster, error) {
	logger := tilog.WithContextL(ctx).With(
		zap.String("cluster", req.ClusterId),
	)

	cluster, err := s.getClusterCR(ctx, req.ClusterId)
	if err != nil {
		logger.Error("failed to get cluster", zap.Error(err))
		return nil, err
	}

	// Get workflow from cluster annotation
	wfRef, err := getWorkflowFromAnnotation(cluster, consts.AnnoKeyScaleClusterWorkflow)
	if err != nil {
		logger.Warn("no ongoing scale cluster workflow found")
		return nil, status.Errorf(codes.NotFound, "no ongoing scale cluster workflow found")
	}

	// Get the actual workflow from Kubernetes
	existingWf, err := s.getWorkflow(ctx, client.ObjectKey{
		Namespace: wfRef.Namespace,
		Name:      wfRef.Name,
	})
	if err != nil {
		if apierrors.IsNotFound(err) {
			logger.Warn("workflow not found, may have been deleted", zap.String("workflow", wfRef.Name))
			return nil, status.Errorf(codes.NotFound, "workflow not found: %s", wfRef.Name)
		}
		logger.Error("failed to get workflow", zap.Error(err))
		return nil, status.Errorf(codes.Internal, "failed to get workflow: %s", err)
	}

	// Check if workflow is already in a terminal state
	if isWorkflowTerminal(existingWf.Status.Phase) {
		logger.Info("workflow is already in terminal state, no need to cancel",
			zap.String("workflow", existingWf.Name),
			zap.String("phase", string(existingWf.Status.Phase)),
		)
		return cvt.Cluster2PB(cluster), nil
	}

	// Cancel the workflow with retry support
	if err := s.cancelWorkflow(ctx, logger, existingWf, req.Reason); err != nil {
		return nil, err
	}

	logger.Info("marked workflow as cancelling", zap.String("workflow", existingWf.Name))
	return cvt.Cluster2PB(cluster), nil
}

const (
	defaultCancelReason = "cancelled by user"
)

// getWorkflowFromAnnotation extracts workflow reference from cluster annotation.
// Returns an error if the annotation is missing or invalid.
func getWorkflowFromAnnotation(cluster *cloudv1.Cluster, annotationKey string) (*cloudv1.Workflow, error) {
	anno, ok := cluster.Annotations[annotationKey]
	if !ok {
		return nil, fmt.Errorf("annotation %s not found", annotationKey)
	}

	var wf cloudv1.Workflow
	if err := yaml.Unmarshal([]byte(anno), &wf); err != nil {
		return nil, fmt.Errorf("failed to unmarshal workflow from annotation: %w", err)
	}

	return &wf, nil
}

// getWorkflow retrieves a workflow from Kubernetes by its ObjectKey.
func (s *ClusterServer) getWorkflow(ctx context.Context, key client.ObjectKey) (*cloudv1.Workflow, error) {
	wf := &cloudv1.Workflow{}
	if err := s.KubeClient.Get(ctx, key, wf); err != nil {
		return nil, err
	}
	return wf, nil
}

// isWorkflowTerminal checks if the workflow is in a terminal state.
// Terminal states include: Cancelling, Cancelled, Completed, and Failed.
func isWorkflowTerminal(phase cloudv1.WorkflowPhase) bool {
	return phase == cloudv1.WorkflowPhaseCancelling ||
		phase == cloudv1.WorkflowPhaseCancelled ||
		phase == cloudv1.WorkflowPhaseCompleted ||
		phase == cloudv1.WorkflowPhaseFailed
}

// cancelWorkflow cancels a workflow by adding a cancel annotation.
func (s *ClusterServer) cancelWorkflow(ctx context.Context, logger *zap.Logger, wf *cloudv1.Workflow, reason string) error {
	// Normalize cancel reason
	if reason == "" {
		reason = defaultCancelReason
	}

	// Get the latest version of the workflow
	latestWf, err := s.getWorkflow(ctx, client.ObjectKey{
		Namespace: wf.Namespace,
		Name:      wf.Name,
	})
	if err != nil {
		if apierrors.IsNotFound(err) {
			logger.Warn("workflow not found", zap.String("workflow", wf.Name))
			return status.Errorf(codes.NotFound, "workflow not found: %s", wf.Name)
		}
		logger.Error("failed to get workflow", zap.Error(err))
		return fmt.Errorf("get workflow: %w", err)
	}

	// Check if workflow is already in a terminal state
	if isWorkflowTerminal(latestWf.Status.Phase) {
		logger.Info("workflow is already in terminal state",
			zap.String("workflow", latestWf.Name),
			zap.String("phase", string(latestWf.Status.Phase)),
		)
		return nil
	}

	// Initialize annotations if needed
	if latestWf.Annotations == nil {
		latestWf.Annotations = make(map[string]string)
	}
	latestWf.Annotations[consts.AnnoCancelWorkflow] = reason

	// Update workflow
	if err := s.KubeClient.Update(ctx, latestWf); err != nil {
		logger.Error("failed to update workflow", zap.Error(err))
		return status.Errorf(codes.Internal, "failed to update workflow: %s", err)
	}

	return nil
}