package cluster

import (
	"context"
	"fmt"
	"sort"
	"strings"

	"github.com/go-logr/logr"
	"github.com/mitchellh/mapstructure"
	"github.com/pingcap/errors"
	"k8s.io/apimachinery/pkg/util/sets"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"
	"sigs.k8s.io/yaml"

	opv1alpha1 "github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"

	"github.com/tidbcloud/eks-provider/controllers/cluster/task/common"
	"github.com/tidbcloud/eks-provider/controllers/utils"
	"github.com/tidbcloud/eks-provider/controllers/utils/configmanager"
	"github.com/tidbcloud/eks-provider/controllers/workflow/action"
	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"
	"github.com/tidbcloud/infra-provider/pkg/k8s"
	"github.com/tidbcloud/infra-provider/pkg/pingcap/label"
)

// scaleClusterAction scale a cluster component according to the componentSpec and specified scopes.
type scaleClusterAction struct {
	cli               client.Client
	cachedProvider    dataplane.EnvProviderManager
	pdAuth            configmanager.PDAuth
	pdServiceResolver PDServiceResolver
	dryRun            bool
}

type scaleClusterInputs struct {
	Namespace                 string   `mapstructure:"namespace" json:"namespace" validate:"required"`
	ClusterID                 string   `mapstructure:"clusterId" json:"clusterId" validate:"required"`
	GroupID                   string   `mapstructure:"groupId" json:"groupId" validate:"required"`
	ComponentType             string   `mapstructure:"componentType" json:"componentType" validate:"required"`
	ComponentSpec             string   `mapstructure:"componentSpec" json:"componentSpec" validate:"required"`
	Scopes                    []string `mapstructure:"scopes" json:"scopes" validate:"required"`
	OriginalComponentSpec     string   `mapstructure:"originalComponentSpec" json:"originalComponentSpec"`
	DesiredConfigRevisionName string   `mapstructure:"desiredConfigRevisionName" json:"desiredConfigRevisionName"` // optional

	// internal use only
	cluster           *cloudv1.Cluster
	tiDBCluster       *opv1alpha1.TidbCluster
	currentComponent  cloudv1.ClusterComponent
	desiredComponent  cloudv1.ClusterComponent
	originalComponent cloudv1.ClusterComponent
}

// newClusterComponent creates a new empty component spec instance based on the component type.
// It returns an error if the component type is not supported.
func newClusterComponent(componentType string) (cloudv1.ClusterComponent, error) {
	switch componentType {
	case string(cloudv1.ComponentPD):
		return &cloudv1.PD{}, nil
	case string(cloudv1.ComponentTiDB):
		return &cloudv1.TiDB{}, nil
	case string(cloudv1.ComponentTiProxy):
		return &cloudv1.TiProxy{}, nil
	case string(cloudv1.ComponentTiKV):
		return &cloudv1.TiKV{}, nil
	case string(cloudv1.ComponentTiFlash):
		return &cloudv1.TiFlash{}, nil
	default:
		return nil, fmt.Errorf("unknown component type: %s", componentType)
	}
}

var (
	_                 action.Action     = &scaleClusterAction{}
	_                 action.ActionPost = &scaleClusterAction{}
	scaleScopeApplier                   = map[string]func(ctx context.Context, wf *cloudv1.Workflow, step cloudv1.WorkflowStep, inputs *scaleClusterInputs) error{
		consts.WfScaleClusterScopeReplicas:       applyReplicasChanges,
		consts.WfScaleClusterScopeInstanceType:   applyInstanceTypeChanges,
		consts.WfScaleClusterScopeStorage:        applyStorageChanges,
		consts.WfScaleClusterScopeConfigRevision: applyConfigRevisionChanges,
	}

	// cancelScaleApplier maps scale scope names to their cancel handlers.
	// Only scopes that support cancellation are included in this map.
	cancelScaleApplier = map[string]func(ctx context.Context, wf *cloudv1.Workflow, step cloudv1.WorkflowStep, inputs *scaleClusterInputs, client, remoteCli client.Client) error{
		consts.WfScaleClusterScopeReplicas: cancelReplicasChanges,
	}

	// cancelPostScaleApplier maps scale scope names to their post-cancellation handler functions.
	// These handlers are executed after a cancellation operation completes.
	cancelPostScaleApplier = map[string]func(ctx context.Context, wf *cloudv1.Workflow, step cloudv1.WorkflowStep, inputs *scaleClusterInputs, localCli, remoteCli client.Client, pdAuth configmanager.PDAuth, pdServiceResolver PDServiceResolver, dryRun bool) error{
		consts.WfScaleClusterScopeReplicas: cancelReplicasPostHandle,
	}
)

