package cluster

import (
	"context"
	stderrors "errors"
	"net/http"
	"testing"

	"github.com/samber/lo"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/types"
	ctrlclient "sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
	"sigs.k8s.io/yaml"

	clusterctrl "github.com/tidbcloud/eks-provider/controllers/cluster"
	"github.com/tidbcloud/eks-provider/controllers/utils/configmanager"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/dataplane"
	"github.com/tidbcloud/infra-provider/pkg/gkeheterogeneous"
	nodepoolutil "github.com/tidbcloud/infra-provider/pkg/nodepool"
	"github.com/tidbcloud/infra-provider/pkg/util"
)

const (
	gkeScaleTestBaseInstanceType  = "n2-standard-4"
	gkeScaleTestExtraInstanceType = "n2d-standard-4"
)

func Test_applyReplicasChanges(t *testing.T) {
	tests := map[string]struct {
		currentReplicas int
		desiredReplicas int
	}{
		"scale out": {currentReplicas: 1, desiredReplicas: 3},
		"scale in":  {currentReplicas: 5, desiredReplicas: 2},
		"no change": {currentReplicas: 2, desiredReplicas: 2},
	}

	for name, tc := range tests {
		t.Run(name, func(t *testing.T) {
			cur := &cloudv1.TiDB{}
			cur.SetReplicas(tc.currentReplicas)
			desired := &cloudv1.TiDB{}
			desired.SetReplicas(tc.desiredReplicas)
			inputs := &scaleClusterInputs{
				currentComponent: cur,
				desiredComponent: desired,
			}
			err := applyReplicasChanges(context.Background(), nil, cloudv1.WorkflowStep{}, inputs)
			assert.NoError(t, err)
			assert.Equal(t, tc.desiredReplicas, cur.GetReplicas())
		})
	}
}

func newGKEScaleTestCluster() *cloudv1.Cluster {
	cluster := &cloudv1.Cluster{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "cluster-a",
			Namespace: "ns-a",
			Labels: map[string]string{
				cloudv1.ProviderTypeLabel: cloudv1.ProviderTypeGKE,
				cloudv1.RegionLabel:       "us-central1",
				cloudv1.TenantTagKey:      "tenant-1",
				cloudv1.ProjectTagKey:     "project-1",
			},
		},
	}
	cluster.UpdateAnnotation(consts.AnnoKeyControlPlaneVersion, "1.35")
	return cluster
}

func scaleTestGKEComputeClassBindingAnnotation(t *testing.T, cluster *cloudv1.Cluster, annotationKey, baseInstanceType, extraInstanceType, extraNodePoolName string) string {
	t.Helper()
	baseNodePoolName := nodepoolutil.GetGKENodePoolName(cluster, baseInstanceType, cluster.GetAnnotations()[consts.AnnoKeyControlPlaneVersion])
	nodePools := []cloudv1.GKEComputeClassNodePoolBinding{
		{
			Role:         cloudv1.GKEComputeClassNodePoolRoleBase,
			InstanceType: baseInstanceType,
			NodePoolName: baseNodePoolName,
		},
		{
			Role:         cloudv1.GKEComputeClassNodePoolRoleExtra,
			InstanceType: extraInstanceType,
			NodePoolName: extraNodePoolName,
		},
	}
	computeClassName, err := cloudv1.GKEComputeClassNameForBindings(nodePools)
	if err != nil {
		t.Fatalf("build gke compute class name: %v", err)
	}
	annotation, err := cloudv1.EncodeGKEComputeClassBindingAnnotation(
		cloudv1.GKEComputeClassBindingAnnotation{
			ComputeClassName: computeClassName,
			NodePools:        nodePools,
		},
		annotationKey,
		baseInstanceType,
		baseNodePoolName,
	)
	if err != nil {
		t.Fatalf("build gke compute class binding annotation: %v", err)
	}
	return annotation
}

