package server

import (
	"context"
	"fmt"
	"strconv"

	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	"github.com/tidbcloud/cluster-service-ng/pkg/k8s"
	"github.com/tidbcloud/cluster-service-ng/pkg/metricsprovider"
	schemav1 "github.com/tidbcloud/cluster-service-ng/pkg/storage/etcd/schema/v1"
	"gomodules.xyz/jsonpatch/v3"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	autoscalingv2 "k8s.io/api/autoscaling/v2"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/resource"
	"k8s.io/apimachinery/pkg/labels"
	client "sigs.k8s.io/controller-runtime/pkg/client"
)

func (s *serverImpl) SetScaleOutConfig(
	ctx context.Context,
	sharedPoolId string,
	groupId string,
	componentType entity.ComponentType,
	config *entity.ScaleOutConfig,
) error {
	if sharedPoolId == "" {
		return status.Errorf(codes.InvalidArgument, "sharedPoolId is required")
	}
	if config.MinCount > config.MaxCount {
		return status.Errorf(codes.InvalidArgument, "minCount can not larger than maxCount")
	}
	unlockFunc, err := s.lockTiKVRemoteWorkerGuardedMutation(ctx, sharedPoolId, componentType)
	if err != nil {
		return err
	}
	defer unlockFunc()

	sharedPool, err := s.sharedPoolRepo.Load(ctx, sharedPoolId, false)
	if err != nil {
		return err
	}
	if err := guardSharedPoolMutation(sharedPool, sharedPoolMutationBlockedWhilePaused); err != nil {
		return err
	}
	if componentType == apiv1.ComponentType_COPROCESSOR_WORKER {
		if err := ensureCoprocessorWorkerComponentExists(sharedPool); err != nil {
			return err
		}
	}
	if isTiKVRemoteWorkerGuardedComponentType(componentType) {
		if err := s.mergeSharedPoolStatusFromDagTask(ctx, sharedPool); err != nil {
			return status.Errorf(codes.Internal, "failed to merge sharedPool status from dag task: %s", err)
		}
		if err := s.guardTiKVRemoteWorkersTopologyStable(ctx, sharedPool); err != nil {
			return err
		}
	}
	remoteCli, err := s.dataplaneManager.GetClient(sharedPool.GetDataplaneName())
	if err != nil {
		return status.Errorf(codes.Internal, "failed to get remote client for shared pool: %s", err)
	}
	namespace := "tidb" + sharedPool.GetId()
	hpas := new(autoscalingv2.HorizontalPodAutoscalerList)
	switch componentType {
	case apiv1.ComponentType_TIDB:
		if err = remoteCli.List(ctx, hpas, &client.ListOptions{
			Namespace: namespace,
			LabelSelector: labels.SelectorFromSet(map[string]string{
				schemav1.LabelPingcapGroupKey: fmt.Sprintf("db-%s", groupId),
				schemav1.LabelComponentKey:    "tidb",
			}),
		}); err != nil {
			return status.Errorf(codes.Internal, "failed to list tidb hpas: %s", err)
		}
	case apiv1.ComponentType_TIDB_WORKER:
		if err = remoteCli.List(ctx, hpas, &client.ListOptions{
			Namespace: namespace,
			LabelSelector: labels.SelectorFromSet(map[string]string{
				schemav1.LabelComponentKey: "worker-tidb",
			}),
		}); err != nil {
			return status.Errorf(codes.Internal, "failed to list tidb-worker hpas: %s", err)
		}
	case apiv1.ComponentType_TIKV_WORKER:
		if err = remoteCli.List(ctx, hpas, &client.ListOptions{
			Namespace: namespace,
			LabelSelector: labels.SelectorFromSet(map[string]string{
				schemav1.LabelComponentKey: entity.TiKVWorkerType,
			}),
		}); err != nil {
			return status.Errorf(codes.Internal, "failed to list tikv-worker hpas: %s", err)
		}
	case apiv1.ComponentType_COPROCESSOR_WORKER:
		if err = remoteCli.List(ctx, hpas, &client.ListOptions{
			Namespace: namespace,
			LabelSelector: labels.SelectorFromSet(map[string]string{
				schemav1.LabelComponentKey: entity.CoprocessorWorkerType,
			}),
		}); err != nil {
			return status.Errorf(codes.Internal, "failed to list coprocessor-worker hpas: %s", err)
		}
	default:
		return status.Errorf(codes.Unimplemented, "%s does not support hpa", componentType)
	}
	if len(hpas.Items) == 0 {
		return status.Errorf(codes.Internal, "%s has no hpa", componentType)
	}
	hpa := &hpas.Items[0]

	if config.Thresholds != nil {
		annotationChanged, err := applyTiDBScaleInThreshold(hpa, componentType, config.Thresholds)
		if err != nil {
			return err
		}
		if annotationChanged || needsEnableTiDBScaleInScaleDown(hpa, componentType, config.Thresholds) {
			return updateScaleOutConfigWithRetry(ctx, remoteCli, hpa, componentType, config)
		}
	}

	ops := make([]*jsonpatch.Operation, 0)

	if config.Thresholds != nil {
		thresholdsOps, err := opsForThresholds(&hpa.Spec, config.Thresholds)
		if err != nil {
			return err
		}
		ops = append(ops, thresholdsOps...)
	}

	ops = append(ops,
		&jsonpatch.Operation{
			Operation: "replace",
			Path:      "spec/minReplicas",
			Value:     config.MinCount,
		},
		&jsonpatch.Operation{
			Operation: "replace",
			Path:      "spec/maxReplicas",
			Value:     config.MaxCount,
		})

	if config.Step != nil {
		stepOps, err := opsForStep(&hpa.Spec, config.Step)
		if err != nil {
			return err
		}
		ops = append(ops, stepOps...)
	}

	patch, err := k8s.ConstructPatch(ops...)
	if err != nil {
		return status.Errorf(codes.Internal, "construct patch failed: %v", err)
	}
	if err := remoteCli.Patch(ctx, hpa, patch); err != nil {
		return status.Errorf(codes.Internal, "failed to patch hpa: %s", err)
	}

	return nil
}

