package scheduler

import (
	"context"
	"fmt"
	"strings"

	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	clientSet "k8s.io/client-go/kubernetes"
	corelisters "k8s.io/client-go/listers/core/v1"
	"k8s.io/klog/v2"
	framework "k8s.io/kubernetes/pkg/scheduler/framework/v1alpha1"

	"github.com/pingcap/test-infra/services/webhook/pkg/utils"
	"github.com/pingcap/test-infra/services/webhook/pkg/validation"
)

const (
	Name = "test-infra-scheduler"
)

var (
	_ framework.QueueSortPlugin  = &Scheduler{}
	_ framework.FilterPlugin     = &Scheduler{}
	_ framework.PostFilterPlugin = &Scheduler{}
)

type Scheduler struct {
	kubeClient      clientSet.Interface
	podLister       corelisters.PodLister
	namespaceLister corelisters.NamespaceLister
}

func New(_ runtime.Object, fh framework.FrameworkHandle) (framework.Plugin, error) {
	return &Scheduler{
		kubeClient:      fh.ClientSet(),
		podLister:       fh.SharedInformerFactory().Core().V1().Pods().Lister(),
		namespaceLister: fh.SharedInformerFactory().Core().V1().Namespaces().Lister(),
	}, nil
}

func (s *Scheduler) Name() string {
	return Name
}

// Less are used to sort pods in the scheduling queue.
func (s *Scheduler) Less(podA *framework.QueuedPodInfo, podB *framework.QueuedPodInfo) (result bool) {
	// 1. holder pod takes precedence
	if validation.IsHolderPod(podA.Pod) != validation.IsHolderPod(podB.Pod) {
		return validation.IsHolderPod(podA.Pod)
	}

	// 2. pod from older namespace takes precedence
	if podA.Pod.Namespace != podB.Pod.Namespace {
		nsA, err := s.namespaceLister.Get(podA.Pod.Namespace)
		if err != nil {
			klog.ErrorS(err, "get namespace", "ns", podA.Pod.Namespace)
			return false
		}
		nsB, err := s.namespaceLister.Get(podB.Pod.Namespace)
		if err != nil {
			klog.ErrorS(err, "get namespace", "ns", podB.Pod.Namespace)
			return false
		}

		return nsA.CreationTimestamp.Before(&nsB.CreationTimestamp)
	}

	// 3. sort by statefuleSetName and ordinal
	// nothing bad will happen if pod isn't a statefulSet pod
	statefulSetNameA, ordinalA := validation.GetStatefulSetNameAndOrdinal(podA.Pod)
	statefuleSetNameB, ordinalB := validation.GetStatefulSetNameAndOrdinal(podB.Pod)
	if strings.Compare(statefulSetNameA, statefuleSetNameB) == 0 {
		return ordinalA < ordinalB
	}
	return strings.Compare(statefulSetNameA, statefuleSetNameB) == -1
}

// Filter is called by the scheduling framework.
// It marks pod with peer holder pod as Unschedulable to force it into PostFilter stage.
func (s *Scheduler) Filter(_ context.Context, _ *framework.CycleState, pod *corev1.Pod, _ *framework.NodeInfo) *framework.Status {
	if !validation.HasPeerHolderPod(pod) {
		return nil
	}

	holderPodName := validation.BuildHolderPodName(pod.Name)
	_, err := s.podLister.Pods(pod.Namespace).Get(holderPodName)
	switch {
	case apierrors.IsNotFound(err): // holder pod already deleted(preempted), normal pod can be scheduled
		// However, if holder pod is deleted unexpectedly(e.g. manually deleted) instead of being preempted,
		// this branch becomes a downgrade of functionality. Normal pod get scheduled as usual, that might fail due to lack of resource.
		return nil

	case err != nil:
		klog.ErrorS(err, "get pod", "pod", klog.KRef(pod.Namespace, holderPodName))
		return framework.NewStatus(framework.Error, err.Error())

	default: // holder pod exist, should preempt it in PostFilter stage
		return framework.NewStatus(framework.Unschedulable, fmt.Sprintf("pod %s/%s should only be scheduled via preemption", pod.Namespace, pod.Name))
	}
}

// PostFilter is called by the scheduling framework.
// For normal pod with peer holder, it preempts peer holder-pod to get scheduled.
// For holder pod, it try to preempts peer normal-pod(if exist) to get scheduled, otherwise it does nothing.
func (s *Scheduler) PostFilter(ctx context.Context, _ *framework.CycleState, pod *corev1.Pod, _ framework.NodeToStatusMap) (*framework.PostFilterResult, *framework.Status) {
	switch {
	case validation.HasPeerHolderPod(pod):
		holderPodName := validation.BuildHolderPodName(pod.Name)
		holderPod, err := utils.GetPodWithCache(ctx, s.podLister, s.kubeClient, pod.Namespace, holderPodName)
		if err != nil {
			klog.ErrorS(err, "get pod", "pod", klog.KRef(pod.Namespace, holderPodName))
			return nil, framework.NewStatus(framework.Error, err.Error())
		}

		if holderPod.Status.Phase != corev1.PodRunning { // holder pod not running, cannot preempt
			return nil, framework.NewStatus(framework.Unschedulable)
		}

		// preempt
		err = s.kubeClient.CoreV1().Pods(pod.Namespace).Delete(ctx, holderPod.Name, metav1.DeleteOptions{})
		if err != nil {
			klog.ErrorS(err, "preempting pod", "pod", klog.KObj(holderPod))
			return nil, framework.NewStatus(framework.Error, err.Error())
		}

		return &framework.PostFilterResult{NominatedNodeName: holderPod.Spec.NodeName}, framework.NewStatus(framework.Success)
	case validation.IsHolderPod(pod):
		normalPodName := validation.GetNormalPodName(pod.Name)
		normalPod, err := utils.GetPodWithCache(ctx, s.podLister, s.kubeClient, pod.Namespace, normalPodName)
		switch {
		case apierrors.IsNotFound(err):
			return nil, framework.NewStatus(framework.Unschedulable)
		case err != nil:
			klog.ErrorS(err, "get pod", "pod", klog.KRef(pod.Namespace, normalPodName))
			return nil, framework.NewStatus(framework.Error, err.Error())
		}

		// preempt
		err = s.kubeClient.CoreV1().Pods(pod.Namespace).Delete(ctx, normalPod.Name, metav1.DeleteOptions{})
		if err != nil {
			klog.ErrorS(err, "preempting pod", "pod", klog.KObj(normalPod))
			return nil, framework.NewStatus(framework.Error, err.Error())
		}

		return &framework.PostFilterResult{NominatedNodeName: normalPod.Spec.NodeName}, framework.NewStatus(framework.Success)
	default:
		return nil, framework.NewStatus(framework.Unschedulable)
	}
}