package nodescorer

import (
	"context"
	"math"
	"os"
	"time"

	"github.com/go-logr/logr"
	"github.com/pingcap/errors"
	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/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"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/query"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

const (
	MaxNodeScore                   = 100
	DrainerOnDemandNodeScoreDeduct = 40
	RebalanceNodeScoreDeduct       = 30
	NotReadyDeduct                 = 100
	DeletionCandidateOfCADeduct    = 160

	DeletionCandidateOfCATaint = "DeletionCandidateOfClusterAutoscaler"
)

type NodeScorer struct {
	logr.Logger

	ctx           context.Context
	client        client.Client
	checkInterval time.Duration
	o11yQuerier   query.PromQuerier
}

func New(client client.Client, checkInterval time.Duration) *NodeScorer {
	return &NodeScorer{
		Logger: ctrl.Log.WithName("nodescorer"),

		client:        client,
		checkInterval: checkInterval,
	}
}

func (ns *NodeScorer) Start(ctx context.Context) error {
	ns.ctx = ctx
	ns.Info("start nodescorer")

	queryURL := os.Getenv(types.O11yGlobalReadURL)
	if queryURL == "" {
		return errors.Errorf("env %s is empty", types.O11yGlobalReadURL)
	}
	ns.o11yQuerier = query.NewPromQueryer(nil, queryURL)

	checkTicker := time.NewTicker(ns.checkInterval)
	defer checkTicker.Stop()

	// call first start before checkTicker
	ns.do()

	for {
		select {
		case <-checkTicker.C:
			ns.do()
		case <-ns.ctx.Done():
			err := ns.ctx.Err()
			ns.Error(err, "stop nodescorer")
			return err
		}
	}
}

//nolint:funlen // ignore function length
func (ns *NodeScorer) do() {
	start := time.Now()
	ns.V(1).Info("start nodescorer check cycle")
	defer func() {
		ns.V(1).Info("finish nodescorer check cycle", "time", time.Since(start))
	}()

	nodeScoreMap := &serverlessv1alpha1.NodeScore{}
	if err := ns.client.Get(ns.ctx, client.ObjectKey{
		Name:      types.NodeScoreMap,
		Namespace: os.Getenv("NAMESPACE"),
	}, nodeScoreMap); err != nil {
		ns.Error(err, "failed to get node score map")
		return
	}

	tidbPoolNodeList := &corev1.NodeList{}
	if err := ns.client.List(ns.ctx, tidbPoolNodeList, client.MatchingLabels{
		labels.ServerlessComponentNodeLabelKey: "tidb",
	}); err != nil {
		ns.Error(err, "failed to list tidb pool nodes")
		return
	}
	tidbPoolNodes := make(map[string]bool, len(tidbPoolNodeList.Items))
	tidbPoolDrainerOnDemandNodes := make([]string, 0)
	tidbPoolRebalanceNodes := make([]string, 0)
	tidbPoolNotReadyNodes := make([]string, 0)
	tidbPoolDeletionCandidateOfCANodes := make([]string, 0)
	for i := 0; i < len(tidbPoolNodeList.Items); i++ {
		node := tidbPoolNodeList.Items[i]
		if node.DeletionTimestamp != nil {
			// don't worry about the terminating node, as the pods above are also terminating and cannot schedule new pods
			ns.Info("skip tidb pool node in terminating", "node", node.Name, "nodegroup", node.Labels[labels.ServerlessNodeGroupLabelKey])
			continue
		}
		tidbPoolNodes[node.Name] = true
		if node.Labels[labels.NodeDrainerCapacityTypeLabelKey] == labels.NodeDrainerCapacityTypeOnDemand {
			tidbPoolDrainerOnDemandNodes = append(tidbPoolDrainerOnDemandNodes, node.Name)
		}
		for j := range node.Spec.Taints {
			if node.Spec.Taints[j].Key == types.RebalanceSignalTaintKey {
				tidbPoolRebalanceNodes = append(tidbPoolRebalanceNodes, node.Name)
				break
			}
		}
		for j := range node.Spec.Taints {
			if node.Spec.Taints[j].Key == DeletionCandidateOfCATaint {
				tidbPoolDeletionCandidateOfCANodes = append(tidbPoolDeletionCandidateOfCANodes, node.Name)
				break
			}
		}
		if !kubeapi.IsNodeReady(&node) {
			tidbPoolNotReadyNodes = append(tidbPoolNotReadyNodes, node.Name)
		}
	}
	if len(tidbPoolNodes) == 0 {
		ns.Info("tidb pool nodes are empty, skip")
		return
	}

	if nodeScoreMap.Spec.CPU == nil && nodeScoreMap.Spec.MEM == nil {
		ns.Info("disable nodescore, skip", "name", nodeScoreMap.Namespace+"/"+nodeScoreMap.Name)
		// we must clean NodeScoreStatus when disable nodescore
		if len(nodeScoreMap.Status.ResourceNodeScore) > 0 || len(nodeScoreMap.Status.SchedulerNodeScore) > 0 ||
			nodeScoreMap.Status.NodeNeedEvict != nil {
			nodeScoreMap.Status = serverlessv1alpha1.NodeScoreStatus{}
			if err := ns.nodeScoreStatusUpdate(nodeScoreMap); err != nil {
				ns.Error(err, "Update NodeScoreMap status")
			}
			metrics.NeedEvictNodesGauge.WithLabelValues("cpu").Set(0)
			metrics.NeedEvictNodesGauge.WithLabelValues("mem").Set(0)
			metrics.NodesGauge.WithLabelValues().Set(float64(len(tidbPoolNodes)))
		}
		return
	}

	var nodeCpuIdle map[string]float32
	if nodeScoreMap.Spec.CPU != nil {
		var err error
		if nodeCpuIdle, err = ns.getCpuIdle(tidbPoolNodes); err != nil {
			ns.Error(err, "getCpuIdle")
			return
		}
		ns.V(3).Info("getCpuIdle", "rlt", nodeCpuIdle)
	}
	var nodeMemIdle map[string]float32
	if nodeScoreMap.Spec.MEM != nil {
		var err error
		if nodeMemIdle, err = ns.getMemIdle(tidbPoolNodes); err != nil {
			ns.Error(err, "getMemIdle")
			return
		}
		ns.V(3).Info("getMemIdle", "rlt", nodeMemIdle)
	}

	nodeScoreMap.Status.ResourceNodeScore = calculateNodeScore(nodeScoreMap, nodeCpuIdle, nodeMemIdle)
	ns.V(3).Info("ResourceNodeScore", "rlt", nodeScoreMap.Status.ResourceNodeScore)
	nodeScoreMap.Status.SchedulerNodeScore = copyNodeScore(nodeScoreMap.Status.ResourceNodeScore)
	for _, drainerOnDemandNode := range tidbPoolDrainerOnDemandNodes {
		if _, ok := nodeScoreMap.Status.SchedulerNodeScore[drainerOnDemandNode]; ok {
			nodeScoreMap.Status.SchedulerNodeScore[drainerOnDemandNode] -= DrainerOnDemandNodeScoreDeduct
		}
	}
	for _, rebalanceNode := range tidbPoolRebalanceNodes {
		if _, ok := nodeScoreMap.Status.SchedulerNodeScore[rebalanceNode]; ok {
			nodeScoreMap.Status.SchedulerNodeScore[rebalanceNode] -= RebalanceNodeScoreDeduct
		}
	}
	for _, notReadyNode := range tidbPoolNotReadyNodes {
		if _, ok := nodeScoreMap.Status.SchedulerNodeScore[notReadyNode]; ok {
			nodeScoreMap.Status.SchedulerNodeScore[notReadyNode] -= NotReadyDeduct
		}
	}
	for _, deletionCandidateOfCANode := range tidbPoolDeletionCandidateOfCANodes {
		if _, ok := nodeScoreMap.Status.SchedulerNodeScore[deletionCandidateOfCANode]; ok {
			nodeScoreMap.Status.SchedulerNodeScore[deletionCandidateOfCANode] -= DeletionCandidateOfCADeduct
		}
	}
	ns.V(3).Info("SchedulerNodeScore", "rlt", nodeScoreMap.Status.SchedulerNodeScore)

	nodeScoreMap.Status.NodeNeedEvict = &serverlessv1alpha1.EvictType{}
	if nodeScoreMap.Spec.CPU != nil {
		nodeScoreMap.Status.NodeNeedEvict.CPU = filterEvictThreshold(nodeCpuIdle, nodeScoreMap.Spec.CPU.EvictThreshold)
	}
	if nodeScoreMap.Spec.MEM != nil {
		nodeScoreMap.Status.NodeNeedEvict.MEM = filterEvictThreshold(nodeMemIdle, nodeScoreMap.Spec.MEM.EvictThreshold)
	}

	if err := ns.nodeScoreStatusUpdate(nodeScoreMap); err != nil {
		ns.Error(err, "Update NodeScoreMap status")
		return
	}
	metrics.NeedEvictNodesGauge.WithLabelValues("cpu").Set(float64(len(nodeScoreMap.Status.NodeNeedEvict.CPU)))
	metrics.NeedEvictNodesGauge.WithLabelValues("mem").Set(float64(len(nodeScoreMap.Status.NodeNeedEvict.MEM)))
	metrics.NodesGauge.WithLabelValues().Set(float64(len(tidbPoolNodes)))
}

func (ns *NodeScorer) nodeScoreStatusUpdate(nodeScoreMap *serverlessv1alpha1.NodeScore) error {
	now := metav1.Now()
	nodeScoreMap.Status.UpdateTime = &now
	return ns.client.Status().Update(ns.ctx, nodeScoreMap)
}

func copyNodeScore(nodeScore map[string]int) map[string]int {
	copiedNodeScore := make(map[string]int, len(nodeScore))
	for k, v := range nodeScore {
		copiedNodeScore[k] = v
	}
	return copiedNodeScore
}

func calculateNodeScore(
	nodeScoreMap *serverlessv1alpha1.NodeScore,
	nodeCpuIdle, nodeMemIdle map[string]float32,
) map[string]int {
	cpuWeight, memWeight := float32(0), float32(0)
	if nodeScoreMap.Spec.CPU != nil {
		cpuWeight = float32(nodeScoreMap.Spec.CPU.Weight)
	}
	if nodeScoreMap.Spec.MEM != nil {
		memWeight = float32(nodeScoreMap.Spec.MEM.Weight)
	}
	if cpuWeight+memWeight == 0 {
		return nil
	}

	// build union of nodes from CPU and MEM maps
	nodeNames := make(map[string]struct{}, utils.Max(len(nodeCpuIdle), len(nodeMemIdle)))
	for name := range nodeCpuIdle {
		nodeNames[name] = struct{}{}
	}
	for name := range nodeMemIdle {
		nodeNames[name] = struct{}{}
	}

	nodeScore := make(map[string]int, len(nodeNames))
	for name := range nodeNames {
		var sumWeight float32
		cpuIdle, ok := nodeCpuIdle[name]
		if ok {
			sumWeight += cpuWeight
		}
		memIdle, ok := nodeMemIdle[name]
		if ok {
			sumWeight += memWeight
		}
		nodeScore[name] = int((cpuIdle*cpuWeight + memIdle*memWeight) * float32(MaxNodeScore) / sumWeight)
	}

	return nodeScore
}

func filterEvictThreshold(m map[string]float32, evictThreshold float32) map[string]float32 {
	filterMap := map[string]float32{}
	for node, idle := range m {
		usage := 1 - idle
		if usage > evictThreshold {
			// round to the nearest number with 2 decimal places
			//nolint: gomnd // nolint
			filterMap[node] = float32(math.Round(float64(usage*100)) / 100)
		}
	}
	return filterMap
}