func Test_applyInstanceTypeChanges(t *testing.T) {
	tests := map[string]struct {
		componentType cloudv1.ComponentType
		currentType   string
		desiredType   string
		groupID       string
		cluster       *cloudv1.Cluster
		disable       bool
		wantPresent   string
		wantErr       string
	}{
		"change type":   {currentType: "t3.medium", desiredType: "m5.large"},
		"no change":     {currentType: "m5.large", desiredType: "m5.large"},
		"empty desired": {currentType: "t3.medium", desiredType: ""},
		"empty current": {currentType: "", desiredType: "m5.large"},
		"gke tikv heterogeneous blocks change": {
			currentType: gkeScaleTestBaseInstanceType,
			desiredType: gkeScaleTestExtraInstanceType,
			cluster: func() *cloudv1.Cluster {
				cluster := newGKEScaleTestCluster()
				cluster.Spec.Cluster.TiKV.InstanceType = gkeScaleTestBaseInstanceType
				cluster.UpdateAnnotation(
					cloudv1.AnnoKeyGKEComputeClassForTiKVBindings,
					scaleTestGKEComputeClassBindingAnnotation(t, cluster, cloudv1.AnnoKeyGKEComputeClassForTiKVBindings, gkeScaleTestBaseInstanceType, gkeScaleTestExtraInstanceType, "tikv-extra"),
				)
				return cluster
			}(),
			wantErr: "cannot change tikv instance type while GKE heterogeneous machine deploy is enabled",
		},
		"gke tikv workflow disable flag allows change and keeps intent during apply": {
			currentType: gkeScaleTestBaseInstanceType,
			desiredType: gkeScaleTestExtraInstanceType,
			cluster: func() *cloudv1.Cluster {
				cluster := newGKEScaleTestCluster()
				cluster.Spec.Cluster.TiKV.InstanceType = gkeScaleTestBaseInstanceType
				cluster.UpdateAnnotation(
					cloudv1.AnnoKeyGKEComputeClassForTiKVBindings,
					scaleTestGKEComputeClassBindingAnnotation(t, cluster, cloudv1.AnnoKeyGKEComputeClassForTiKVBindings, gkeScaleTestBaseInstanceType, gkeScaleTestExtraInstanceType, "tikv-extra"),
				)
				return cluster
			}(),
			disable:     true,
			wantPresent: cloudv1.AnnoKeyGKEComputeClassForTiKVBindings,
		},
		"gke tikv workflow disable flag does not clear intent without type change": {
			currentType: gkeScaleTestBaseInstanceType,
			desiredType: gkeScaleTestBaseInstanceType,
			cluster: func() *cloudv1.Cluster {
				cluster := newGKEScaleTestCluster()
				cluster.Spec.Cluster.TiKV.InstanceType = gkeScaleTestBaseInstanceType
				cluster.UpdateAnnotation(
					cloudv1.AnnoKeyGKEComputeClassForTiKVBindings,
					scaleTestGKEComputeClassBindingAnnotation(t, cluster, cloudv1.AnnoKeyGKEComputeClassForTiKVBindings, gkeScaleTestBaseInstanceType, gkeScaleTestExtraInstanceType, "tikv-extra"),
				)
				return cluster
			}(),
			disable:     true,
			wantPresent: cloudv1.AnnoKeyGKEComputeClassForTiKVBindings,
		},
		"gke tikv no binding intent does not block change": {
			currentType: gkeScaleTestBaseInstanceType,
			desiredType: gkeScaleTestExtraInstanceType,
			cluster: func() *cloudv1.Cluster {
				cluster := newGKEScaleTestCluster()
				cluster.Spec.Cluster.TiKV.InstanceType = gkeScaleTestBaseInstanceType
				return cluster
			}(),
		},
		"gke tikv malformed binding blocks change": {
			currentType: gkeScaleTestBaseInstanceType,
			desiredType: gkeScaleTestExtraInstanceType,
			cluster: func() *cloudv1.Cluster {
				cluster := newGKEScaleTestCluster()
				cluster.Spec.Cluster.TiKV.InstanceType = gkeScaleTestBaseInstanceType
				cluster.UpdateAnnotation(cloudv1.AnnoKeyGKEComputeClassForTiKVBindings, `not-json`)
				return cluster
			}(),
			wantErr: "cannot change tikv instance type while GKE heterogeneous machine deploy is enabled",
		},
		"gke tidb default group heterogeneous blocks change": {
			componentType: cloudv1.ComponentTiDB,
			currentType:   gkeScaleTestBaseInstanceType,
			desiredType:   gkeScaleTestExtraInstanceType,
			cluster: func() *cloudv1.Cluster {
				cluster := newGKEScaleTestCluster()
				cluster.Spec.Cluster.TiDB.InstanceType = gkeScaleTestBaseInstanceType
				bindingKey := cluster.TiDBGKEComputeClassBindingsAnnoKey("")
				cluster.UpdateAnnotation(
					bindingKey,
					scaleTestGKEComputeClassBindingAnnotation(t, cluster, bindingKey, gkeScaleTestBaseInstanceType, gkeScaleTestExtraInstanceType, "tidb-default-extra"),
				)
				return cluster
			}(),
			wantErr: "cannot change tidb instance type while GKE heterogeneous machine deploy is enabled",
		},
		"gke tidb extra group heterogeneous blocks change": {
			componentType: cloudv1.ComponentTiDB,
			currentType:   gkeScaleTestBaseInstanceType,
			desiredType:   gkeScaleTestExtraInstanceType,
			groupID:       "group1",
			cluster: func() *cloudv1.Cluster {
				cluster := newGKEScaleTestCluster()
				cluster.Spec.ExtraGroups = []cloudv1.ExtraGroupSpec{{
					Name: "group1",
					TiDB: cloudv1.TiDB{
						InstanceType: gkeScaleTestBaseInstanceType,
					},
				}}
				bindingKey := cluster.TiDBGKEComputeClassBindingsAnnoKey("group1")
				cluster.UpdateAnnotation(
					bindingKey,
					scaleTestGKEComputeClassBindingAnnotation(t, cluster, bindingKey, gkeScaleTestBaseInstanceType, gkeScaleTestExtraInstanceType, "tidb-extra"),
				)
				return cluster
			}(),
			wantErr: "cannot change tidb instance type while GKE heterogeneous machine deploy is enabled",
		},
		"gke tidb no binding intent does not block change": {
			componentType: cloudv1.ComponentTiDB,
			currentType:   gkeScaleTestBaseInstanceType,
			desiredType:   gkeScaleTestExtraInstanceType,
			groupID:       "group1",
			cluster: func() *cloudv1.Cluster {
				cluster := newGKEScaleTestCluster()
				cluster.Spec.ExtraGroups = []cloudv1.ExtraGroupSpec{{Name: "group1"}}
				return cluster
			}(),
		},
		"gke tidb extra group workflow disable flag allows change and keeps group intent during apply": {
			componentType: cloudv1.ComponentTiDB,
			currentType:   gkeScaleTestBaseInstanceType,
			desiredType:   gkeScaleTestExtraInstanceType,
			groupID:       "group1",
			cluster: func() *cloudv1.Cluster {
				cluster := newGKEScaleTestCluster()
				cluster.Spec.ExtraGroups = []cloudv1.ExtraGroupSpec{{
					Name: "group1",
					TiDB: cloudv1.TiDB{
						InstanceType: gkeScaleTestBaseInstanceType,
					},
				}}
				bindingKey := cluster.TiDBGKEComputeClassBindingsAnnoKey("group1")
				cluster.UpdateAnnotation(
					bindingKey,
					scaleTestGKEComputeClassBindingAnnotation(t, cluster, bindingKey, gkeScaleTestBaseInstanceType, gkeScaleTestExtraInstanceType, "tidb-extra"),
				)
				return cluster
			}(),
			disable:     true,
			wantPresent: (&cloudv1.Cluster{}).TiDBGKEComputeClassBindingsAnnoKey("group1"),
		},
		"gke tidb malformed extra group binding blocks change": {
			componentType: cloudv1.ComponentTiDB,
			currentType:   gkeScaleTestBaseInstanceType,
			desiredType:   gkeScaleTestExtraInstanceType,
			groupID:       "group1",
			cluster: func() *cloudv1.Cluster {
				cluster := newGKEScaleTestCluster()
				cluster.Spec.ExtraGroups = []cloudv1.ExtraGroupSpec{{
					Name: "group1",
					TiDB: cloudv1.TiDB{
						InstanceType: gkeScaleTestBaseInstanceType,
					},
				}}
				cluster.UpdateAnnotation(cluster.TiDBGKEComputeClassBindingsAnnoKey("group1"), `not-json`)
				return cluster
			}(),
			wantErr: "cannot change tidb instance type while GKE heterogeneous machine deploy is enabled",
		},
		"gke tidb valid sibling intent does not block unrelated group change": {
			componentType: cloudv1.ComponentTiDB,
			currentType:   gkeScaleTestBaseInstanceType,
			desiredType:   gkeScaleTestExtraInstanceType,
			groupID:       "group2",
			cluster: func() *cloudv1.Cluster {
				cluster := newGKEScaleTestCluster()
				cluster.Spec.ExtraGroups = []cloudv1.ExtraGroupSpec{
					{
						Name: "group1",
						TiDB: cloudv1.TiDB{
							InstanceType: gkeScaleTestBaseInstanceType,
						},
					},
					{
						Name: "group2",
						TiDB: cloudv1.TiDB{
							InstanceType: gkeScaleTestBaseInstanceType,
						},
					},
				}
				bindingKey := cluster.TiDBGKEComputeClassBindingsAnnoKey("group1")
				cluster.UpdateAnnotation(
					bindingKey,
					scaleTestGKEComputeClassBindingAnnotation(t, cluster, bindingKey, gkeScaleTestBaseInstanceType, gkeScaleTestExtraInstanceType, "tidb-extra"),
				)
				return cluster
			}(),
		},
	}

	for name, tc := range tests {
		t.Run(name, func(t *testing.T) {
			componentType := tc.componentType
			if componentType == "" {
				componentType = cloudv1.ComponentTiKV
			}
			var cur cloudv1.ClusterComponent
			var desired cloudv1.ClusterComponent
			switch componentType {
			case cloudv1.ComponentTiDB:
				cur = &cloudv1.TiDB{}
				desired = &cloudv1.TiDB{}
			default:
				cur = &cloudv1.TiKV{}
				desired = &cloudv1.TiKV{}
			}
			cur.SetInstanceType(tc.currentType)
			desired.SetInstanceType(tc.desiredType)
			inputs := &scaleClusterInputs{
				currentComponent: cur,
				desiredComponent: desired,
				cluster:          tc.cluster,
				GroupID:          tc.groupID,
				ComponentType:    string(componentType),
			}
			var wf *cloudv1.Workflow
			if tc.disable {
				wf = &cloudv1.Workflow{
					Spec: cloudv1.WorkflowSpec{
						Inputs: map[string]string{
							consts.WfInputKeyScaleClusterDisableGKEHeterogeneous: "true",
						},
					},
				}
			}
			err := applyInstanceTypeChanges(context.Background(), wf, cloudv1.WorkflowStep{}, inputs)
			if tc.wantErr != "" {
				assert.Error(t, err)
				assert.ErrorContains(t, err, tc.wantErr)
				assert.Equal(t, tc.currentType, cur.GetInstanceType())
				return
			}
			assert.NoError(t, err)
			assert.Equal(t, tc.desiredType, cur.GetInstanceType())
			if tc.wantPresent != "" {
				assert.Contains(t, tc.cluster.GetAnnotations(), tc.wantPresent)
			}
		})
	}
}

