package server

import (
	"context"
	"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/autoscale/resolve"
	"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"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/resource"
	"k8s.io/apimachinery/pkg/labels"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

// getSelectPolicy resolves the HPA SelectPolicy for the legacy write
// path (SetHPAScaleBehavior / writeTidbGroupHPABehavior). It is a thin
// converter that maps the local (isScaleUp bool) signature onto the
// shared resolve.SelectPolicyFromComponentState helper, which is the
// single source of truth also used by the bundle-aware reconcile path
// (resolve.ResolveHPABehaviorFromBundle). Centralising the translation
// in resolve keeps the Disabled→MinChange flip for SCALE_IN on
// Disabled-baseline components (TIPROXY) consistent across both paths
// — see design doc §HPA Behavior Assembly Rule.
func getSelectPolicy(enabled bool, componentType entity.ComponentType, isScaleUp bool) autoscalingv2.ScalingPolicySelect {
	direction := apiv1.AutoscalePolicy_SCALE_IN
	if isScaleUp {
		direction = apiv1.AutoscalePolicy_SCALE_OUT
	}
	return resolve.SelectPolicyFromComponentState(componentType, direction, enabled)
}

func applyScaleBehaviorInPlace(
	spec *autoscalingv2.HorizontalPodAutoscalerSpec,
	componentType entity.ComponentType,
	scaleUpEnabled *bool,
	scaleDownEnabled *bool,
) bool {
	if spec.Behavior == nil {
		spec.Behavior = &autoscalingv2.HorizontalPodAutoscalerBehavior{}
	}

	changed := false

	if scaleUpEnabled != nil {
		if spec.Behavior.ScaleUp == nil {
			spec.Behavior.ScaleUp = &autoscalingv2.HPAScalingRules{}
		}
		newPolicy := getSelectPolicy(*scaleUpEnabled, componentType, true)
		if spec.Behavior.ScaleUp.SelectPolicy == nil || *spec.Behavior.ScaleUp.SelectPolicy != newPolicy {
			spec.Behavior.ScaleUp.SelectPolicy = &newPolicy
			changed = true
		}
	}

	if scaleDownEnabled != nil {
		if spec.Behavior.ScaleDown == nil {
			spec.Behavior.ScaleDown = &autoscalingv2.HPAScalingRules{}
		}
		newPolicy := getSelectPolicy(*scaleDownEnabled, componentType, false)
		if spec.Behavior.ScaleDown.SelectPolicy == nil || *spec.Behavior.ScaleDown.SelectPolicy != newPolicy {
			spec.Behavior.ScaleDown.SelectPolicy = &newPolicy
			changed = true
		}
	}

	return changed
}

func ensureCPUScaleInGateForHPA(hpa *autoscalingv2.HorizontalPodAutoscaler, cfg metricsprovider.CPUScaleInGateConfig) bool {
	if hpa == nil {
		return false
	}

	changed := false
	if hpa.Annotations == nil {
		hpa.Annotations = make(map[string]string)
	}
	threshold := hpa.Annotations[cfg.ThresholdAnnotation]
	if threshold == "" {
		threshold = strconv.FormatInt(cfg.DefaultThreshold, 10)
		hpa.Annotations[cfg.ThresholdAnnotation] = threshold
		changed = true
	}

	for i := range hpa.Spec.Metrics {
		metric := &hpa.Spec.Metrics[i]
		if metric.Type != autoscalingv2.ObjectMetricSourceType || metric.Object == nil {
			continue
		}
		if metric.Object.Metric.Name != cfg.MetricName {
			continue
		}

		if metric.Object.DescribedObject.APIVersion != "autoscaling/v2" {
			metric.Object.DescribedObject.APIVersion = "autoscaling/v2"
			changed = true
		}
		if metric.Object.DescribedObject.Kind != "HorizontalPodAutoscaler" {
			metric.Object.DescribedObject.Kind = "HorizontalPodAutoscaler"
			changed = true
		}
		if metric.Object.DescribedObject.Name != hpa.Name {
			metric.Object.DescribedObject.Name = hpa.Name
			changed = true
		}
		if metric.Object.Target.Type != autoscalingv2.ValueMetricType {
			metric.Object.Target.Type = autoscalingv2.ValueMetricType
			changed = true
		}
		if metric.Object.Target.Value == nil || metric.Object.Target.Value.Value() != 1 {
			metric.Object.Target.Value = resource.NewQuantity(1, resource.DecimalSI)
			changed = true
		}
		return changed
	}

	hpa.Spec.Metrics = append([]autoscalingv2.MetricSpec{{
		Type: autoscalingv2.ObjectMetricSourceType,
		Object: &autoscalingv2.ObjectMetricSource{
			Metric: autoscalingv2.MetricIdentifier{
				Name: cfg.MetricName,
			},
			DescribedObject: autoscalingv2.CrossVersionObjectReference{
				APIVersion: "autoscaling/v2",
				Kind:       "HorizontalPodAutoscaler",
				Name:       hpa.Name,
			},
			Target: autoscalingv2.MetricTarget{
				Type:  autoscalingv2.ValueMetricType,
				Value: resource.NewQuantity(1, resource.DecimalSI),
			},
		},
	}}, hpa.Spec.Metrics...)
	return true
}

func (s *serverImpl) SetHPAScaleBehavior(
	ctx context.Context,
	sharedPoolId string,
	groupId string,
	componentType entity.ComponentType,
	scaleUpEnabled *bool,
	scaleDownEnabled *bool,
) error {
	// groupId is overloaded by componentType — callers MUST pass:
	//   - TIDB_WORKER / TIKV_WORKER / COPROCESSOR_WORKER → "" (pool-level HPA, no group label)
	//   - TIDB                                            → tidbGroupID   (HPA label `db-<tidbGroupID>`)
	//   - TIPROXY                                         → tidbClusterID (HPA label `db-<tidbClusterID>`,
	//                                                                       created at generatedTiProxyGroup
	//                                                                       with comp.Id = tidbCluster.GetId())
	// The mismatch between TIDB and TIPROXY is intentional: per-group
	// TIPROXY HPA is keyed by tidbClusterID (one TIPROXY per cluster),
	// per-group TIDB HPA is keyed by tidbGroupID (one TIDB per tidbGroup).
	// See writeTidbGroupHPABehavior for the dispatch logic.
	if sharedPoolId == "" {
		return status.Errorf(codes.InvalidArgument, "sharedPoolId is required")
	}
	if componentUsesGroupScopedHPA(componentType) && groupId == "" {
		return status.Errorf(codes.InvalidArgument, "groupId is required for %s component type", componentType)
	}
	if scaleUpEnabled == nil && scaleDownEnabled == nil {
		return nil
	}
	if componentType != apiv1.ComponentType_TIKV && !componentSupportsHPA(componentType) {
		return status.Errorf(codes.Unimplemented, "%s does not support hpa", componentType)
	}
	unlockFunc, lockErr := s.lockTiKVRemoteWorkerGuardedMutation(ctx, sharedPoolId, componentType)
	if lockErr != nil {
		return lockErr
	}
	defer unlockFunc()

	var sharedPool *entity.SharedPool
	var err error
	if componentType == apiv1.ComponentType_COPROCESSOR_WORKER {
		sharedPool, err = s.sharedPoolRepo.Load(ctx, sharedPoolId, false)
		if err != nil {
			return err
		}
		if err := ensureCoprocessorWorkerComponentExists(sharedPool); err != nil {
			return err
		}
	}

	if sharedPool == nil {
		sharedPool, err = s.sharedPoolRepo.Load(ctx, sharedPoolId, false)
		if err != nil {
			return err
		}
	}
	if err := guardSharedPoolMutation(sharedPool, sharedPoolMutationBlockedWhilePaused); 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: entity.GetComponentGroupName(&entity.Component{
					Id:            groupId,
					ComponentType: componentType,
				}),
				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)
		}
	case apiv1.ComponentType_TIPROXY:
		if err = remoteCli.List(ctx, hpas, &client.ListOptions{
			Namespace: namespace,
			LabelSelector: labels.SelectorFromSet(map[string]string{
				schemav1.LabelPingcapGroupKey: entity.GetComponentGroupName(&entity.Component{
					Id:            groupId,
					ComponentType: componentType,
				}),
				schemav1.LabelComponentKey: "tiproxy",
			}),
		}); err != nil {
			return status.Errorf(codes.Internal, "failed to list tiproxy 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 scaleDownEnabled != nil && *scaleDownEnabled {
		switch componentType {
		case apiv1.ComponentType_TIDB:
			return updateHPAWithRetry(ctx, remoteCli, hpa, componentType, scaleUpEnabled, scaleDownEnabled, metricsprovider.TiDBCPUScaleInGateConfig)
		case apiv1.ComponentType_TIPROXY:
			return updateHPAWithRetry(ctx, remoteCli, hpa, componentType, scaleUpEnabled, scaleDownEnabled, metricsprovider.TiProxyCPUScaleInGateConfig)
		}
	}
	ops := opsForScaleBehavior(&hpa.Spec, componentType, scaleUpEnabled, scaleDownEnabled)
	if len(ops) == 0 {
		return nil
	}
	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 opsForScaleBehavior(
	spec *autoscalingv2.HorizontalPodAutoscalerSpec,
	componentType entity.ComponentType,
	scaleUpEnabled *bool,
	scaleDownEnabled *bool,
) []*jsonpatch.Operation {
	behaviorWasNil := spec.Behavior == nil
	if behaviorWasNil {
		spec.Behavior = &autoscalingv2.HorizontalPodAutoscalerBehavior{}
	}

	scaleUpWasNil := scaleUpEnabled != nil && spec.Behavior.ScaleUp == nil
	scaleDownWasNil := scaleDownEnabled != nil && spec.Behavior.ScaleDown == nil
	// SelectPolicy may be nil even when the ScaleUp/ScaleDown object exists
	// (legacy/manual HPA). omitempty drops the key from JSON, so "replace"
	// would fail — must use "add" for the selectPolicy path in that case.
	scaleUpSelectPolicyWasNil := scaleUpEnabled != nil && !scaleUpWasNil && spec.Behavior.ScaleUp.SelectPolicy == nil
	scaleDownSelectPolicyWasNil := scaleDownEnabled != nil && !scaleDownWasNil && spec.Behavior.ScaleDown.SelectPolicy == nil

	// Apply all policy changes in-memory first.
	if scaleUpEnabled != nil {
		newPolicy := getSelectPolicy(*scaleUpEnabled, componentType, true)
		if scaleUpWasNil {
			spec.Behavior.ScaleUp = &autoscalingv2.HPAScalingRules{SelectPolicy: &newPolicy}
		} else {
			spec.Behavior.ScaleUp.SelectPolicy = &newPolicy
		}
	}
	if scaleDownEnabled != nil {
		newPolicy := getSelectPolicy(*scaleDownEnabled, componentType, false)
		if scaleDownWasNil {
			spec.Behavior.ScaleDown = &autoscalingv2.HPAScalingRules{SelectPolicy: &newPolicy}
		} else {
			spec.Behavior.ScaleDown.SelectPolicy = &newPolicy
		}
	}

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

	if behaviorWasNil {
		// Single "add" with the complete Behavior (scaleUp/scaleDown already
		// populated above) — avoids pointer aliasing where later sub-field
		// writes would leak into this op's serialised value.
		ops = append(ops, &jsonpatch.Operation{
			Operation: "add",
			Path:      "/spec/behavior",
			Value:     spec.Behavior,
		})
		return ops
	}

	if scaleUpEnabled != nil {
		switch {
		case scaleUpWasNil:
			ops = append(ops, &jsonpatch.Operation{
				Operation: "add",
				Path:      "/spec/behavior/scaleUp",
				Value:     spec.Behavior.ScaleUp,
			})
		case scaleUpSelectPolicyWasNil:
			ops = append(ops, &jsonpatch.Operation{
				Operation: "add",
				Path:      "/spec/behavior/scaleUp/selectPolicy",
				Value:     *spec.Behavior.ScaleUp.SelectPolicy,
			})
		default:
			ops = append(ops, &jsonpatch.Operation{
				Operation: "replace",
				Path:      "/spec/behavior/scaleUp/selectPolicy",
				Value:     *spec.Behavior.ScaleUp.SelectPolicy,
			})
		}
	}

	if scaleDownEnabled != nil {
		switch {
		case scaleDownWasNil:
			ops = append(ops, &jsonpatch.Operation{
				Operation: "add",
				Path:      "/spec/behavior/scaleDown",
				Value:     spec.Behavior.ScaleDown,
			})
		case scaleDownSelectPolicyWasNil:
			ops = append(ops, &jsonpatch.Operation{
				Operation: "add",
				Path:      "/spec/behavior/scaleDown/selectPolicy",
				Value:     *spec.Behavior.ScaleDown.SelectPolicy,
			})
		default:
			ops = append(ops, &jsonpatch.Operation{
				Operation: "replace",
				Path:      "/spec/behavior/scaleDown/selectPolicy",
				Value:     *spec.Behavior.ScaleDown.SelectPolicy,
			})
		}
	}

	return ops
}

func updateHPAWithRetry(
	ctx context.Context,
	remoteCli client.Client,
	hpa *autoscalingv2.HorizontalPodAutoscaler,
	componentType entity.ComponentType,
	scaleUpEnabled *bool,
	scaleDownEnabled *bool,
	gateConfig metricsprovider.CPUScaleInGateConfig,
) error {
	for i := 0; i < 3; i++ {
		gateChanged := ensureCPUScaleInGateForHPA(hpa, gateConfig)
		behaviorChanged := applyScaleBehaviorInPlace(&hpa.Spec, componentType, scaleUpEnabled, scaleDownEnabled)
		if !gateChanged && !behaviorChanged {
			return nil
		}
		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")
}