func NewScaleClusterAction(cli client.Client, cachedProvider dataplane.EnvProviderManager, pdAuth configmanager.PDAuth, pdServiceResolver PDServiceResolver, dryRun bool) action.Action {
	return &scaleClusterAction{
		cli:               cli,
		cachedProvider:    cachedProvider,
		pdAuth:            pdAuth,
		pdServiceResolver: pdServiceResolver,
		dryRun:            dryRun,
	}
}

func (a *scaleClusterAction) Name() cloudv1.WorkflowActionName {
	return consts.WorkflowActionScaleCluster
}

func (a *scaleClusterAction) PreCheck(ctx context.Context, wf *cloudv1.Workflow, step cloudv1.WorkflowStep) error {
	inputs := &scaleClusterInputs{}
	if err := mapstructure.Decode(step.Inputs.Object, inputs); err != nil {
		return fmt.Errorf("decode inputs: %w", err)
	}
	if err := validate.Struct(inputs); err != nil {
		return fmt.Errorf("validate inputs: %w", err)
	}
	return nil
}

func (a *scaleClusterAction) Handle(ctx context.Context, wf *cloudv1.Workflow, step cloudv1.WorkflowStep) error {
	inputs := &scaleClusterInputs{}
	if err := mapstructure.Decode(step.Inputs.Object, inputs); err != nil {
		return fmt.Errorf("decode inputs: %w", err)
	}
	if err := a.fulfillInputs(ctx, wf, step, inputs); err != nil {
		return fmt.Errorf("fulfill inputs: %w", err)
	}
	for _, scope := range inputs.Scopes {
		if _, ok := scaleScopeApplier[scope]; !ok {
			return cloudv1.NewWFNonRetriableError(fmt.Errorf("unknown scale scope: %s", scope))
		}
	}

	for _, scope := range inputs.Scopes {
		applier := scaleScopeApplier[scope]
		if err := applier(ctx, wf, step, inputs); err != nil {
			return fmt.Errorf("apply scale scope %s: %w", scope, err)
		}
	}
	if shouldDisableGKEHeterogeneousForInstanceTypeScale(wf, inputs) {
		if err := a.deleteGKEHeterogeneousGatekeeperForInstanceTypeScale(ctx, inputs); err != nil {
			return fmt.Errorf("delete gke heterogeneous gatekeeper before instance type scale: %w", err)
		}
		clearGKEHeterogeneousIntentInScaleInputs(inputs)
	}
	return a.cli.Update(ctx, inputs.cluster)
}

func (a *scaleClusterAction) PostHandle(ctx context.Context, wf *cloudv1.Workflow, step cloudv1.WorkflowStep) error {
	inputs := &scaleClusterInputs{}
	if err := mapstructure.Decode(step.Inputs.Object, inputs); err != nil {
		return fmt.Errorf("decode inputs: %w", err)
	}
	if err := a.fulfillInputs(ctx, wf, step, inputs); err != nil {
		return fmt.Errorf("fulfill inputs: %w", err)
	}

	if err := inputs.cluster.FullSyncedWithDetail(); err != nil {
		return fmt.Errorf("not fully synced: %w", err)
	}
	return nil
}

