package tidbnodescaler

import (
	"context"
	"time"

	"github.com/tidbcloud/aws-shared-provider/pkg/metrics"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/kubeapi"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/labels"
	corev1 "k8s.io/api/core/v1"
	"sigs.k8s.io/controller-runtime/pkg/log"
)

func (r *Reconciler) trySetScaleDownDisabledNode(ctx context.Context, node *corev1.Node, score int) {
	logger := log.FromContext(ctx, "nodeName", node.Name, "score", score).WithName("trySetScaleDownDisabledNode")
	if node.Annotations[types.AnnoNodeScaleDownDisabledTimestamp] != "" {
		logger.V(2).Info("node already has scale down disabled annotation")
		return
	}

	if err := kubeapi.UpdateNode(ctx, r.Client, func() *corev1.Node {
		if node.Annotations == nil {
			node.Annotations = map[string]string{}
		}
		node.Annotations[types.AnnoNodeScaleDownDisabledTimestamp] = time.Now().Format(time.RFC3339)
		if r.ProviderType == types.ProviderTypeAliCloudFreeTier {
			node.Annotations[types.GoatScaleDownDisabledKey] = "true"
		} else {
			node.Annotations[types.CAScaleDownDisabledKey] = "true"
		}
		return node
	}); err != nil {
		logger.Error(err, "failed to update node")
		return
	}
	r.Recorder.Eventf(node, corev1.EventTypeNormal, "SetScaleDownDisabled", "Set scale down disabled annotation, score=%d", score)
	metrics.SetScaleDownDisabledCounter.WithLabelValues(node.Labels[labels.ServerlessNodeGroupLabelKey]).Inc()
}

func (r *Reconciler) tryRemoveScaleDownDisabledNode(ctx context.Context, node *corev1.Node, score int) {
	logger := log.FromContext(ctx, "nodeName", node.Name, "score", score).WithName("tryRemoveScaleDownDisabledNode")
	if node.Annotations[types.AnnoNodeScaleDownDisabledTimestamp] == "" {
		logger.V(2).Info("node does not have scale down disabled annotation")
		return
	}

	if err := kubeapi.UpdateNode(ctx, r.Client, func() *corev1.Node {
		delete(node.Annotations, types.AnnoNodeScaleDownDisabledTimestamp)
		if r.ProviderType == types.ProviderTypeAliCloudFreeTier {
			delete(node.Annotations, types.GoatScaleDownDisabledKey)
		} else {
			delete(node.Annotations, types.CAScaleDownDisabledKey)
		}
		return node
	}); err != nil {
		logger.Error(err, "failed to update node")
		return
	}
	r.Recorder.Eventf(node, corev1.EventTypeNormal, "RemoveScaleDownDisabled", "Removed scale down disabled annotation, score=%d", score)
	metrics.RemoveScaleDownDisabledCounter.WithLabelValues(node.Labels[labels.ServerlessNodeGroupLabelKey]).Inc()
}