package nodedrainer

import (
	"context"
	"strings"
	"time"

	"github.com/go-logr/logr"
	"github.com/tidbcloud/aws-shared-provider/pkg/manager"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"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"
)

const (
	activeOnDemandNodeDuration = 4 * time.Hour
)

// NodeDrainer is responsible for draining on-demand nodes to spot nodes
type NodeDrainer struct {
	logr.Logger

	ctx           context.Context
	client        client.Client
	recorder      record.EventRecorder
	checkInterval time.Duration
}

func New(mgr manager.Manager, checkInterval time.Duration) *NodeDrainer {
	return &NodeDrainer{
		Logger: ctrl.Log.WithName("nodedrainer"),

		client:        mgr.GetClient(),
		recorder:      mgr.GetEventRecorderFor("NodeDrainer"),
		checkInterval: checkInterval,
	}
}

func (dn *NodeDrainer) Start(ctx context.Context) error {
	dn.ctx = ctx
	dn.Info("start nodedrainer")

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

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

func (dn *NodeDrainer) do() {
	start := time.Now()
	dn.V(1).Info("start nodedrainer check cycle")
	defer func() {
		dn.V(1).Info("finish nodedrainer check cycle", "duration", time.Since(start))
	}()

	onDemandNodeList := &corev1.NodeList{}
	if err := dn.client.List(dn.ctx, onDemandNodeList, client.MatchingLabels{
		labels.NodeDrainerCapacityTypeLabelKey: labels.NodeDrainerCapacityTypeOnDemand,
	}); err != nil {
		dn.Error(err, "failed to list on-demand nodes")
		return
	}
	spotNodeList := &corev1.NodeList{}
	if err := dn.client.List(dn.ctx, spotNodeList, client.MatchingLabels{
		labels.NodeDrainerCapacityTypeLabelKey: labels.NodeDrainerCapacityTypeSpot,
	}); err != nil {
		dn.Error(err, "failed to list spot nodes")
		return
	}

	for nodeGroupComponentZone, nodeGroupList := range dn.filterNodeGroup(onDemandNodeList) {
		dn.ctx = log.IntoContext(dn.ctx, dn.WithValues("nodeGroupComponentZone", nodeGroupComponentZone))
		if haveActiveOnDemandNode(nodeGroupList) || havePreventCAScaleUpNode(nodeGroupComponentZone, spotNodeList) {
			dn.recover(nodeGroupList)
		} else {
			dn.drain(nodeGroupList)
		}
	}
}

func (dn *NodeDrainer) filterNodeGroup(nodeList *corev1.NodeList) map[string][]*corev1.Node {
	var nodeGroupMap = make(map[string][]*corev1.Node)
	for i := range nodeList.Items {
		node := &nodeList.Items[i]
		if node.DeletionTimestamp != nil {
			continue
		}
		nodeGroupName := node.Labels[labels.ServerlessNodeGroupLabelKey]
		if nodeGroupName == "" {
			continue
		}
		nodeGroupComponentZone := nodeGroupName + "/" +
			node.Labels[labels.ServerlessComponentNodeLabelKey] + "/" + node.Labels[corev1.LabelTopologyZone]
		nodeGroupMap[nodeGroupComponentZone] = append(nodeGroupMap[nodeGroupComponentZone], node)
	}
	return nodeGroupMap
}

func haveActiveOnDemandNode(nodeList []*corev1.Node) bool {
	for _, node := range nodeList {
		if time.Since(node.CreationTimestamp.Time) < activeOnDemandNodeDuration {
			return true
		}
	}
	return false
}

func havePreventCAScaleUpNode(nodeGroupZone string, nodeList *corev1.NodeList) bool {
	rlt := strings.Split(nodeGroupZone, "/")
	for i := range nodeList.Items {
		node := &nodeList.Items[i]
		if node.DeletionTimestamp != nil {
			continue
		}
		if node.Labels[labels.ServerlessComponentNodeLabelKey] == rlt[1] && node.Labels[corev1.LabelTopologyZone] == rlt[2] {
			for j := range node.Spec.Taints {
				if node.Spec.Taints[j].Key == types.PreventCAScaleUpTaintKey {
					return true
				}
			}
		}
	}
	return false
}