package noderebalance

import (
	"context"
	"fmt"
	"strconv"
	"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"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/labels"
	corev1 "k8s.io/api/core/v1"
	policyv1 "k8s.io/api/policy/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"
)

func getRebalanceSignalNodes(allNodes []*corev1.Node) []*corev1.Node {
	rebalanceSignalNodes := make([]*corev1.Node, 0)
	for _, node := range allNodes {
		for i := range node.Spec.Taints {
			if node.Spec.Taints[i].Key == types.RebalanceSignalTaintKey {
				rebalanceSignalNodes = append(rebalanceSignalNodes, node)
				break
			}
		}
	}
	return rebalanceSignalNodes
}

func needEvictCurrentNode(node *corev1.Node) bool {
	for i := range node.Spec.Taints {
		if node.Spec.Taints[i].Key == types.RebalanceSignalTaintKey {
			return true
		}
	}
	return false
}

func rebalanceNoticeTimestamp(node *corev1.Node) (int64, error) {
	for i := range node.Spec.Taints {
		if node.Spec.Taints[i].Key == types.RebalanceSignalTaintKey {
			timestamp, err := strconv.ParseInt(node.Spec.Taints[i].Value, 10, 64)
			return timestamp, err
		}
	}
	return 0, nil
}

func checkAllEvicted(evictedList []bool) bool {
	if len(evictedList) == 0 {
		return false
	}
	for _, evicted := range evictedList {
		if !evicted {
			return false
		}
	}
	return true
}

func (r *Reconciler) evictNode(ctx context.Context, node *corev1.Node) (bool, bool, error) {
	var retry bool
	podList := &corev1.PodList{}
	if err := r.Client.List(ctx, podList, client.MatchingFields{"spec.nodeName": node.Name}); err != nil {
		return false, retry, err
	}

	evictedList := make([]bool, 0, len(podList.Items))
	for i := range podList.Items {
		pod := &podList.Items[i]
		if pod.DeletionTimestamp != nil {
			continue
		}
		if ownerRef := metav1.GetControllerOf(pod); ownerRef != nil && ownerRef.Kind == "DaemonSet" {
			continue
		}

		if err := client.IgnoreNotFound(r.Client.SubResource("eviction").Create(
			ctx, pod, &policyv1.Eviction{}),
		); err != nil {
			if !apierrors.IsTooManyRequests(err) {
				return false, retry, err
			}
			retry = true
			evictedList = append(evictedList, false)
			log.FromContext(ctx).Error(err, "evict pod error", "pod", pod.Name+"/"+pod.Namespace, "node", node.Name)
		} else {
			evictedList = append(evictedList, true)
			log.FromContext(ctx).Info("evict pod success", "pod", pod.Name+"/"+pod.Namespace, "node", node.Name)
		}
	}

	return checkAllEvicted(evictedList), retry, nil
}

func (r *Reconciler) preventCAScaleUpAndEvict(ctx context.Context, allNodes []*corev1.Node, node *corev1.Node) error {
	for _, node := range allNodes {
		needUpdate := true
		for i := range node.Spec.Taints {
			if node.Spec.Taints[i].Key == types.PreventCAScaleUpTaintKey {
				needUpdate = false
				break
			}
		}
		if needUpdate {
			node.Spec.Taints = append(node.Spec.Taints, corev1.Taint{
				Key:    types.PreventCAScaleUpTaintKey,
				Value:  fmt.Sprint(time.Now().Unix()),
				Effect: corev1.TaintEffectNoSchedule,
			})
			if err := r.Client.Update(ctx, node); err != nil {
				return err
			}
			log.FromContext(ctx).Info("set PreventCAScaleUpTaintKey success", "node", node.Name)
		}
	}
	if len(allNodes) > 0 {
		nodeGroup := allNodes[0].Labels[labels.ServerlessNodeGroupLabelKey]
		nodeZone := allNodes[0].Labels[corev1.LabelTopologyZone]
		metrics.PreventCAScaleUpGauge.WithLabelValues(nodeGroup, nodeZone).Set(1)
	}

	if !needEvictCurrentNode(node) {
		return nil
	}

	allEvicted, retry, err := r.evictNode(ctx, node)
	if err != nil {
		return err
	}
	if retry {
		return &utils.RequeueAfterError{RequeueAfter: retryTooManyRequestsTime}
	}
	if allEvicted {
		nodeGroup := node.Labels[labels.ServerlessNodeGroupLabelKey]
		nodeZone := node.Labels[corev1.LabelTopologyZone]
		r.Recorder.Event(node, corev1.EventTypeNormal, "EvictRebalanceNode", "evict rebalance signal node success")
		metrics.RebalanceEvictNodeCounter.WithLabelValues(node.Name, nodeGroup, nodeZone).Inc()
	}

	return nil
}

func (r *Reconciler) permitCAScaleUp(ctx context.Context, allNodes []*corev1.Node) error {
	for _, node := range allNodes {
		for i := range node.Spec.Taints {
			if node.Spec.Taints[i].Key == types.PreventCAScaleUpTaintKey {
				node.Spec.Taints = append(node.Spec.Taints[:i], node.Spec.Taints[i+1:]...)
				if err := r.Client.Update(ctx, node); err != nil {
					return err
				}
				log.FromContext(ctx).Info("del PreventCAScaleUpTaintKey success", "node", node.Name)
				break
			}
		}
	}
	if len(allNodes) > 0 {
		nodeGroup := allNodes[0].Labels[labels.ServerlessNodeGroupLabelKey]
		nodeZone := allNodes[0].Labels[corev1.LabelTopologyZone]
		metrics.PreventCAScaleUpGauge.WithLabelValues(nodeGroup, nodeZone).Set(0)
	}
	return nil
}

func (r *Reconciler) preventCAScaleUpIfNeed(ctx context.Context, allNodes []*corev1.Node) error {
	unPreventNodes := make([]*corev1.Node, 0)
	for _, node := range allNodes {
		isUnPreventNode := true
		for i := range node.Spec.Taints {
			if node.Spec.Taints[i].Key == types.PreventCAScaleUpTaintKey {
				isUnPreventNode = false
				break
			}
		}
		if isUnPreventNode {
			unPreventNodes = append(unPreventNodes, node)
		}
	}

	nodeGroup := allNodes[0].Labels[labels.ServerlessNodeGroupLabelKey]
	nodeZone := allNodes[0].Labels[corev1.LabelTopologyZone]
	if len(unPreventNodes) < len(allNodes) {
		for _, node := range unPreventNodes {
			node.Spec.Taints = append(node.Spec.Taints, corev1.Taint{
				Key:    types.PreventCAScaleUpTaintKey,
				Value:  fmt.Sprint(time.Now().Unix()),
				Effect: corev1.TaintEffectNoSchedule,
			})
			if err := r.Client.Update(ctx, node); err != nil {
				return err
			}
			log.FromContext(ctx).Info("complete PreventCAScaleUpTaintKey success", "node", node.Name)
		}
		metrics.PreventCAScaleUpGauge.WithLabelValues(nodeGroup, nodeZone).Set(1)
	} else {
		metrics.PreventCAScaleUpGauge.WithLabelValues(nodeGroup, nodeZone).Set(0)
	}

	return nil
}