package tidbpodrescheduler

import (
	"context"
	"math"
	"slices"
	"sort"
	"time"

	serverlessv1alpha1 "github.com/tidbcloud/aws-shared-provider/apis/serverless/v1alpha1"
	"github.com/tidbcloud/aws-shared-provider/pkg/metrics"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbmanager"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/kubeapi"
	"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"
)

const (
	rescheduleFailTimeout = 20 * time.Minute
)

type podInfo struct {
	podName  string
	nodeName string
}

func (r *Reconciler) reschedule(ctx context.Context, tManager tidbmanager.TidbManager, nodeScore serverlessv1alpha1.NodeScoreStatus) {
	tidbPool := tManager.GetTidbPool()
	pool := tidbPool.GetPool(ctx)
	logger := log.FromContext(ctx, "tidbPool", tidbPool.Namespace()).WithName("reschedule")

	patchedNodes, err := r.deleteRescheduledPendingPods(ctx, tidbPool.Namespace())
	if err != nil {
		logger.Error(err, "failed to delete rescheduled pending pods")
		return
	}

	tidbPoolNodeList := &corev1.NodeList{}
	labelSelector := &metav1.LabelSelector{MatchLabels: map[string]string{
		labels.ServerlessComponentNodeLabelKey: "tidb",
	}}
	if len(pool.NodeGroups) > 0 {
		labelSelector.MatchExpressions = []metav1.LabelSelectorRequirement{{
			Key:      labels.ServerlessNodeGroupLabelKey,
			Operator: metav1.LabelSelectorOpIn,
			Values:   pool.NodeGroups,
		}}
	}
	sel, err := metav1.LabelSelectorAsSelector(labelSelector)
	if err != nil {
		logger.Error(err, "invalid label selector for tidb nodes")
		return
	}
	if err := r.Client.List(ctx, tidbPoolNodeList, client.MatchingLabelsSelector{Selector: sel}); err != nil {
		logger.Error(err, "failed to list tidb pool nodes")
		return
	}
	if len(tidbPoolNodeList.Items) == 0 {
		logger.Info("tidb pool nodes are empty, skip")
		return
	}
	sort.Slice(tidbPoolNodeList.Items, func(i, j int) bool {
		return nodeScore.SchedulerNodeScore[tidbPoolNodeList.Items[i].Name] < nodeScore.SchedulerNodeScore[tidbPoolNodeList.Items[j].Name]
	})

	var rescheduleToNodes []*corev1.Node     // find the nodes within 20 points of the highest score
	var rescheduleFromNodes []*corev1.Node   // find the bottom 20% nodes with the lowest score and nodes that need eviction
	var rescheduleMiddleNodes []*corev1.Node // find the middle nodes that are left
	highestScore := nodeScore.SchedulerNodeScore[tidbPoolNodeList.Items[len(tidbPoolNodeList.Items)-1].Name]
	for i := range tidbPoolNodeList.Items {
		tidbNode := &tidbPoolNodeList.Items[i]
		if patchedNode := patchedNodes[tidbNode.Name]; patchedNode != nil && patchedNode.ResourceVersion != tidbNode.ResourceVersion {
			patchedNode.DeepCopyInto(tidbNode)
		}

		if tidbNode.DeletionTimestamp != nil {
			continue
		}
		if nodeNeedEvict := nodeScore.NodeNeedEvict; nodeNeedEvict != nil &&
			(nodeNeedEvict.CPU[tidbNode.Name] > 0 || nodeNeedEvict.MEM[tidbNode.Name] > 0) {
			rescheduleFromNodes = append(rescheduleFromNodes, tidbNode)
			continue
		}
		if nodeScore.SchedulerNodeScore[tidbNode.Name] > highestScore-20 &&
			kubeapi.IsNodeReady(tidbNode) && !tidbNode.Spec.Unschedulable {
			rescheduleToNodes = append(rescheduleToNodes, tidbNode)
			continue
		}
		if i < divCeil(len(tidbPoolNodeList.Items), 5) {
			rescheduleFromNodes = append(rescheduleFromNodes, tidbNode)
			continue
		}
		rescheduleMiddleNodes = append(rescheduleMiddleNodes, tidbNode)
	}
	slices.Reverse(rescheduleToNodes) // reverse the order to make sure the highest score node is first

	if len(rescheduleToNodes) == 0 {
		metrics.RescheduleSkipCounter.WithLabelValues(tidbPool.Namespace(), "NoToNodes").Inc()
		logger.V(1).Info("no rescheduleToNodes, skip")
		return
	}
	if len(rescheduleMiddleNodes) == 0 {
		metrics.RescheduleSkipCounter.WithLabelValues(tidbPool.Namespace(), "NoMiddleNodes").Inc()
		logger.V(1).Info("no rescheduleMiddleNodes means reschedule too close, skip")
		return
	}
	if len(rescheduleFromNodes) == 0 {
		metrics.RescheduleSkipCounter.WithLabelValues(tidbPool.Namespace(), "NoFromNodes").Inc()
		logger.V(1).Info("no rescheduleFromNodes, skip")
		return
	}

	freeCachePerNode := divCeil(tManager.FreeCacheSize(), len(tidbPoolNodeList.Items))
	freePods := getFreePodsFromNodes(tManager, rescheduleFromNodes, mulCeil(freeCachePerNode, 0.8))
	if len(freePods) == 0 {
		metrics.RescheduleSkipCounter.WithLabelValues(tidbPool.Namespace(), "NoFreePods").Inc()
		logger.V(2).Info("no free pods to reschedule, skip", "freeCachePerNode", freeCachePerNode)
		return
	}
	if len(freePods) > types.MaxAbnormalTidbPods {
		freePods = freePods[:types.MaxAbnormalTidbPods]
	}
	freePodsIndex := 0
	for _, toNode := range rescheduleToNodes {
		logger := logger.WithValues("freeCachePerNode", freeCachePerNode,
			"toNodeName", toNode.Name, "toNodeScore", nodeScore.SchedulerNodeScore[toNode.Name])
		if !canNodeReschedule(toNode) {
			metrics.RescheduleSkipCounter.WithLabelValues(tidbPool.Namespace(), "NodeNotReschedulable").Inc()
			logger.Info("node is not reschedulable, skip")
			continue
		}
		abnormalList, err := tidbPool.GetAbnormalPodList(ctx)
		if err != nil {
			logger.Error(err, "get abnormal tidb pod list")
			return
		}
		freeCacheToNode := len(tManager.GetTidbFromFreeCache(toNode.Name, true))
		incomingPodToNode := incomingPodCount(ctx, toNode.Name, abnormalList)
		needPods := freeCachePerNode - freeCacheToNode - incomingPodToNode
		if needPods <= 0 {
			metrics.RescheduleSkipCounter.WithLabelValues(tidbPool.Namespace(), "EnoughFreePods").Inc()
			logger.V(2).Info("node has enough free pods, skip", "freeCache", freeCacheToNode, "incomingPod", incomingPodToNode)
			continue
		}

		for i := 0; i < needPods; i++ {
			if freePodsIndex >= len(freePods) {
				metrics.RescheduleSkipCounter.WithLabelValues(tidbPool.Namespace(), "NoMoreFreePods").Inc()
				logger.Info("no more free pods to reschedule", "freeCache", freeCacheToNode, "incomingPod", incomingPodToNode)
				return
			}

			freePod := freePods[freePodsIndex]
			logger := logger.WithValues("fromNodeName", freePod.nodeName, "fromNodeScore", nodeScore.SchedulerNodeScore[freePod.nodeName])
			if err := r.Client.SubResource("eviction").Create(ctx, &corev1.Pod{ObjectMeta: metav1.ObjectMeta{
				Name:      freePod.podName,
				Namespace: tidbPool.Namespace(),
			}}, &policyv1.Eviction{}); err != nil {
				metrics.RescheduleCounter.WithLabelValues(tidbPool.Namespace(), "EvictFreePodFail").Inc()
				logger.Error(err, "failed to evict free cache pod", "podName", freePod.podName)
				if apierrors.IsTooManyRequests(err) {
					return
				}
				freePodsIndex++
				continue
			}
			metrics.RescheduleCounter.WithLabelValues(tidbPool.Namespace(), "EvictFreePodSuccess").Inc()
			logger.Info("[pod lifecycle] evict free cache pod", "podName", freePod.podName)

			pod := pool.TidbPod.DeepCopy()
			pod.Annotations = utils.UpdateMap(pod.Annotations, types.AnnoTidbReschedule, "true")
			ensureNodeNameAffinity(pod, toNode.Name) // force schedule to the target node
			if err := r.Client.Create(ctx, pod); err != nil {
				metrics.RescheduleCounter.WithLabelValues(tidbPool.Namespace(), "CreateNewPodFail").Inc()
				logger.Error(err, "failed to create new pod", "podName", pod.Name)
			} else {
				metrics.RescheduleCounter.WithLabelValues(tidbPool.Namespace(), "CreateNewPodSuccess").Inc()
				logger.Info("create new pod", "podName", pod.Name)
			}
			freePodsIndex++
		}
	}
}

func canNodeReschedule(node *corev1.Node) bool {
	rescheduleFailTimestamp := node.Annotations[types.AnnoNodeRescheduleFailTimestamp]
	if rescheduleFailTimestamp == "" {
		return true
	}
	rescheduleFailTime, err := time.Parse(time.RFC3339, rescheduleFailTimestamp)
	if err != nil {
		return false
	}
	return time.Since(rescheduleFailTime) > rescheduleFailTimeout
}

// getFreePodsFromNodes retrieves free pods from the specified nodes.
func getFreePodsFromNodes(tManager tidbmanager.TidbManager, nodes []*corev1.Node, reservePerNode int) []podInfo {
	freePods := make([]podInfo, 0)
	for _, node := range nodes {
		pods := tManager.GetTidbFromFreeCache(node.Name, false)
		if len(pods) > reservePerNode {
			for _, pod := range pods[reservePerNode:] {
				freePods = append(freePods, podInfo{pod, node.Name})
			}
		}
	}
	return freePods
}

func divCeil(a, b int) int {
	return (a + b - 1) / b
}

func mulCeil(a int, b float64) int {
	return int(math.Ceil(float64(a) * b))
}