func (a *scaleClusterAction) Cancel(ctx context.Context, wf *cloudv1.Workflow, step cloudv1.WorkflowStep, dryRun bool) error {
	inputs := &scaleClusterInputs{}

	if err := mapstructure.Decode(step.Inputs.Object, inputs); err != nil {
		return fmt.Errorf("decode inputs: %w", err)
	}
	if err := a.fulfillInputs(ctx, wf, step, inputs); err != nil {
		return fmt.Errorf("fulfill inputs: %w", err)
	}

	if inputs.cluster.FeatureEnabled(consts.ClusterFeatureUseTiDBOperatorV2) {
		return action.NewActionStepNotCancelableError(
			fmt.Errorf("cancel is not supported for TiDB Operator v2"),
		)
	}

	// Check if all scopes support cancellation
	for _, scope := range inputs.Scopes {
		if _, ok := cancelScaleApplier[scope]; !ok {
			return action.NewActionStepNotCancelableError(fmt.Errorf("scale scope %q does not support cancellation", scope))
		}
	}

	envProvider, err := a.cachedProvider.GetEnvProvider(ctx, inputs.cluster)
	if err != nil {
		return fmt.Errorf("get env provider: %w", err)
	}
	remoteCli := envProvider.GetClient()
	localCli := a.cli
	if dryRun {
		localCli = client.NewDryRunClient(localCli)
		remoteCli = client.NewDryRunClient(remoteCli)
	}

	// Apply cancellation for each scope.
	for _, scope := range inputs.Scopes {
		// safe to index directly because we validated all scopes above
		cancelApplier := cancelScaleApplier[scope]
		if err := cancelApplier(ctx, wf, step, inputs, localCli, remoteCli); err != nil {
			return err
		}
	}
	if dryRun {
		return nil
	}

	inputs.tiDBCluster.Spec.Paused = true
	if err := remoteCli.Update(ctx, inputs.tiDBCluster); err != nil {
		return fmt.Errorf("failed to update TiDBCluster: %w", err)
	}

	return localCli.Update(ctx, inputs.cluster)
}

func (a *scaleClusterAction) PostCancel(ctx context.Context, wf *cloudv1.Workflow, step cloudv1.WorkflowStep) error {
	inputs := &scaleClusterInputs{}
	if err := mapstructure.Decode(step.Inputs.Object, inputs); err != nil {
		return fmt.Errorf("decode inputs: %w", err)
	}
	if err := a.fulfillInputs(ctx, wf, step, inputs); err != nil {
		return fmt.Errorf("fulfill inputs: %w", err)
	}
	envProvider, err := a.cachedProvider.GetEnvProvider(ctx, inputs.cluster)
	if err != nil {
		return fmt.Errorf("get env provider: %w", err)
	}
	remoteCli := envProvider.GetClient()
	for _, scope := range inputs.Scopes {
		postCancelApplier, ok := cancelPostScaleApplier[scope]
		if !ok {
			return fmt.Errorf("scale scope %q does not support post-cancel", scope)
		}
		if err := postCancelApplier(ctx, wf, step, inputs, a.cli, remoteCli, a.pdAuth, a.pdServiceResolver, a.dryRun); err != nil {
			return fmt.Errorf("cancel post scale scope %s: %w", scope, err)
		}
	}

	inputs.tiDBCluster.Spec.Paused = false
	if err := remoteCli.Update(ctx, inputs.tiDBCluster); err != nil {
		return fmt.Errorf("failed to update TiDBCluster: %w", err)
	}

	if err := inputs.cluster.FullSyncedWithDetail(); err != nil {
		return fmt.Errorf("not fully synced: %w", err)
	}

	// Collect all cancelled steps and record them on cluster annotation.
	cancelledSteps := wf.StepsByPhase(cloudv1.WorkflowStepPhaseCancelEnding)
	// TODO: Consider whether there is a better implementation method
	inputs.cluster.MutateAnnotations(func(anno map[string]string) {
		if len(cancelledSteps) > 0 {
			sort.Strings(cancelledSteps)
			cancelledStepsStr := strings.Join(cancelledSteps, ",")
			anno[consts.AnnoClusterResourceStatusDiff] = cancelledStepsStr
		} else {
			delete(anno, consts.AnnoClusterResourceStatusDiff)
		}
	})

	return a.cli.Update(ctx, inputs.cluster)
}