func Test_clearGKEHeterogeneousIntentForInstanceTypeScale(t *testing.T) {
	tests := map[string]struct {
		provider    string
		component   cloudv1.ComponentType
		groupID     string
		scopes      []string
		disable     bool
		original    string
		desired     string
		wantDeleted bool
	}{
		"tikv instance type disable clears intent": {
			provider:    cloudv1.ProviderTypeGKE,
			component:   cloudv1.ComponentTiKV,
			scopes:      []string{consts.WfScaleClusterScopeInstanceType},
			disable:     true,
			original:    gkeScaleTestBaseInstanceType,
			desired:     gkeScaleTestExtraInstanceType,
			wantDeleted: true,
		},
		"tikv storage disable keeps intent": {
			provider:  cloudv1.ProviderTypeGKE,
			component: cloudv1.ComponentTiKV,
			scopes:    []string{consts.WfScaleClusterScopeStorage},
			disable:   true,
			original:  gkeScaleTestBaseInstanceType,
			desired:   gkeScaleTestExtraInstanceType,
		},
		"tikv instance type without disable keeps intent": {
			provider:  cloudv1.ProviderTypeGKE,
			component: cloudv1.ComponentTiKV,
			scopes:    []string{consts.WfScaleClusterScopeInstanceType},
			original:  gkeScaleTestBaseInstanceType,
			desired:   gkeScaleTestExtraInstanceType,
		},
		"tikv instance type disable without type change keeps intent": {
			provider:  cloudv1.ProviderTypeGKE,
			component: cloudv1.ComponentTiKV,
			scopes:    []string{consts.WfScaleClusterScopeInstanceType},
			disable:   true,
			original:  gkeScaleTestBaseInstanceType,
			desired:   gkeScaleTestBaseInstanceType,
		},
		"non gke instance type disable keeps intent": {
			provider:  cloudv1.ProviderTypeAWS,
			component: cloudv1.ComponentTiKV,
			scopes:    []string{consts.WfScaleClusterScopeInstanceType},
			disable:   true,
			original:  gkeScaleTestBaseInstanceType,
			desired:   gkeScaleTestExtraInstanceType,
		},
		"tidb group instance type disable clears group intent": {
			provider:    cloudv1.ProviderTypeGKE,
			component:   cloudv1.ComponentTiDB,
			groupID:     "group1",
			scopes:      []string{consts.WfScaleClusterScopeInstanceType},
			disable:     true,
			original:    gkeScaleTestBaseInstanceType,
			desired:     gkeScaleTestExtraInstanceType,
			wantDeleted: true,
		},
	}

	for name, tc := range tests {
		t.Run(name, func(t *testing.T) {
			cluster := newGKEScaleTestCluster()
			cluster.Labels[cloudv1.ProviderTypeLabel] = tc.provider

			annotationKey := cloudv1.AnnoKeyGKEComputeClassForTiKVBindings
			if tc.component == cloudv1.ComponentTiDB {
				annotationKey = cluster.TiDBGKEComputeClassBindingsAnnoKey(tc.groupID)
			}
			cluster.UpdateAnnotation(annotationKey, "binding")

			var wf *cloudv1.Workflow
			if tc.disable {
				wf = &cloudv1.Workflow{
					Spec: cloudv1.WorkflowSpec{
						Inputs: map[string]string{
							consts.WfInputKeyScaleClusterDisableGKEHeterogeneous: "true",
						},
					},
				}
			}
			inputs := &scaleClusterInputs{
				cluster:       cluster,
				ComponentType: string(tc.component),
				GroupID:       tc.groupID,
				Scopes:        tc.scopes,
			}
			switch tc.component {
			case cloudv1.ComponentTiDB:
				inputs.originalComponent = &cloudv1.TiDB{InstanceType: tc.original}
				inputs.desiredComponent = &cloudv1.TiDB{InstanceType: tc.desired}
			default:
				inputs.originalComponent = &cloudv1.TiKV{InstanceType: tc.original}
				inputs.desiredComponent = &cloudv1.TiKV{InstanceType: tc.desired}
			}

			if shouldDisableGKEHeterogeneousForInstanceTypeScale(wf, inputs) {
				clearGKEHeterogeneousIntentInScaleInputs(inputs)
			}
			if tc.wantDeleted {
				assert.NotContains(t, cluster.GetAnnotations(), annotationKey)
			} else {
				assert.Contains(t, cluster.GetAnnotations(), annotationKey)
			}
		})
	}
}

