/*
Copyright 2021 The Volcano Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package schedulingaction

import (
	"context"
	"fmt"
	"strings"
	"time"

	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/klog/v2"

	schedulingv1beta1 "volcano.sh/apis/pkg/apis/scheduling/v1beta1"

	e2eutil "volcano.sh/volcano/test/e2e/util"
)

const (
	highPriority        = "high-priority"
	middlePriority      = "middle-priority"
	lowPriority         = "low-priority"
	highPriorityValue   = 100
	middlePriorityValue = 50
	lowPriorityValue    = 10
)

var _ = Describe("Job E2E Test", func() {

	var ctx *e2eutil.TestContext
	AfterEach(func() {
		e2eutil.CleanupTestContext(ctx)
	})

	It("schedule high priority job without preemption when resource is enough", func() {
		ctx = e2eutil.InitTestContext(e2eutil.Options{
			PriorityClasses: map[string]int32{
				highPriority: highPriorityValue,
				lowPriority:  lowPriorityValue,
			},
		})

		slot := e2eutil.OneCPU

		job := &e2eutil.JobSpec{
			Tasks: []e2eutil.TaskSpec{
				{
					Img:    e2eutil.DefaultNginxImage,
					Req:    slot,
					Min:    1,
					Rep:    1,
					Labels: map[string]string{schedulingv1beta1.PodPreemptable: "true"},
				},
			},
		}

		job.Name = "preemptee-0"
		job.Pri = lowPriority
		preempteeJob := e2eutil.CreateJob(ctx, job)
		err := e2eutil.WaitTasksReady(ctx, preempteeJob, 1)
		Expect(err).NotTo(HaveOccurred())

		job.Name = "preemptor-0"
		job.Pri = highPriority
		preemptorJob := e2eutil.CreateJob(ctx, job)
		err = e2eutil.WaitTasksReady(ctx, preempteeJob, 1)
		Expect(err).NotTo(HaveOccurred())

		err = e2eutil.WaitTasksReady(ctx, preemptorJob, 1)
		Expect(err).NotTo(HaveOccurred())
	})

	It("schedule high priority job with preemption when idle resource is NOT enough but preemptee resource is enough", func() {
		// Remove enqueue action first because it conflicts with preempt.
		cmc := e2eutil.NewConfigMapCase("volcano-system", "integration-scheduler-configmap")
		modifier := func(sc *e2eutil.SchedulerConfiguration) bool {
			newActions := strings.TrimPrefix(sc.Actions, "enqueue, ")
			if newActions == sc.Actions {
				klog.Warning("There is already no enqueue action")
				return false
			}

			sc.Actions = newActions
			return true
		}
		cmc.ChangeBy(func(data map[string]string) (changed bool, changedBefore map[string]string) {
			return e2eutil.ModifySchedulerConfig(data, modifier)
		})
		defer cmc.UndoChanged()

		ctx = e2eutil.InitTestContext(e2eutil.Options{
			PriorityClasses: map[string]int32{
				highPriority: highPriorityValue,
				lowPriority:  lowPriorityValue,
			},
		})

		slot := e2eutil.OneCPU
		rep := e2eutil.ClusterSize(ctx, slot)

		job := &e2eutil.JobSpec{
			Tasks: []e2eutil.TaskSpec{
				{
					Img:    e2eutil.DefaultNginxImage,
					Req:    slot,
					Min:    1,
					Rep:    rep,
					Labels: map[string]string{schedulingv1beta1.PodPreemptable: "true"},
				},
			},
		}

		job.Name = "preemptee-1"
		job.Pri = lowPriority
		preempteeJob := e2eutil.CreateJob(ctx, job)
		err := e2eutil.WaitTasksReady(ctx, preempteeJob, int(rep))
		Expect(err).NotTo(HaveOccurred())

		job.Name = "preemptor-1"
		job.Pri = highPriority
		job.Min = rep / 2
		preemptorJob := e2eutil.CreateJob(ctx, job)
		err = e2eutil.WaitTasksReady(ctx, preempteeJob, int(rep)/2)
		Expect(err).NotTo(HaveOccurred())

		err = e2eutil.WaitTasksReady(ctx, preemptorJob, int(rep)/2)
		Expect(err).NotTo(HaveOccurred())
	})

	It("preemption doesn't work when podgroup is pending due to insufficient resource", func() {
		ctx = e2eutil.InitTestContext(e2eutil.Options{
			PriorityClasses: map[string]int32{
				highPriority: highPriorityValue,
				lowPriority:  lowPriorityValue,
			},
		})

		pgName := "pending-pg"
		pg := &schedulingv1beta1.PodGroup{
			ObjectMeta: v1.ObjectMeta{
				Namespace: ctx.Namespace,
				Name:      pgName,
			},
			Spec: schedulingv1beta1.PodGroupSpec{
				MinMember:    1,
				MinResources: &e2eutil.ThirtyCPU,
			},
		}
		_, err := ctx.Vcclient.SchedulingV1beta1().PodGroups(ctx.Namespace).Create(context.TODO(), pg, v1.CreateOptions{})
		Expect(err).NotTo(HaveOccurred())

		slot := e2eutil.OneCPU
		rep := e2eutil.ClusterSize(ctx, slot)
		job := &e2eutil.JobSpec{
			Tasks: []e2eutil.TaskSpec{
				{
					Img: e2eutil.DefaultNginxImage,
					Req: slot,
					Min: 1,
					Rep: rep,
				},
			},
		}
		job.Name = "preemptee-2"
		job.Pri = lowPriority
		preempteeJob := e2eutil.CreateJob(ctx, job)
		err = e2eutil.WaitTasksReady(ctx, preempteeJob, int(rep))
		Expect(err).NotTo(HaveOccurred())

		pod := &corev1.Pod{
			TypeMeta: v1.TypeMeta{
				APIVersion: "v1",
				Kind:       "Pod",
			},
			ObjectMeta: v1.ObjectMeta{
				Namespace:   ctx.Namespace,
				Name:        "preemptor-pod",
				Annotations: map[string]string{schedulingv1beta1.KubeGroupNameAnnotationKey: pgName},
			},
			Spec: corev1.PodSpec{
				SchedulerName:     "volcano",
				Containers:        e2eutil.CreateContainers(e2eutil.DefaultNginxImage, "", "", e2eutil.OneCPU, e2eutil.OneCPU, 0),
				PriorityClassName: highPriority,
			},
		}
		// Pod is allowed to be created, preemption does not happen due to PodGroup is in pending state
		_, err = ctx.Kubeclient.CoreV1().Pods(ctx.Namespace).Create(context.TODO(), pod, v1.CreateOptions{})
		Expect(err).NotTo(HaveOccurred())
		// Make sure preempteeJob is not preempted as expected
		err = e2eutil.WaitTasksReady(ctx, preempteeJob, int(rep))
		Expect(err).NotTo(HaveOccurred())
	})

	It("preemption only works in the same queue", func() {
		cmc := e2eutil.NewConfigMapCase("volcano-system", "integration-scheduler-configmap")
		modifier := func(sc *e2eutil.SchedulerConfiguration) bool {
			actions := strings.Split(sc.Actions, ",")
			newActions := make([]string, 0)
			// remove reclaim action
			for _, action := range actions {
				action = strings.TrimSpace(action)
				if action != "reclaim" {
					newActions = append(newActions, action)
				}
			}

			if len(newActions) == len(actions) {
				klog.Warning("There is already no reclaim action")
				return false
			}
			return true
		}
		cmc.ChangeBy(func(data map[string]string) (changed bool, changedBefore map[string]string) {
			return e2eutil.ModifySchedulerConfig(data, modifier)
		})
		defer cmc.UndoChanged()

		ctx = e2eutil.InitTestContext(e2eutil.Options{
			Queues: []string{"q1-preemption", "q2-reference"},
			PriorityClasses: map[string]int32{
				highPriority: highPriorityValue,
				lowPriority:  lowPriorityValue,
			},
		})

		slot := e2eutil.OneCPU
		rep := e2eutil.ClusterSize(ctx, slot)
		job := &e2eutil.JobSpec{
			Tasks: []e2eutil.TaskSpec{
				{
					Img:    e2eutil.DefaultNginxImage,
					Req:    slot,
					Min:    rep / 2,
					Rep:    rep / 2,
					Labels: map[string]string{schedulingv1beta1.PodPreemptable: "true"},
				},
			},
		}

		job.Name = "j1-q1"
		job.Pri = lowPriority
		job.Queue = "q1-preemption"
		queue1Job := e2eutil.CreateJob(ctx, job)
		err := e2eutil.WaitTasksReady(ctx, queue1Job, int(rep)/2)
		Expect(err).NotTo(HaveOccurred())

		job.Name = "j2-q2"
		job.Pri = lowPriority
		job.Queue = "q2-reference"
		queue2Job := e2eutil.CreateJob(ctx, job)
		err = e2eutil.WaitTasksReady(ctx, queue2Job, int(rep)/2)
		Expect(err).NotTo(HaveOccurred())

		job.Name = "j3-q1"
		job.Pri = highPriority
		job.Queue = "q1-preemption"
		job.Tasks[0].Rep = rep
		queue1Job3 := e2eutil.CreateJob(ctx, job)
		err = e2eutil.WaitTasksReady(ctx, queue1Job3, int(rep)/2)
		Expect(err).NotTo(HaveOccurred())
		err = e2eutil.WaitTasksReady(ctx, queue1Job, 0)
		Expect(err).NotTo(HaveOccurred())
	})

	It("preemption doesn't work when total resource of idle resource and preemptee is NOT enough", func() {
		ctx = e2eutil.InitTestContext(e2eutil.Options{
			Queues: []string{"q1-preemption", "q2-reference"},
			PriorityClasses: map[string]int32{
				highPriority: highPriorityValue,
				lowPriority:  lowPriorityValue,
			},
		})

		slot := e2eutil.OneCPU
		rep := e2eutil.ClusterSize(ctx, slot)
		job := &e2eutil.JobSpec{
			Tasks: []e2eutil.TaskSpec{
				{
					Img:    e2eutil.DefaultNginxImage,
					Req:    slot,
					Min:    1,
					Rep:    1,
					Labels: map[string]string{schedulingv1beta1.PodPreemptable: "true"},
				},
			},
		}

		job.Name = "j1-q1"
		job.Pri = lowPriority
		job.Queue = "q1-preemption"
		queue1Job := e2eutil.CreateJob(ctx, job)
		err := e2eutil.WaitTasksReady(ctx, queue1Job, 1)
		Expect(err).NotTo(HaveOccurred())

		job.Name = "j2-q2"
		job.Pri = lowPriority
		job.Queue = "q2-reference"
		job.Tasks[0].Min = rep / 2
		job.Tasks[0].Rep = rep / 2
		queue2Job := e2eutil.CreateJob(ctx, job)
		err = e2eutil.WaitTasksReady(ctx, queue2Job, int(rep)/2)
		Expect(err).NotTo(HaveOccurred())

		job.Name = "j3-q1"
		job.Pri = highPriority
		job.Queue = "q1-preemption"
		job.Tasks[0].Min = rep
		job.Tasks[0].Rep = rep
		queue1Job3 := e2eutil.CreateJob(ctx, job)
		err = e2eutil.WaitTasksReady(ctx, queue1Job3, int(rep))
		Expect(err).To(HaveOccurred())
		err = e2eutil.WaitTasksReady(ctx, queue1Job, 1)
		Expect(err).NotTo(HaveOccurred())
		err = e2eutil.WaitTasksReady(ctx, queue2Job, int(rep)/2)
		Expect(err).NotTo(HaveOccurred())
	})

	It("multi-preemptor-jobs who are in different priority", func() {
		Skip("https://github.com/volcano-sh/volcano/issues/911")
		ctx = e2eutil.InitTestContext(e2eutil.Options{
			Queues: []string{"q1-preemption"},
			PriorityClasses: map[string]int32{
				highPriority:   highPriorityValue,
				middlePriority: middlePriorityValue,
				lowPriority:    lowPriorityValue,
			},
		})

		slot := e2eutil.OneCPU
		rep := e2eutil.ClusterSize(ctx, slot)
		job := &e2eutil.JobSpec{
			Tasks: []e2eutil.TaskSpec{
				{
					Img:    e2eutil.DefaultNginxImage,
					Req:    slot,
					Min:    1,
					Rep:    rep,
					Labels: map[string]string{schedulingv1beta1.PodPreemptable: "true"},
				},
			},
		}

		job.Name = "low-priority-job"
		job.Pri = lowPriority
		job.Queue = "q1-preemption"
		lowPriorityJob := e2eutil.CreateJob(ctx, job)
		err := e2eutil.WaitTasksReady(ctx, lowPriorityJob, int(rep))
		Expect(err).NotTo(HaveOccurred())

		job.Name = "middle-prority-job"
		job.Pri = middlePriority
		job.Queue = "q1-preemption"
		job.Tasks[0].Rep = rep / 2
		job.Tasks[0].Min = rep / 2
		middlePriorityJob := e2eutil.CreateJob(ctx, job)
		err = e2eutil.WaitTasksReady(ctx, middlePriorityJob, int(rep)/2)
		Expect(err).NotTo(HaveOccurred())
		err = e2eutil.WaitTasksReady(ctx, lowPriorityJob, int(rep)/2)
		Expect(err).NotTo(HaveOccurred())

		job.Name = "high-priority-job"
		job.Pri = highPriority
		job.Queue = "q1-preemption"
		job.Tasks[0].Rep = rep
		job.Tasks[0].Min = rep
		highPriorityJob := e2eutil.CreateJob(ctx, job)
		err = e2eutil.WaitTasksReady(ctx, highPriorityJob, int(rep))
		Expect(err).NotTo(HaveOccurred())
		err = e2eutil.WaitTasksReady(ctx, lowPriorityJob, 0)
		Expect(err).NotTo(HaveOccurred())
		err = e2eutil.WaitTasksReady(ctx, middlePriorityJob, 0)
		Expect(err).NotTo(HaveOccurred())
	})

	It("Jobs unschedulable due to scheduling gates will not preempt other jobs despite sufficient preemptor", func() {
		// Remove enqueue action first because it conflicts with preempt.
		cmc := e2eutil.NewConfigMapCase("volcano-system", "integration-scheduler-configmap")
		modifier := func(sc *e2eutil.SchedulerConfiguration) bool {
			newActions := strings.TrimPrefix(sc.Actions, "enqueue, ")
			if newActions == sc.Actions {
				klog.Warning("There is already no enqueue action")
				return false
			}

			sc.Actions = newActions
			return true
		}
		cmc.ChangeBy(func(data map[string]string) (changed bool, changedBefore map[string]string) {
			return e2eutil.ModifySchedulerConfig(data, modifier)
		})
		defer cmc.UndoChanged()

		ctx = e2eutil.InitTestContext(e2eutil.Options{
			PriorityClasses: map[string]int32{
				highPriority: highPriorityValue,
				lowPriority:  lowPriorityValue,
			},
		})

		slot := e2eutil.OneCPU
		rep := e2eutil.ClusterSize(ctx, slot)

		job := &e2eutil.JobSpec{
			Tasks: []e2eutil.TaskSpec{
				{
					Img:    e2eutil.DefaultNginxImage,
					Req:    slot,
					Min:    1,
					Rep:    rep,
					Labels: map[string]string{schedulingv1beta1.PodPreemptable: "true"},
				},
			},
		}

		job.Name = "preemptee-1"
		job.Pri = lowPriority
		preempteeJob := e2eutil.CreateJob(ctx, job)
		err := e2eutil.WaitTasksReady(ctx, preempteeJob, int(rep))
		Expect(err).NotTo(HaveOccurred())

		job.Name = "preemptor-1"
		job.Pri = highPriority
		job.Min = rep / 2
		preemptorJob := e2eutil.CreateJob(ctx, job)
		// All pods of preemptorJob will be created and remain in pending state
		err = e2eutil.WaitTasksPending(ctx, preemptorJob, int(rep)/2)
		Expect(err).NotTo(HaveOccurred())
		// None of the tasks of preemptee should be evicted
		err = e2eutil.WaitTasksReady(ctx, preempteeJob, int(rep))
		Expect(err).NotTo(HaveOccurred())

		// remove gate
		err = e2eutil.RemovePodSchGates(ctx, preemptorJob)
		Expect(err).NotTo(HaveOccurred())

		// half jobs of preemptee will be evicted to make space for preemptorJob
		err = e2eutil.WaitTasksReady(ctx, preempteeJob, int(rep)/2)
		Expect(err).NotTo(HaveOccurred())
		err = e2eutil.WaitTasksReady(ctx, preemptorJob, int(rep)/2)
		Expect(err).NotTo(HaveOccurred())
	})

	It("should preempt low priority pod with anti-affinity constraint", func() {
		// Remove enqueue action first because it conflicts with preempt.
		cmc := e2eutil.NewConfigMapCase("volcano-system", "integration-scheduler-configmap")
		modifier := func(sc *e2eutil.SchedulerConfiguration) bool {
			newActions := strings.TrimPrefix(sc.Actions, "enqueue, ")
			if newActions == sc.Actions {
				klog.Warning("There is already no enqueue action")
				return false
			}

			sc.Configurations = append(sc.Configurations, e2eutil.Configuration{
				Name: "preempt",
				Arguments: map[string]string{
					"enableTopologyAwarePreemption": "true",
				},
			})
			sc.Actions = newActions
			return true
		}
		cmc.ChangeBy(func(data map[string]string) (changed bool, changedBefore map[string]string) {
			return e2eutil.ModifySchedulerConfig(data, modifier)
		})
		defer cmc.UndoChanged()
		// Wait for the scheduler file watcher to reload the refreshed ConfigMap volume.
		time.Sleep(10 * time.Second)

		ctx = e2eutil.InitTestContext(e2eutil.Options{
			PriorityClasses: map[string]int32{
				highPriority:   highPriorityValue,
				middlePriority: middlePriorityValue,
				lowPriority:    lowPriorityValue,
			},
		})

		rep := int32(e2eutil.ClusterNodeNumber(ctx))
		Expect(rep).Should(BeNumerically(">=", 2))
		slots := e2eutil.ClusterSize(ctx, e2eutil.OneCPU) / rep
		Expect(slots).Should(BeNumerically(">=", 3))
		highRep := rep - 1
		highReq := e2eutil.CPUResource(fmt.Sprintf("%d", slots-1))

		blockerLabels := map[string]string{
			"app": "blocker",
		}
		bystanderLabels := map[string]string{
			"app": "bystander",
		}

		blockerJob := &e2eutil.JobSpec{
			Tasks: []e2eutil.TaskSpec{
				{
					Img: e2eutil.DefaultNginxImage,
					Req: e2eutil.OneCPU,
					Min: 1,
					Rep: rep,
					Labels: map[string]string{
						schedulingv1beta1.PodPreemptable: "true",
						"app":                            "blocker",
					},
					Affinity: &corev1.Affinity{
						PodAntiAffinity: &corev1.PodAntiAffinity{
							RequiredDuringSchedulingIgnoredDuringExecution: []corev1.PodAffinityTerm{
								{
									LabelSelector: &metav1.LabelSelector{
										MatchLabels: blockerLabels,
									},
									TopologyKey: "kubernetes.io/hostname",
								},
							},
						},
					},
				},
			},
		}

		blockerJob.Name = "blocker-job"
		blockerJob.Pri = middlePriority
		blockerJobCreated := e2eutil.CreateJob(ctx, blockerJob)
		err := e2eutil.WaitTasksReady(ctx, blockerJobCreated, int(rep))
		Expect(err).NotTo(HaveOccurred())

		bystanderJob := &e2eutil.JobSpec{
			Tasks: []e2eutil.TaskSpec{
				{
					Img: e2eutil.DefaultNginxImage,
					Req: e2eutil.OneCPU,
					Min: 1,
					Rep: rep,
					Labels: map[string]string{
						schedulingv1beta1.PodPreemptable: "true",
						"app":                            "bystander",
					},
					Affinity: &corev1.Affinity{
						PodAffinity: &corev1.PodAffinity{
							RequiredDuringSchedulingIgnoredDuringExecution: []corev1.PodAffinityTerm{
								{
									LabelSelector: &metav1.LabelSelector{
										MatchLabels: blockerLabels,
									},
									TopologyKey: "kubernetes.io/hostname",
								},
							},
						},
						PodAntiAffinity: &corev1.PodAntiAffinity{
							RequiredDuringSchedulingIgnoredDuringExecution: []corev1.PodAffinityTerm{
								{
									LabelSelector: &metav1.LabelSelector{
										MatchLabels: bystanderLabels,
									},
									TopologyKey: "kubernetes.io/hostname",
								},
							},
						},
					},
				},
			},
		}

		bystanderJob.Name = "bystander-job"
		bystanderJob.Pri = lowPriority
		bystanderJobCreated := e2eutil.CreateJob(ctx, bystanderJob)
		err = e2eutil.WaitTasksReady(ctx, bystanderJobCreated, int(rep))
		Expect(err).NotTo(HaveOccurred())

		highPriorityJob := &e2eutil.JobSpec{
			Tasks: []e2eutil.TaskSpec{
				{
					Img: e2eutil.DefaultNginxImage,
					Req: highReq,
					Min: highRep,
					Rep: highRep,
					Labels: map[string]string{
						schedulingv1beta1.PodPreemptable: "true",
						"app":                            "preemptor",
					},
					Affinity: &corev1.Affinity{
						PodAntiAffinity: &corev1.PodAntiAffinity{
							RequiredDuringSchedulingIgnoredDuringExecution: []corev1.PodAffinityTerm{
								{
									LabelSelector: &metav1.LabelSelector{
										MatchLabels: blockerLabels,
									},
									TopologyKey: "kubernetes.io/hostname",
								},
							},
						},
					},
				},
			},
		}

		highPriorityJob.Name = "high-priority-job"
		highPriorityJob.Pri = highPriority
		highPriorityJobCreated := e2eutil.CreateJob(ctx, highPriorityJob)

		err = e2eutil.WaitTasksReady(ctx, highPriorityJobCreated, int(highRep))
		Expect(err).NotTo(HaveOccurred())
		err = e2eutil.WaitTasksReady(ctx, bystanderJobCreated, int(rep))
		Expect(err).NotTo(HaveOccurred())
		err = e2eutil.WaitTasksReadyEx(ctx, blockerJobCreated, map[string]int{
			middlePriority: 1,
		})
		Expect(err).NotTo(HaveOccurred())
	})

})