func applyTiDBScaleInThreshold(
	hpa *autoscalingv2.HorizontalPodAutoscaler,
	componentType entity.ComponentType,
	thresholds *apiv1.ScaleOutThresholds,
) (bool, error) {
	if thresholds == nil || thresholds.CpuScaleInUtilizationPercentage == 0 {
		return false, nil
	}
	if componentType != apiv1.ComponentType_TIDB {
		return false, status.Errorf(codes.InvalidArgument, "cpuScaleInUtilizationPercentage is only supported for tidb")
	}
	if thresholds.CpuScaleInUtilizationPercentage > 100 {
		return false, status.Errorf(codes.InvalidArgument, "cpuScaleInUtilizationPercentage should be in range [1,100]")
	}

	changed := false
	if hpa.Annotations == nil {
		hpa.Annotations = make(map[string]string)
	}
	value := strconv.FormatUint(thresholds.CpuScaleInUtilizationPercentage, 10)
	if hpa.Annotations[metricsprovider.AnnotationTiDBCPUScaleInThreshold] != value {
		hpa.Annotations[metricsprovider.AnnotationTiDBCPUScaleInThreshold] = value
		changed = true
	}
	if ensureCPUScaleInGateForHPA(hpa, metricsprovider.TiDBCPUScaleInGateConfig) {
		changed = true
	}
	return changed, nil
}

func needsEnableTiDBScaleInScaleDown(
	hpa *autoscalingv2.HorizontalPodAutoscaler,
	componentType entity.ComponentType,
	thresholds *apiv1.ScaleOutThresholds,
) bool {
	if componentType != apiv1.ComponentType_TIDB || thresholds == nil || thresholds.CpuScaleInUtilizationPercentage == 0 {
		return false
	}
	if hpa == nil || hpa.Spec.Behavior == nil || hpa.Spec.Behavior.ScaleDown == nil || hpa.Spec.Behavior.ScaleDown.SelectPolicy == nil {
		return true
	}
	return *hpa.Spec.Behavior.ScaleDown.SelectPolicy == autoscalingv2.DisabledPolicySelect
}

func applyConfigInPlace(hpa *autoscalingv2.HorizontalPodAutoscaler, componentType entity.ComponentType, config *entity.ScaleOutConfig) error {
	if config.Thresholds != nil {
		if err := applyThresholdsInPlace(&hpa.Spec, config.Thresholds); err != nil {
			return err
		}
	}
	minCount := int32(config.MinCount)
	hpa.Spec.MinReplicas = &minCount
	hpa.Spec.MaxReplicas = int32(config.MaxCount)
	if config.Step != nil {
		if err := applyStepInPlace(&hpa.Spec, config.Step); err != nil {
			return err
		}
	}
	// When enabling TiDB scale-in threshold, ensure scaleDown is not Disabled
	// so the gate can actually allow scale-in.
	if componentType == apiv1.ComponentType_TIDB {
		if hpa.Spec.Behavior == nil {
			hpa.Spec.Behavior = &autoscalingv2.HorizontalPodAutoscalerBehavior{}
		}
		if hpa.Spec.Behavior.ScaleDown == nil {
			hpa.Spec.Behavior.ScaleDown = &autoscalingv2.HPAScalingRules{}
		}
		if hpa.Spec.Behavior.ScaleDown.SelectPolicy == nil || *hpa.Spec.Behavior.ScaleDown.SelectPolicy == autoscalingv2.DisabledPolicySelect {
			minPolicy := autoscalingv2.MinChangePolicySelect
			hpa.Spec.Behavior.ScaleDown.SelectPolicy = &minPolicy
		}
	}
	return nil
}