func TestScaleClusterHandleDeletesGKEHeterogeneousGatekeeperBeforeClearingIntent(t *testing.T) {
	cluster := newGKEScaleTestCluster()
	cluster.EnableFeature(consts.ClusterFeatureUseTiDBOperatorV2)
	cluster.Spec.Cluster.TiKV.InstanceType = gkeScaleTestBaseInstanceType
	cluster.UpdateAnnotation(
		cloudv1.AnnoKeyGKEComputeClassForTiKVBindings,
		scaleTestGKEComputeClassBindingAnnotation(t, cluster, cloudv1.AnnoKeyGKEComputeClassForTiKVBindings, gkeScaleTestBaseInstanceType, gkeScaleTestExtraInstanceType, "tikv-extra"),
	)
	primary := scaleTestGatekeeperAssign(cluster.Name, cloudv1.ComponentTiKV, "", "compute-class-assign", gkeheterogeneous.GatekeeperAssignPurposeNodeSelector)
	companion := scaleTestGatekeeperAssign(cluster.Name, cloudv1.ComponentTiKV, "", "topology-spread-assign", gkeheterogeneous.GatekeeperAssignPurposeTopologySpread)

	localCli, remoteCli, action := newScaleClusterActionHandleTest(t, cluster, primary, companion)
	step := scaleClusterHandleTestStep(t, cluster, cloudv1.ComponentTiKV, "", gkeScaleTestBaseInstanceType, gkeScaleTestExtraInstanceType)
	wf := scaleClusterDisableGKEHeterogeneousWorkflow()

	require.NoError(t, action.Handle(context.Background(), wf, step))

	gotCluster := &cloudv1.Cluster{}
	require.NoError(t, localCli.Get(context.Background(), ctrlclient.ObjectKey{Namespace: cluster.Namespace, Name: cluster.Name}, gotCluster))
	require.Equal(t, gkeScaleTestExtraInstanceType, gotCluster.Spec.Cluster.TiKV.InstanceType)
	require.NotContains(t, gotCluster.GetAnnotations(), cloudv1.AnnoKeyGKEComputeClassForTiKVBindings)
	err := remoteCli.Get(context.Background(), ctrlclient.ObjectKey{Name: primary.GetName()}, &unstructured.Unstructured{Object: map[string]interface{}{"apiVersion": primary.GetAPIVersion(), "kind": primary.GetKind()}})
	require.True(t, apierrors.IsNotFound(err), "primary gatekeeper assign should be deleted, got %v", err)
	err = remoteCli.Get(context.Background(), ctrlclient.ObjectKey{Name: companion.GetName()}, &unstructured.Unstructured{Object: map[string]interface{}{"apiVersion": companion.GetAPIVersion(), "kind": companion.GetKind()}})
	require.True(t, apierrors.IsNotFound(err), "topology spread gatekeeper assign should be deleted, got %v", err)
}

