package nodedrainer

import (
	"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"
	corev1 "k8s.io/api/core/v1"
	"sigs.k8s.io/controller-runtime/pkg/log"
)

func (dn *NodeDrainer) recoverNode(node *corev1.Node) error {
	if err := kubeapi.UpdateNode(dn.ctx, dn.client, func() *corev1.Node {
		delete(node.Annotations, types.AnnoNodeDrainTimestamp)
		node.Spec.Unschedulable = false
		return node
	}); err != nil {
		metrics.UpdateRecoverNodeCounter.WithLabelValues("fail").Inc()
		return err
	}
	metrics.UpdateRecoverNodeCounter.WithLabelValues("success").Inc()
	dn.recorder.Eventf(node, corev1.EventTypeNormal, "RecoverDrainedNode", "Drained node %s has been recovered", node.Name)

	return nil
}

func (dn *NodeDrainer) recover(nodeList []*corev1.Node) {
	for _, node := range nodeList {
		if node.Annotations[types.AnnoNodeDrainTimestamp] == "" {
			continue
		}
		if err := dn.recoverNode(node); err != nil {
			log.FromContext(dn.ctx).Error(err, "failed to recover on-demand node", "nodeName", node.Name)
		}
	}
}