func applyThresholdsInPlace(spec *autoscalingv2.HorizontalPodAutoscalerSpec, thresholds *apiv1.ScaleOutThresholds) error {
	if value := thresholds.CpuUtilizationPercentage; value != 0 {
		index := -1
		for i, m := range spec.Metrics {
			if m.Type == autoscalingv2.ResourceMetricSourceType && m.Resource != nil && m.Resource.Name == corev1.ResourceCPU {
				index = i
				break
			}
		}
		if index == -1 {
			return status.Errorf(codes.FailedPrecondition, "resource metric cpu is not found")
		}
		v := int32(value)
		spec.Metrics[index].Resource.Target.AverageUtilization = &v
	}
	if value := thresholds.MemoryUtilizationPercentage; value != 0 {
		index := -1
		for i, m := range spec.Metrics {
			if m.Type == autoscalingv2.ResourceMetricSourceType && m.Resource != nil && m.Resource.Name == corev1.ResourceMemory {
				index = i
				break
			}
			if m.Type == autoscalingv2.ObjectMetricSourceType &&
				m.Object != nil &&
				m.Object.Metric.Name == metricsprovider.MetricNameExtAvgMemoryUtilization {
				index = i
				break
			}
		}
		if index == -1 {
			return status.Errorf(codes.FailedPrecondition, "resource metric memory is not found")
		}
		if spec.Metrics[index].Type == autoscalingv2.ResourceMetricSourceType {
			v := int32(value)
			spec.Metrics[index].Resource.Target.AverageUtilization = &v
		} else {
			spec.Metrics[index].Object.Target.Value = resource.NewQuantity(int64(value), resource.DecimalSI)
		}
	}
	return nil
}

func applyStepInPlace(spec *autoscalingv2.HorizontalPodAutoscalerSpec, step *apiv1.ScaleOutStep) error {
	if spec.Behavior == nil {
		spec.Behavior = &autoscalingv2.HorizontalPodAutoscalerBehavior{}
	}
	if spec.Behavior.ScaleUp == nil {
		spec.Behavior.ScaleUp = &autoscalingv2.HPAScalingRules{}
	}
	if len(spec.Behavior.ScaleUp.Policies) > 1 {
		return status.Errorf(codes.FailedPrecondition, "behavior scaleUp has more than one policies")
	}
	if value := step.Count; value != 0 {
		index := -1
		for i, p := range spec.Behavior.ScaleUp.Policies {
			if p.Type == autoscalingv2.PodsScalingPolicy {
				index = i
				break
			}
		}
		if index == -1 {
			return status.Errorf(codes.FailedPrecondition, "PodsScalingPolicy is not found in scaleUp behavior")
		}
		spec.Behavior.ScaleUp.Policies[index].Value = int32(value)
	}
	if value := step.Percent; value != 0 {
		index := -1
		for i, p := range spec.Behavior.ScaleUp.Policies {
			if p.Type == autoscalingv2.PercentScalingPolicy {
				index = i
				break
			}
		}
		if index == -1 {
			return status.Errorf(codes.FailedPrecondition, "PercentScalingPolicy is not found in scaleUp behavior")
		}
		spec.Behavior.ScaleUp.Policies[index].Value = int32(value)
	}
	return nil
}