func TestScaleClusterHandleDeletesTiDBGKEHeterogeneousGatekeeperForGroup(t *testing.T) {
	cluster := newGKEScaleTestCluster()
	cluster.EnableFeature(consts.ClusterFeatureUseTiDBOperatorV2)
	cluster.Spec.ExtraGroups = []cloudv1.ExtraGroupSpec{
		{
			Name: "group1",
			TiDB: cloudv1.TiDB{InstanceType: gkeScaleTestBaseInstanceType},
		},
	}
	annotationKey := cluster.TiDBGKEComputeClassBindingsAnnoKey("group1")
	cluster.UpdateAnnotation(
		annotationKey,
		scaleTestGKEComputeClassBindingAnnotation(t, cluster, annotationKey, gkeScaleTestBaseInstanceType, gkeScaleTestExtraInstanceType, "tidb-extra"),
	)
	primary := scaleTestGatekeeperAssign(cluster.Name, cloudv1.ComponentTiDB, "group1", "compute-class-assign", gkeheterogeneous.GatekeeperAssignPurposeNodeSelector)
	companion := scaleTestGatekeeperAssign(cluster.Name, cloudv1.ComponentTiDB, "group1", "topology-spread-assign", gkeheterogeneous.GatekeeperAssignPurposeTopologySpread)

	localCli, remoteCli, action := newScaleClusterActionHandleTest(t, cluster, primary, companion)
	step := scaleClusterHandleTestStep(t, cluster, cloudv1.ComponentTiDB, "group1", gkeScaleTestBaseInstanceType, gkeScaleTestExtraInstanceType)
	wf := scaleClusterDisableGKEHeterogeneousWorkflow()

	require.NoError(t, action.Handle(context.Background(), wf, step))

	gotCluster := &cloudv1.Cluster{}
	require.NoError(t, localCli.Get(context.Background(), ctrlclient.ObjectKey{Namespace: cluster.Namespace, Name: cluster.Name}, gotCluster))
	gotTiDB, err := gotCluster.GetComponentSpecInGroup("group1", cloudv1.ComponentTiDB)
	require.NoError(t, err)
	require.Equal(t, gkeScaleTestExtraInstanceType, gotTiDB.GetInstanceType())
	require.NotContains(t, gotCluster.GetAnnotations(), annotationKey)
	err = remoteCli.Get(context.Background(), ctrlclient.ObjectKey{Name: primary.GetName()}, &unstructured.Unstructured{Object: map[string]interface{}{"apiVersion": primary.GetAPIVersion(), "kind": primary.GetKind()}})
	require.True(t, apierrors.IsNotFound(err), "primary gatekeeper assign should be deleted, got %v", err)
	err = remoteCli.Get(context.Background(), ctrlclient.ObjectKey{Name: companion.GetName()}, &unstructured.Unstructured{Object: map[string]interface{}{"apiVersion": companion.GetAPIVersion(), "kind": companion.GetKind()}})
	require.True(t, apierrors.IsNotFound(err), "topology spread gatekeeper assign should be deleted, got %v", err)
}