func applyReplicasChanges(ctx context.Context, wf *cloudv1.Workflow, step cloudv1.WorkflowStep, inputs *scaleClusterInputs) error {
	logger := log.FromContext(ctx)
	logger.Info("scaling replicas", "from", inputs.currentComponent.GetReplicas(), "to", inputs.desiredComponent.GetReplicas())

	inputs.currentComponent.SetReplicas(inputs.desiredComponent.GetReplicas())
	return nil
}

// cancelReplicasChanges reverts the in-memory replicas of the current component back to the original value.
// The updated spec will be persisted when the cluster object is updated by the caller.
func cancelReplicasChanges(ctx context.Context, wf *cloudv1.Workflow, step cloudv1.WorkflowStep, inputs *scaleClusterInputs, client, remoteCli client.Client) error {
	logger := log.FromContext(ctx)
	currentReplicas := inputs.currentComponent.GetReplicas()
	originalReplicas := inputs.originalComponent.GetReplicas()
	desiredReplicas := inputs.desiredComponent.GetReplicas()

	if desiredReplicas > originalReplicas {
		logger.Info("cancel scale out is not supported")
		return action.NewActionStepNotCancelableError(fmt.Errorf("cancel scale out is not supported"))
	}

	logger.Info("cancel replicas changes", "from", currentReplicas, "to", originalReplicas)
	inputs.currentComponent.SetReplicas(originalReplicas)

	return nil
}

// syncDeleteSlotsAnnotation rebuilds the delete-slots annotation based on current pod state.
// It removes slots for pods that already exist, ensuring the annotation only contains slots for non-existent pods.
// This function is skipped for TiDB Operator v2 as it handles delete slots automatically.
func syncDeleteSlotsAnnotation(ctx context.Context, inputs *scaleClusterInputs, remoteCli client.Client, tc *opv1alpha1.TidbCluster) error {
	// Skip if using TiDB Operator v2, which automatically handles delete slots.
	if inputs.cluster != nil && inputs.cluster.FeatureEnabled(consts.ClusterFeatureUseTiDBOperatorV2) {
		return nil
	}
	logger := log.FromContext(ctx)

	componentType := cloudv1.ComponentType(inputs.ComponentType)
	currentSlots := common.GetDeleteSlots(tc, componentType)
	if currentSlots.Len() == 0 {
		logger.Info("no delete slots to sync", "component", componentType)
		return nil
	}

	// Get existing pods to determine which slots should be removed from the annotation
	pods, err := k8s.SelectPods(ctx, remoteCli, inputs.cluster.PrefixedId(), map[string]string{
		label.InstanceLabelKey:  consts.ClusterPrefix,
		label.ComponentLabelKey: string(componentType),
	})
	if err != nil {
		return fmt.Errorf("failed to list component pods: %w", err)
	}

	slotsToRemove := sets.New[int]()
	for _, pod := range pods {
		statefulPod, err := utils.NewStatefulPodFromName(pod.Name)
		if err != nil {
			return fmt.Errorf("failed to parse pod name: %w", err)
		}

		slotsToRemove.Insert(statefulPod.GetOrdinal())
	}

	// Calculate new slots by removing slots that have existing pods
	newSlots := currentSlots.Difference(slotsToRemove)
	if newSlots.Equal(currentSlots) {
		logger.Info("delete slots unchanged, skip update", "component", componentType, "slots", currentSlots)
		return nil
	}

	logger.Info("updating delete slots", "component", componentType, "currentSlots", currentSlots, "slotsToRemove", slotsToRemove, "newSlots", newSlots)
	if err := common.SetDeleteSlots(tc, componentType, newSlots); err != nil {
		return fmt.Errorf("failed to set delete slots for component %s: %w", componentType, err)
	}
	return nil
}

