package noderebalance

import (
	"context"
	"time"

	"github.com/tidbcloud/aws-shared-provider/pkg/manager"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/labels"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/client-go/tools/record"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"
	"sigs.k8s.io/controller-runtime/pkg/predicate"
)

const (
	maxRebalanceSignalTime   = 300 // better to be the same as the node's heartbeat.
	retryTooManyRequestsTime = 3 * time.Second
)

// Reconciler reconciles a nodedrainer spot Node object.
type Reconciler struct {
	Client   client.Client
	Recorder record.EventRecorder
}

// Reconcile is part of the main kubernetes reconciliation loop which aims to
// move the current state of the cluster closer to the desired state.
func (r *Reconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	logger := log.FromContext(ctx)
	logger.V(2).Info("start reconcile noderebalance")
	start := time.Now()
	defer func() {
		logger.V(2).Info("finish reconcile noderebalance", "time", time.Since(start))
	}()

	currentNode := &corev1.Node{}
	if err := r.Client.Get(ctx, req.NamespacedName, currentNode); err != nil {
		return utils.ReconcileErr(logger, "NodeRebalance", req.Name, client.IgnoreNotFound(err))
	}

	nodeGroup := currentNode.Labels[labels.ServerlessNodeGroupLabelKey]
	nodeZone := currentNode.Labels[corev1.LabelTopologyZone]
	logger = logger.WithValues("nodeGroup", nodeGroup, "nodeZone", nodeZone)
	ctx = log.IntoContext(ctx, logger)

	allNodes, err := r.getAllNodes(ctx, currentNode)
	if err != nil {
		return utils.ReconcileErr(logger, "NodeRebalance", req.Name, err)
	}

	rebalanceSignalNodes := getRebalanceSignalNodes(allNodes)
	switch {
	case len(rebalanceSignalNodes) == 0:
		err = r.permitCAScaleUp(ctx, allNodes)
	case len(rebalanceSignalNodes) > 1:
		err = r.preventCAScaleUpAndEvict(ctx, allNodes, currentNode)
	case len(rebalanceSignalNodes) == 1:
		var timestamp int64
		if timestamp, err = rebalanceNoticeTimestamp(rebalanceSignalNodes[0]); err == nil {
			if time.Now().Unix()-timestamp > maxRebalanceSignalTime {
				err = r.preventCAScaleUpAndEvict(ctx, allNodes, currentNode)
				break
			}
			err = r.preventCAScaleUpIfNeed(ctx, allNodes)
		}
	}

	return utils.ReconcileErr(logger, "NodeRebalance", req.Name, err)
}

// SetupWithManager sets up the controller with the Manager.
func (r *Reconciler) SetupWithManager(mgr manager.Manager) error {
	r.Recorder = mgr.GetEventRecorderFor("NodeRebalance")

	byAnnoFuncs := predicate.NewPredicateFuncs(func(object client.Object) bool {
		return object.GetLabels()[labels.NodeDrainerCapacityTypeLabelKey] == labels.NodeDrainerCapacityTypeSpot
	})
	return ctrl.NewControllerManagedBy(mgr).Named("NodeRebalance").For(&corev1.Node{}).WithEventFilter(byAnnoFuncs).
		Complete(r)
}

func (r *Reconciler) getAllNodes(ctx context.Context, currentNode *corev1.Node) ([]*corev1.Node, error) {
	nodeList := &corev1.NodeList{}
	if err := r.Client.List(ctx, nodeList, client.MatchingLabels{
		labels.ServerlessNodeGroupLabelKey: currentNode.Labels[labels.ServerlessNodeGroupLabelKey],
		corev1.LabelTopologyZone:           currentNode.Labels[corev1.LabelTopologyZone],
	}); err != nil {
		return nil, err
	}

	allNodes := make([]*corev1.Node, 0, len(nodeList.Items))
	for i := range nodeList.Items {
		if nodeList.Items[i].DeletionTimestamp == nil {
			allNodes = append(allNodes, &nodeList.Items[i])
		}
	}
	return allNodes, nil
}