func TestScaleClusterHandleConvergesOnWorkflowRetryAfterClusterUpdateConflict(t *testing.T) {
	cluster := newGKEScaleTestCluster()
	cluster.EnableFeature(consts.ClusterFeatureUseTiDBOperatorV2)
	cluster.Spec.Cluster.TiKV.InstanceType = gkeScaleTestBaseInstanceType
	cluster.UpdateAnnotation(
		cloudv1.AnnoKeyGKEComputeClassForTiKVBindings,
		scaleTestGKEComputeClassBindingAnnotation(t, cluster, cloudv1.AnnoKeyGKEComputeClassForTiKVBindings, gkeScaleTestBaseInstanceType, gkeScaleTestExtraInstanceType, "tikv-extra"),
	)
	primary := scaleTestGatekeeperAssign(cluster.Name, cloudv1.ComponentTiKV, "", "compute-class-assign", gkeheterogeneous.GatekeeperAssignPurposeNodeSelector)

	localCli, remoteCli, action := newScaleClusterActionHandleTest(t, cluster, primary)
	conflictCli := &conflictOnceClient{Client: localCli, conflictsLeft: 1}
	action.cli = conflictCli
	step := scaleClusterHandleTestStep(t, cluster, cloudv1.ComponentTiKV, "", gkeScaleTestBaseInstanceType, gkeScaleTestExtraInstanceType)
	wf := scaleClusterDisableGKEHeterogeneousWorkflow()

	err := action.Handle(context.Background(), wf, step)
	require.True(t, apierrors.IsConflict(err), "first handle should surface the cluster update conflict, got %v", err)
	require.Zero(t, conflictCli.conflictsLeft)

	gotCluster := &cloudv1.Cluster{}
	require.NoError(t, localCli.Get(context.Background(), ctrlclient.ObjectKey{Namespace: cluster.Namespace, Name: cluster.Name}, gotCluster))
	require.Equal(t, gkeScaleTestBaseInstanceType, gotCluster.Spec.Cluster.TiKV.InstanceType)
	require.Contains(t, gotCluster.GetAnnotations(), cloudv1.AnnoKeyGKEComputeClassForTiKVBindings)
	err = remoteCli.Get(context.Background(), ctrlclient.ObjectKey{Name: primary.GetName()}, &unstructured.Unstructured{Object: map[string]interface{}{"apiVersion": primary.GetAPIVersion(), "kind": primary.GetKind()}})
	require.True(t, apierrors.IsNotFound(err), "gatekeeper assign should be deleted before the conflicted cluster update, got %v", err)

	require.NoError(t, action.Handle(context.Background(), wf, step))
	require.NoError(t, localCli.Get(context.Background(), ctrlclient.ObjectKey{Namespace: cluster.Namespace, Name: cluster.Name}, gotCluster))
	require.Equal(t, gkeScaleTestExtraInstanceType, gotCluster.Spec.Cluster.TiKV.InstanceType)
	require.NotContains(t, gotCluster.GetAnnotations(), cloudv1.AnnoKeyGKEComputeClassForTiKVBindings)
	err = remoteCli.Get(context.Background(), ctrlclient.ObjectKey{Name: primary.GetName()}, &unstructured.Unstructured{Object: map[string]interface{}{"apiVersion": primary.GetAPIVersion(), "kind": primary.GetKind()}})
	require.True(t, apierrors.IsNotFound(err), "gatekeeper assign should stay deleted across retry, got %v", err)
}

func TestScaleClusterHandleKeepsIntentWhenGKEHeterogeneousGatekeeperDeleteFails(t *testing.T) {
	cluster := newGKEScaleTestCluster()
	cluster.EnableFeature(consts.ClusterFeatureUseTiDBOperatorV2)
	cluster.Spec.Cluster.TiKV.InstanceType = gkeScaleTestBaseInstanceType
	cluster.UpdateAnnotation(
		cloudv1.AnnoKeyGKEComputeClassForTiKVBindings,
		scaleTestGKEComputeClassBindingAnnotation(t, cluster, cloudv1.AnnoKeyGKEComputeClassForTiKVBindings, gkeScaleTestBaseInstanceType, gkeScaleTestExtraInstanceType, "tikv-extra"),
	)
	primary := scaleTestGatekeeperAssign(cluster.Name, cloudv1.ComponentTiKV, "", "compute-class-assign", gkeheterogeneous.GatekeeperAssignPurposeNodeSelector)
	labels := primary.GetLabels()
	labels[consts.ClusterLabelKey] = "other-cluster"
	primary.SetLabels(labels)

	localCli, remoteCli, action := newScaleClusterActionHandleTest(t, cluster, primary)
	step := scaleClusterHandleTestStep(t, cluster, cloudv1.ComponentTiKV, "", gkeScaleTestBaseInstanceType, gkeScaleTestExtraInstanceType)
	wf := scaleClusterDisableGKEHeterogeneousWorkflow()

	err := action.Handle(context.Background(), wf, step)
	require.Error(t, err)
	require.ErrorContains(t, err, consts.ClusterLabelKey)

	gotCluster := &cloudv1.Cluster{}
	require.NoError(t, localCli.Get(context.Background(), ctrlclient.ObjectKey{Namespace: cluster.Namespace, Name: cluster.Name}, gotCluster))
	require.Equal(t, gkeScaleTestBaseInstanceType, gotCluster.Spec.Cluster.TiKV.InstanceType)
	require.Contains(t, gotCluster.GetAnnotations(), cloudv1.AnnoKeyGKEComputeClassForTiKVBindings)
	gotAssign := &unstructured.Unstructured{}
	gotAssign.SetGroupVersionKind(primary.GroupVersionKind())
	require.NoError(t, remoteCli.Get(context.Background(), ctrlclient.ObjectKey{Name: primary.GetName()}, gotAssign))
}