func applyInstanceTypeChanges(ctx context.Context, wf *cloudv1.Workflow, step cloudv1.WorkflowStep, inputs *scaleClusterInputs) error {
	logger := log.FromContext(ctx)
	logger.Info("scaling instance type", "from", inputs.currentComponent.GetInstanceType(), "to", inputs.desiredComponent.GetInstanceType())

	if inputs.currentComponent.GetInstanceType() == inputs.desiredComponent.GetInstanceType() {
		return nil
	}
	disablesGKEHeterogeneous := scaleWorkflowDisablesGKEHeterogeneous(wf)
	if inputs.cluster != nil && inputs.cluster.GetProviderType() == cloudv1.ProviderTypeGKE {
		componentType := inputs.currentComponent.GetComponentType()
		blocked := componentType == cloudv1.ComponentTiKV && inputs.cluster.TiKVGKEComputeClassIntentConfigured()
		blocked = blocked || componentType == cloudv1.ComponentTiDB && inputs.cluster.TiDBGKEComputeClassIntentConfigured(inputs.GroupID)
		if blocked && !disablesGKEHeterogeneous {
			return cloudv1.NewWFNonRetriableError(errors.Errorf("cannot change %s instance type while GKE heterogeneous machine deploy is enabled; disable heterogeneous deploy first", componentType))
		}
	}

	inputs.currentComponent.SetInstanceType(inputs.desiredComponent.GetInstanceType())
	return nil
}

func scaleWorkflowDisablesGKEHeterogeneous(wf *cloudv1.Workflow) bool {
	return wf != nil && strings.EqualFold(strings.TrimSpace(wf.Spec.Inputs[consts.WfInputKeyScaleClusterDisableGKEHeterogeneous]), "true")
}

func scaleInputsIncludeScope(inputs *scaleClusterInputs, scope string) bool {
	return inputs != nil && sets.New[string](inputs.Scopes...).Has(scope)
}

func scaleInputsChangeInstanceType(inputs *scaleClusterInputs) bool {
	return inputs != nil && inputs.originalComponent != nil && inputs.desiredComponent != nil &&
		strings.TrimSpace(inputs.originalComponent.GetInstanceType()) != strings.TrimSpace(inputs.desiredComponent.GetInstanceType())
}

func shouldDisableGKEHeterogeneousForInstanceTypeScale(wf *cloudv1.Workflow, inputs *scaleClusterInputs) bool {
	if !scaleWorkflowDisablesGKEHeterogeneous(wf) ||
		!scaleInputsIncludeScope(inputs, consts.WfScaleClusterScopeInstanceType) ||
		!scaleInputsChangeInstanceType(inputs) {
		return false
	}
	_, ok := scaleInputsGKEHeterogeneousIntentAnnotationKey(inputs)
	return ok
}

func (a *scaleClusterAction) deleteGKEHeterogeneousGatekeeperForInstanceTypeScale(ctx context.Context, inputs *scaleClusterInputs) error {
	envProvider, err := a.cachedProvider.GetEnvProvider(ctx, inputs.cluster)
	if err != nil {
		return fmt.Errorf("get env provider: %w", err)
	}
	componentType := cloudv1.ComponentType(inputs.ComponentType)
	tidbGroupID := ""
	if componentType == cloudv1.ComponentTiDB {
		tidbGroupID = scaleInputsTiDBGroupID(inputs)
	}
	return gkeheterogeneous.DeleteGatekeeperAssignForComponentIfExists(ctx, envProvider.GetClient(), inputs.cluster.Name, componentType, tidbGroupID)
}

func scaleInputsGKEHeterogeneousIntentAnnotationKey(inputs *scaleClusterInputs) (string, bool) {
	if inputs == nil || inputs.cluster == nil || inputs.cluster.GetProviderType() != cloudv1.ProviderTypeGKE {
		return "", false
	}

	switch cloudv1.ComponentType(inputs.ComponentType) {
	case cloudv1.ComponentTiKV:
		return cloudv1.AnnoKeyGKEComputeClassForTiKVBindings, true
	case cloudv1.ComponentTiDB:
		return inputs.cluster.TiDBGKEComputeClassBindingsAnnoKey(scaleInputsTiDBGroupID(inputs)), true
	default:
		return "", false
	}
}

func scaleInputsTiDBGroupID(inputs *scaleClusterInputs) string {
	if inputs == nil {
		return ""
	}
	groupID := strings.TrimSpace(inputs.GroupID)
	if groupID == "" && inputs.cluster != nil {
		groupID = strings.TrimSpace(inputs.cluster.GetDefaultGroupID())
	}
	return groupID
}

