package nodedrainer

import (
	"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/kubeapi"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/labels"
	corev1 "k8s.io/api/core/v1"
	policyv1 "k8s.io/api/policy/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"
)

func haveNotSafeToEvictPod(podList *corev1.PodList) string {
	for i := range podList.Items {
		pod := &podList.Items[i]
		if pod.DeletionTimestamp == nil && pod.Annotations[types.CASafeToEvict] == "false" {
			return pod.Namespace + "/" + pod.Name
		}
	}
	return ""
}

func (dn *NodeDrainer) drainNode(node *corev1.Node) (evictCount int, err error) {
	podList := &corev1.PodList{}
	if err = dn.client.List(dn.ctx, podList, client.MatchingFields{"spec.nodeName": node.Name}, client.MatchingLabels{
		labels.NameLabelKey:      labels.TidbPodNameLabelValue,
		labels.ComponentLabelKey: labels.TidbPodComponentLabelValue,
	}); err != nil {
		metrics.EvictPodCounter.WithLabelValues("ListPodsFail").Inc()
		return
	}

	if name := haveNotSafeToEvictPod(podList); name != "" {
		dn.recorder.Eventf(node, corev1.EventTypeWarning, "SkipNotSafeToEvictPod",
			"found not safe to evict pod %s, skip drain", name)
		metrics.EvictPodCounter.WithLabelValues("SkipNotSafeToEvictPod").Inc()
		return
	}

	for i := range podList.Items {
		pod := &podList.Items[i]
		if pod.DeletionTimestamp != nil {
			continue
		}
		log.FromContext(dn.ctx).Info("evicting pod", "pod name", pod.Name, "pod ns", pod.Namespace, "node name", node.Name)
		if err = client.IgnoreNotFound(dn.client.SubResource("eviction").Create(
			dn.ctx, pod, &policyv1.Eviction{}),
		); err != nil {
			metrics.EvictPodCounter.WithLabelValues("fail").Inc()
			return
		}
		evictCount++
		metrics.EvictPodCounter.WithLabelValues("success").Inc()
	}
	return
}

func (dn *NodeDrainer) drain(nodeList []*corev1.Node) {
	// first must unschedule all nodes
	for _, node := range nodeList {
		if node.Annotations[types.AnnoNodeDrainTimestamp] == "" || !node.Spec.Unschedulable {
			if err := kubeapi.UpdateNode(dn.ctx, dn.client, func() *corev1.Node {
				node.Annotations = utils.UpdateMap(node.Annotations, types.AnnoNodeDrainTimestamp, time.Now().Format(time.RFC3339))
				node.Spec.Unschedulable = true
				return node
			}); err != nil {
				metrics.UpdateDrainNodeCounter.WithLabelValues("fail").Inc()
				log.FromContext(dn.ctx).Error(err, "failed to unschedule drain node, give up this round", "nodeName", node.Name)
				return
			}
			metrics.UpdateDrainNodeCounter.WithLabelValues("success").Inc()
			dn.recorder.Eventf(node, corev1.EventTypeNormal, "UnscheduleDrainNode", "unschedule drain node %s", node.Name)
		}
	}
	// then drain one node at a time
	for _, node := range nodeList {
		if evictCount, err := dn.drainNode(node); err != nil {
			log.FromContext(dn.ctx).Error(err, "failed to drain on-demand node", "nodeName", node.Name)
		} else if evictCount == 0 {
			// no pod evicted, continue to drain next node
			continue
		}
		return // drain one node at a time
	}
}