func Test_applyStorageChanges(t *testing.T) {
	tests := map[string]struct {
		currentSize  resource.Quantity
		desiredSize  resource.Quantity
		currentClass string
		desiredClass string
		currentAdd   []cloudv1.ComponentStorageVolume
		desiredAdd   []cloudv1.ComponentStorageVolume
	}{
		"change size and class": {
			currentSize: *util.StorageSizeGi2Quantity(100), desiredSize: *util.StorageSizeGi2Quantity(200),
			currentClass: "gp2", desiredClass: "gp3",
			currentAdd: nil, desiredAdd: []cloudv1.ComponentStorageVolume{{
				Name:             "extra",
				StorageClassName: "gp3",
				StorageSize:      *util.StorageSizeGi2Quantity(100),
				MountPath:        "/mnt/extra",
			}},
		},
		"no change": {
			currentSize: *util.StorageSizeGi2Quantity(100), desiredSize: *util.StorageSizeGi2Quantity(100),
			currentClass: "gp2", desiredClass: "gp2",
			currentAdd: nil, desiredAdd: nil,
		},
		"only additional storage": {
			currentSize: *util.StorageSizeGi2Quantity(100), desiredSize: *util.StorageSizeGi2Quantity(100),
			currentClass: "gp2", desiredClass: "gp2",
			currentAdd: nil, desiredAdd: []cloudv1.ComponentStorageVolume{{
				Name:             "extra",
				StorageClassName: "gp3",
				StorageSize:      *util.StorageSizeGi2Quantity(10),
				MountPath:        "/mnt/extra",
			}},
		},
	}

	for name, tc := range tests {
		t.Run(name, func(t *testing.T) {
			// Use TiKV for testing since it supports additional storages
			cur := &cloudv1.TiKV{}
			cur.SetStorageSize(tc.currentSize)
			cur.SetStorageClassName(tc.currentClass)
			cur.SetAdditionalStorages(tc.currentAdd)

			desired := &cloudv1.TiKV{}
			desired.SetStorageSize(tc.desiredSize)
			desired.SetStorageClassName(tc.desiredClass)
			desired.SetAdditionalStorages(tc.desiredAdd)

			inputs := &scaleClusterInputs{
				currentComponent: cur,
				desiredComponent: desired,
			}
			err := applyStorageChanges(context.Background(), nil, cloudv1.WorkflowStep{}, inputs)
			assert.NoError(t, err)
			assert.Equal(t, tc.desiredSize, mustGetStorageSize(cur))
			assert.Equal(t, tc.desiredClass, mustGetStorageClassName(cur))
			assert.Equal(t, tc.desiredAdd, cur.GetAdditionalStorages())
		})
	}
}

func Test_applyConfigRevisionChanges(t *testing.T) {
	tests := map[string]struct {
		desiredConfigRevisionName string
		expectSet                 bool
	}{
		"set revision":   {desiredConfigRevisionName: "rev-123", expectSet: true},
		"empty revision": {desiredConfigRevisionName: "", expectSet: false},
	}

	for name, tc := range tests {
		t.Run(name, func(t *testing.T) {
			cluster := &cloudv1.Cluster{
				Spec: cloudv1.ClusterSpec{
					ConfigRevisionRef: &cloudv1.Reference{
						Name: "rev-old",
					},
				},
			}
			inputs := &scaleClusterInputs{
				cluster:                   cluster,
				DesiredConfigRevisionName: tc.desiredConfigRevisionName,
			}
			err := applyConfigRevisionChanges(context.Background(), nil, cloudv1.WorkflowStep{}, inputs)
			assert.NoError(t, err)
			if tc.expectSet {
				assert.NotNil(t, cluster.Spec.ConfigRevisionRef)
				assert.Equal(t, tc.desiredConfigRevisionName, cluster.Spec.ConfigRevisionRef.Name)
			} else {
				assert.Equal(t, cluster.Spec.ConfigRevisionRef.Name, "rev-old")
			}
		})
	}
}

func Test_getPDClientUsesResolverAndAuth(t *testing.T) {
	oldFactory := newServiceHTTPEndpoint
	t.Cleanup(func() {
		newServiceHTTPEndpoint = oldFactory
	})

	wantService := types.NamespacedName{
		Namespace: "tidb-admin",
		Name:      "clustergateway-kong-proxy",
	}
	gotService := types.NamespacedName{}
	gotClusterName := ""
	gotDryRun := false
	authCalled := false
	newServiceHTTPEndpoint = func(_ ctrlclient.Client, serviceKey types.NamespacedName) clusterctrl.DataplaneHTTPEndpoint {
		gotService = serviceKey
		return fakeDataplaneHTTPEndpoint{
			getPDClientFn: func(ctx context.Context, httpCli *http.Client, pdAuthFunc configmanager.PDAuth, clusterName string, dryRun bool) (configmanager.DynamicPDClient, error) {
				gotClusterName = clusterName
				gotDryRun = dryRun
				req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.com", nil)
				assert.NoError(t, err)
				assert.NoError(t, pdAuthFunc(req))
				authCalled = req.Header.Get("Authorization") == "Bearer test-token"
				return configmanager.NewDefaultPDClient("http://example.com"), nil
			},
		}
	}

	clusterObj := &cloudv1.Cluster{}
	clusterObj.Name = "cluster-a"
	resolver := func(cluster *cloudv1.Cluster) types.NamespacedName {
		assert.Same(t, clusterObj, cluster)
		return wantService
	}
	pdAuth := func(r *http.Request) error {
		r.Header.Set("Authorization", "Bearer test-token")
		return nil
	}

	_, err := getPDClient(context.Background(), clusterObj, nil, pdAuth, resolver, true)
	assert.NoError(t, err)
	assert.Equal(t, wantService, gotService)
	assert.Equal(t, "cluster-a", gotClusterName)
	assert.True(t, gotDryRun)
	assert.True(t, authCalled)
}

type fakeDataplaneHTTPEndpoint struct {
	getPDClientFn func(ctx context.Context, httpCli *http.Client, pdAuthFunc configmanager.PDAuth, clusterName string, dryRun bool) (configmanager.DynamicPDClient, error)
}

func (e fakeDataplaneHTTPEndpoint) GetProxyURL(ctx context.Context, clusterName string) (string, error) {
	return "", nil
}

func (e fakeDataplaneHTTPEndpoint) GetPDClient(ctx context.Context, httpCli *http.Client, pdAuthFunc configmanager.PDAuth, clusterName string, dryRun bool) (configmanager.DynamicPDClient, error) {
	return e.getPDClientFn(ctx, httpCli, pdAuthFunc, clusterName, dryRun)
}

// helpers
func mustGetStorageSize(comp cloudv1.ClusterComponent) resource.Quantity {
	return *lo.Must1(comp.GetStorageSize())
}