func clearGKEHeterogeneousIntentInScaleInputs(inputs *scaleClusterInputs) bool {
	key, ok := scaleInputsGKEHeterogeneousIntentAnnotationKey(inputs)
	if !ok {
		return false
	}
	deleted := false
	inputs.cluster.MutateAnnotations(func(annotations map[string]string) {
		if _, ok := annotations[key]; !ok {
			return
		}
		delete(annotations, key)
		deleted = true
	})
	return deleted
}

func ignoreOk[T any](value T, _ bool) T {
	return value
}

func applyStorageChanges(ctx context.Context, wf *cloudv1.Workflow, step cloudv1.WorkflowStep, inputs *scaleClusterInputs) error {
	logger := log.FromContext(ctx)

	desiredComp := inputs.desiredComponent
	currentComp := inputs.currentComponent
	logger.Info("scaling storage",
		"fromSize", ignoreOk(currentComp.GetStorageSize()), "to", ignoreOk(desiredComp.GetStorageSize()),
		"fromStorageClass", ignoreOk(currentComp.GetStorageClassName()), "to", ignoreOk(desiredComp.GetStorageClassName()),
		"fromAdditionalStorages", currentComp.GetAdditionalStorages(), "to", desiredComp.GetAdditionalStorages(),
	)
	size, ok := desiredComp.GetStorageSize()
	if ok {
		currentComp.SetStorageSize(*size)
	}
	sc, ok := desiredComp.GetStorageClassName()
	if ok {
		currentComp.SetStorageClassName(sc)
	}
	currentComp.SetAdditionalStorages(desiredComp.GetAdditionalStorages())

	return nil
}

func applyConfigRevisionChanges(ctx context.Context, wf *cloudv1.Workflow, step cloudv1.WorkflowStep, inputs *scaleClusterInputs) error {
	logger := log.FromContext(ctx)

	if inputs.DesiredConfigRevisionName == "" {
		return nil
	}
	logger.Info("updating config revision", "from", inputs.cluster.Spec.ConfigRevisionRef.Name, "to", inputs.DesiredConfigRevisionName)
	inputs.cluster.Spec.ConfigRevisionRef = &cloudv1.Reference{
		Name: inputs.DesiredConfigRevisionName,
	}
	return nil
}

// fulfillInputs populates the internal fields of scaleClusterInputs by fetching the cluster
// and unmarshaling the component specifications from YAML strings.
func (a *scaleClusterAction) fulfillInputs(ctx context.Context, wf *cloudv1.Workflow, step cloudv1.WorkflowStep, inputs *scaleClusterInputs) error {
	logger := log.FromContext(ctx)
	cluster, err := getCluster(ctx, a.cli, client.ObjectKey{Namespace: inputs.Namespace, Name: inputs.ClusterID})
	if err != nil {
		return fmt.Errorf("get cluster: %w", err)
	}
	desiredComp, err := newClusterComponent(inputs.ComponentType)
	if err != nil {
		return err
	}
	originalComp, err := newClusterComponent(inputs.ComponentType)
	if err != nil {
		return err
	}

	inputs.cluster = cluster

	envProvider, err := a.cachedProvider.GetEnvProvider(ctx, inputs.cluster)
	if err != nil {
		return fmt.Errorf("get env provider: %w", err)
	}
	remoteCli := envProvider.GetClient()
	if !inputs.cluster.FeatureEnabled(consts.ClusterFeatureUseTiDBOperatorV2) {
		tc, err := getTiDBCluster(ctx, inputs.cluster, remoteCli, inputs.GroupID)
		if err != nil {
			return fmt.Errorf("get TiDBCluster: %w", err)
		}
		inputs.tiDBCluster = tc
	}

	if err := yaml.Unmarshal([]byte(inputs.OriginalComponentSpec), originalComp); err != nil {
		return fmt.Errorf("failed to unmarshal original componentSpec: %w", err)
	}
	if err := yaml.Unmarshal([]byte(inputs.ComponentSpec), desiredComp); err != nil {
		return fmt.Errorf("failed to unmarshal componentSpec: %w", err)
	}
	logger.Info("desired component", "comp", desiredComp, "workflow", wf)
	currentComp, err := cluster.GetComponentSpecInGroup(inputs.GroupID, cloudv1.ComponentType(inputs.ComponentType))
	if err != nil {
		return fmt.Errorf("failed to get current component spec in group: %w", err)
	}
	inputs.currentComponent = currentComp
	inputs.desiredComponent = desiredComp
	inputs.originalComponent = originalComp
	return nil
}