func opsForThresholds(spec *autoscalingv2.HorizontalPodAutoscalerSpec, thresholds *apiv1.ScaleOutThresholds) ([]*jsonpatch.Operation, error) {
	ops := make([]*jsonpatch.Operation, 0)
	if value := thresholds.CpuUtilizationPercentage; value != 0 {
		index := -1
		for i, m := range spec.Metrics {
			if m.Type == autoscalingv2.ResourceMetricSourceType && m.Resource != nil && m.Resource.Name == corev1.ResourceCPU {
				index = i
				break
			}
		}
		if index == -1 {
			return nil, status.Errorf(codes.FailedPrecondition, "resource metric cpu is not found")
		}
		ops = append(ops, &jsonpatch.Operation{
			Operation: "replace",
			Path:      fmt.Sprintf("spec/metrics/%v/resource/target/averageUtilization", index),
			Value:     value,
		})
	}
	if value := thresholds.MemoryUtilizationPercentage; value != 0 {
		index := -1
		path := ""
		for i, m := range spec.Metrics {
			if m.Type == autoscalingv2.ResourceMetricSourceType && m.Resource != nil && m.Resource.Name == corev1.ResourceMemory {
				index = i
				path = fmt.Sprintf("spec/metrics/%v/resource/target/averageUtilization", index)
				break
			}
			if m.Type == autoscalingv2.ObjectMetricSourceType &&
				m.Object != nil &&
				m.Object.Metric.Name == metricsprovider.MetricNameExtAvgMemoryUtilization {
				index = i
				path = fmt.Sprintf("spec/metrics/%v/object/target/value", index)
				break
			}
		}
		if index == -1 {
			return nil, status.Errorf(codes.FailedPrecondition, "resource metric memory is not found")
		}
		ops = append(ops, &jsonpatch.Operation{
			Operation: "replace",
			Path:      path,
			Value:     value,
		})
	}
	return ops, nil
}

func opsForStep(spec *autoscalingv2.HorizontalPodAutoscalerSpec, step *apiv1.ScaleOutStep) ([]*jsonpatch.Operation, error) {
	ops := make([]*jsonpatch.Operation, 0)
	if value := step.Count; value != 0 {
		if spec.Behavior.ScaleUp == nil {
			return nil, status.Errorf(codes.FailedPrecondition, "behavior scaleUp is nil")
		}
		if len(spec.Behavior.ScaleUp.Policies) > 1 {
			return nil, status.Errorf(codes.FailedPrecondition, "behavior scaleUp has more than one policies")
		}

		index := -1
		for i, p := range spec.Behavior.ScaleUp.Policies {
			if p.Type == autoscalingv2.PodsScalingPolicy {
				index = i
				break
			}
		}
		if index == -1 {
			return nil, status.Errorf(codes.FailedPrecondition, "PodsScalingPolicy is not found in scaleUp behavior")
		}
		ops = append(ops, &jsonpatch.Operation{
			Operation: "replace",
			Path:      fmt.Sprintf("spec/behavior/scaleUp/policies/%v/value", index),
			Value:     value,
		})
	}

	if value := step.Percent; value != 0 {
		if spec.Behavior.ScaleUp == nil {
			return nil, status.Errorf(codes.FailedPrecondition, "behavior scaleUp is nil")
		}
		if len(spec.Behavior.ScaleUp.Policies) > 1 {
			return nil, status.Errorf(codes.FailedPrecondition, "behavior scaleUp has more than one policies")
		}

		index := -1
		for i, p := range spec.Behavior.ScaleUp.Policies {
			if p.Type == autoscalingv2.PercentScalingPolicy {
				index = i
				break
			}
		}
		if index == -1 {
			return nil, status.Errorf(codes.FailedPrecondition, "PercentScalingPolicy is not found in scaleUp behavior")
		}
		ops = append(ops, &jsonpatch.Operation{
			Operation: "replace",
			Path:      fmt.Sprintf("spec/behavior/scaleUp/policies/%v/value", index),
			Value:     value,
		})
	}
	return ops, nil
}

func updateScaleOutConfigWithRetry(
	ctx context.Context,
	remoteCli client.Client,
	hpa *autoscalingv2.HorizontalPodAutoscaler,
	componentType entity.ComponentType,
	config *entity.ScaleOutConfig,
) error {
	for i := 0; i < 3; i++ {
		if config.Thresholds != nil {
			if _, err := applyTiDBScaleInThreshold(hpa, componentType, config.Thresholds); err != nil {
				return err
			}
		}
		if err := applyConfigInPlace(hpa, componentType, config); err != nil {
			return err
		}
		if err := remoteCli.Update(ctx, hpa); err != nil {
			if !apierrors.IsConflict(err) {
				return status.Errorf(codes.Internal, "failed to update hpa: %s", err)
			}
			if err := remoteCli.Get(ctx, client.ObjectKeyFromObject(hpa), hpa); err != nil {
				return status.Errorf(codes.Internal, "failed to re-fetch hpa for retry: %s", err)
			}
			continue
		}
		return nil
	}
	return status.Errorf(codes.Aborted, "failed to update hpa after retries")
}