func mustGetStorageClassName(comp cloudv1.ClusterComponent) string {
	return lo.Must1(comp.GetStorageClassName())
}

func newScaleClusterActionHandleTest(t *testing.T, cluster *cloudv1.Cluster, remoteObjects ...runtime.Object) (ctrlclient.Client, ctrlclient.Client, *scaleClusterAction) {
	t.Helper()
	scheme := runtime.NewScheme()
	require.NoError(t, cloudv1.AddToScheme(scheme))
	localCli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(cluster).Build()
	remoteCli := fake.NewClientBuilder().WithRuntimeObjects(remoteObjects...).Build()
	envProviderManager := dataplane.NewFakeEnvProviderManager(runtime.NewScheme(), func(ctx context.Context, scheme *runtime.Scheme, cluster *cloudv1.Cluster) (dataplane.EnvProvider, error) {
		return &dataplane.FakeEnvProvider{Client: remoteCli}, nil
	})
	return localCli, remoteCli, &scaleClusterAction{
		cli:            localCli,
		cachedProvider: envProviderManager,
	}
}

type conflictOnceClient struct {
	ctrlclient.Client
	conflictsLeft int
}

func (c *conflictOnceClient) Update(ctx context.Context, obj ctrlclient.Object, opts ...ctrlclient.UpdateOption) error {
	if c.conflictsLeft > 0 && obj.GetName() == "cluster-a" {
		c.conflictsLeft--
		return apierrors.NewConflict(schema.GroupResource{
			Group:    cloudv1.GroupVersion.Group,
			Resource: "clusters",
		}, obj.GetName(), stderrors.New("conflict"))
	}
	return c.Client.Update(ctx, obj, opts...)
}

func scaleClusterHandleTestStep(t *testing.T, cluster *cloudv1.Cluster, componentType cloudv1.ComponentType, groupID, originalInstanceType, desiredInstanceType string) cloudv1.WorkflowStep {
	t.Helper()
	originalComponent, desiredComponent := scaleClusterHandleTestComponentSpecs(componentType, originalInstanceType, desiredInstanceType)
	originalSpec, err := yaml.Marshal(originalComponent)
	require.NoError(t, err)
	desiredSpec, err := yaml.Marshal(desiredComponent)
	require.NoError(t, err)
	return cloudv1.WorkflowStep{
		Inputs: cloudv1.InputsWrapper{Object: map[string]interface{}{
			"namespace":             cluster.Namespace,
			"clusterId":             cluster.Name,
			"groupId":               groupID,
			"componentType":         string(componentType),
			"originalComponentSpec": string(originalSpec),
			"componentSpec":         string(desiredSpec),
			"scopes":                []string{consts.WfScaleClusterScopeInstanceType},
		}},
	}
}

func scaleClusterHandleTestComponentSpecs(componentType cloudv1.ComponentType, originalInstanceType, desiredInstanceType string) (interface{}, interface{}) {
	switch componentType {
	case cloudv1.ComponentTiDB:
		return &cloudv1.TiDB{InstanceType: originalInstanceType}, &cloudv1.TiDB{InstanceType: desiredInstanceType}
	default:
		return &cloudv1.TiKV{InstanceType: originalInstanceType}, &cloudv1.TiKV{InstanceType: desiredInstanceType}
	}
}

func scaleClusterDisableGKEHeterogeneousWorkflow() *cloudv1.Workflow {
	return &cloudv1.Workflow{
		Spec: cloudv1.WorkflowSpec{
			Inputs: map[string]string{
				consts.WfInputKeyScaleClusterDisableGKEHeterogeneous: "true",
			},
		},
	}
}

func scaleTestGatekeeperAssign(clusterID string, componentType cloudv1.ComponentType, tidbGroupID, suffix, purpose string) *unstructured.Unstructured {
	obj := &unstructured.Unstructured{
		Object: map[string]interface{}{
			"apiVersion": "mutations.gatekeeper.sh/v1",
			"kind":       "Assign",
			"metadata": map[string]interface{}{
				"name": gkeheterogeneous.ResourceName(clusterID, componentType, tidbGroupID, suffix),
			},
			"spec": map[string]interface{}{
				"match": map[string]interface{}{
					"scope":      "Namespaced",
					"namespaces": []interface{}{"tidb" + clusterID},
				},
			},
		},
	}
	labels := map[string]string{
		consts.ManagedByLabelKey:                         gkeheterogeneous.ManagedByClusterAgentLabelValue,
		consts.ClusterLabelKey:                           clusterID,
		gkeheterogeneous.HeterogeneousDeployLabelKey:     gkeheterogeneous.GatekeeperAssignDeployLabelValue,
		gkeheterogeneous.HeterogeneousComponentLabelKey:  string(componentType),
		gkeheterogeneous.GatekeeperAssignPurposeLabelKey: purpose,
	}
	if componentType == cloudv1.ComponentTiDB && tidbGroupID != "" {
		labels[consts.LabelKeyTiDBGroupId] = tidbGroupID
	}
	obj.SetLabels(labels)
	if purpose == gkeheterogeneous.GatekeeperAssignPurposeTopologySpread {
		_ = unstructured.SetNestedField(obj.Object, "spec.topologySpreadConstraints[topologyKey:'topology.kubernetes.io/zone'].nodeAffinityPolicy", "spec", "location")
		_ = unstructured.SetNestedField(obj.Object, "Ignore", "spec", "parameters", "assign", "value")
	}
	return obj
}