// cleanupScaleInStoreState cleans up store-related state when canceling a scale-in operation.
// This function performs the following steps:
//  1. Cancels delete operations for all offline stores via PD API
//  2. Removes the evict-leader-scheduler that was used during scale-in
//
// Note: This function is skipped for TiDB Operator v2, which handles store operations automatically.
func cleanupScaleInStoreState(ctx context.Context, inputs *scaleClusterInputs, remoteCli client.Client, pdAuth configmanager.PDAuth, pdServiceResolver PDServiceResolver, dryRun bool) error {
	logger := log.FromContext(ctx)

	// Skip for TiDB Operator v2, which handles store operations automatically
	if inputs.cluster.FeatureEnabled(consts.ClusterFeatureUseTiDBOperatorV2) {
		logger.V(1).Info("skipping offline store cleanup for TiDB Operator v2")
		return nil
	}

	componentType := cloudv1.ComponentType(inputs.ComponentType)
	offlineStoreIDs, err := getStoreIDsByState(ctx, inputs.cluster, componentType, remoteCli, opv1alpha1.TiKVStateOffline)
	if err != nil {
		return fmt.Errorf("failed to get offline store IDs: %w", err)
	}
	if len(offlineStoreIDs) == 0 {
		logger.Info("no offline stores found, skipping cleanup",
			"cluster", inputs.ClusterID,
			"component", componentType)
		return nil
	}

	pdClient, err := getPDClient(ctx, inputs.cluster, remoteCli, pdAuth, pdServiceResolver, dryRun)
	if err != nil {
		return fmt.Errorf("failed to get PD client: %w", err)
	}

	logger.Info("canceling delete operations for offline stores",
		"cluster", inputs.ClusterID,
		"component", componentType,
		"storeCount", len(offlineStoreIDs),
		"storeIDs", offlineStoreIDs)

	// Cancel delete operations for all offline stores
	for _, storeID := range offlineStoreIDs {
		if err := pdClient.CancelDeleteStore(ctx, storeID); err != nil {
			return fmt.Errorf("failed to cancel deletion for store %s: %w", storeID, err)
		}
		logger.V(1).Info("canceled delete operation for store", "storeID", storeID)
	}

	// Remove evict-leader-scheduler after canceling all store deletions.
	// Note: RemoveEvictLeaderScheduler removes all evict-leader-schedulers globally,
	// not just for specific stores. We verify removal after attempting to remove,
	// as PD may return non-standard error responses.
	if err := removeAndVerifyEvictLeaderScheduler(ctx, pdClient, logger); err != nil {
		return fmt.Errorf("failed to remove evict-leader-scheduler: %w", err)
	}

	logger.Info("successfully canceled offline store deletions",
		"cluster", inputs.ClusterID,
		"component", componentType,
		"storeCount", len(offlineStoreIDs))
	return nil
}

// removeAndVerifyEvictLeaderScheduler removes the evict-leader-scheduler and verifies it was removed.
// PD may return an error with "scheduler not found" in the body, which is not a standard response.
// This verification ensures the scheduler is truly removed.
// TODO: Remove this workaround when PD returns standard JSON responses.
func removeAndVerifyEvictLeaderScheduler(ctx context.Context, pdClient configmanager.DynamicPDClient, logger logr.Logger) error {
	// Attempt to remove the evict-leader-scheduler.
	// We continue with verification even if removal fails, as the verification step will confirm the actual state.
	if err := pdClient.RemoveEvictLeaderScheduler(ctx); err != nil {
		logger.Error(err, "failed to remove evict-leader-scheduler, proceeding with verification")
	}

	// Verify that all evict-leader-schedulers have been removed
	evictLeaderSchedulers, err := pdClient.GetEvictLeaderSchedulers(ctx)
	if err != nil {
		return fmt.Errorf("failed to get evict-leader-schedulers for verification: %w", err)
	}

	if len(evictLeaderSchedulers) > 0 {
		return fmt.Errorf("evict-leader-scheduler removal failed: %d scheduler(s) still exist: %v",
			len(evictLeaderSchedulers), evictLeaderSchedulers)
	}

	logger.Info("successfully verified all evict-leader-schedulers have been removed")
	return nil
}

// getComponentReplicas gets the replicas count for a component from TiDBCluster spec.
// Returns an error if the component is nil or unsupported.
func getComponentReplicas(tc *opv1alpha1.TidbCluster, componentType cloudv1.ComponentType) (int32, error) {
	switch componentType {
	case cloudv1.ComponentPD:
		if tc.Spec.PD == nil {
			return 0, fmt.Errorf("PD component is nil in TiDBCluster")
		}
		return tc.Spec.PD.Replicas, nil
	case cloudv1.ComponentTiDB:
		if tc.Spec.TiDB == nil {
			return 0, fmt.Errorf("TiDB component is nil in TiDBCluster")
		}
		return tc.Spec.TiDB.Replicas, nil
	case cloudv1.ComponentTiKV:
		if tc.Spec.TiKV == nil {
			return 0, fmt.Errorf("TiKV component is nil in TiDBCluster")
		}
		return tc.Spec.TiKV.Replicas, nil
	case cloudv1.ComponentTiFlash:
		if tc.Spec.TiFlash == nil {
			return 0, fmt.Errorf("TiFlash component is nil in TiDBCluster")
		}
		return tc.Spec.TiFlash.Replicas, nil
	case cloudv1.ComponentTiProxy:
		if tc.Spec.TiProxy == nil {
			return 0, fmt.Errorf("TiProxy component is nil in TiDBCluster")
		}
		return tc.Spec.TiProxy.Replicas, nil
	default:
		return 0, fmt.Errorf("unsupported component type: %s", componentType)
	}
}

// cancelReplicasPostHandle performs post-cancellation cleanup after canceling a replicas scaling operation.
// It verifies replicas match, stops offline stores for store components, and syncs delete slots.
// Note: localCli parameter is currently unused but kept for consistency with other handler signatures.
func cancelReplicasPostHandle(ctx context.Context, wf *cloudv1.Workflow, step cloudv1.WorkflowStep, inputs *scaleClusterInputs, localCli, remoteCli client.Client, pdAuth configmanager.PDAuth, pdServiceResolver PDServiceResolver, dryRun bool) error {
	logger := log.FromContext(ctx)
	// Confirm that the number of replicas in the tc has been updated to the previous value
	componentType := cloudv1.ComponentType(inputs.ComponentType)
	originalReplicas := inputs.originalComponent.GetReplicas()
	actualReplicas, err := getComponentReplicas(inputs.tiDBCluster, componentType)
	if err != nil {
		return fmt.Errorf("get component replicas from TiDBCluster: %w", err)
	}
	if originalReplicas != int(actualReplicas) {
		return fmt.Errorf("TiDBCluster %s replicas (%d) does not match original replicas (%d)", componentType, actualReplicas, originalReplicas)
	}

	// Cleanup store components (TiKV/TiFlash) after cancellation.
	// When canceling a scale-in operation, we need to:
	// 1. Stop offline stores to prevent them from being deleted
	// 2. Remove evict-leader-scheduler that was used during scale-in
	if componentType == cloudv1.ComponentTiKV || componentType == cloudv1.ComponentTiFlash {
		// Cleanup store state (cancel deletions and remove scheduler)
		if err := cleanupScaleInStoreState(ctx, inputs, remoteCli, pdAuth, pdServiceResolver, dryRun); err != nil {
			logger.Error(err, "failed to cleanup scale-in store state", "component", componentType)
			return fmt.Errorf("cleanup scale-in store state for %s: %w", componentType, err)
		}
	}

	// Sync delete slots annotation based on current pod state
	if err := syncDeleteSlotsAnnotation(ctx, inputs, remoteCli, inputs.tiDBCluster); err != nil {
		logger.Error(err, "failed to sync delete slots annotation")
		return fmt.Errorf("sync delete slots annotation: %w", err)
	}
	return nil
}
