/*
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 jobseq

import (
	"context"
	"fmt"
	"strings"

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

	cv1 "k8s.io/api/core/v1"
	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	batch "volcano.sh/apis/pkg/apis/batch/v1alpha1"
	"volcano.sh/volcano/pkg/controllers/job/helpers"
	"volcano.sh/volcano/pkg/controllers/job/plugins/env"
	"volcano.sh/volcano/pkg/controllers/job/plugins/svc"

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

var _ = Describe("Job E2E Test: Test Job Plugins", func() {
	var testCtx *e2eutil.TestContext

	BeforeEach(func() {
		testCtx = e2eutil.InitTestContext(e2eutil.Options{})
		DeferCleanup(e2eutil.CleanupTestContext, testCtx)
	})

	JustAfterEach(func() {
		e2eutil.DumpTestContextIfFailed(testCtx, CurrentSpecReport())
	})

	It("Test SVC Plugin with Node Affinity", func() {
		jobName := "job-with-svc-plugin"
		taskName := "task"
		foundVolume := false

		nodeName, rep := e2eutil.ComputeNode(testCtx, e2eutil.OneCPU)
		Expect(rep).NotTo(Equal(0))

		affinity := &cv1.Affinity{
			NodeAffinity: &cv1.NodeAffinity{
				RequiredDuringSchedulingIgnoredDuringExecution: &cv1.NodeSelector{
					NodeSelectorTerms: []cv1.NodeSelectorTerm{
						{
							MatchFields: []cv1.NodeSelectorRequirement{
								{
									Key:      "metadata.name",
									Operator: cv1.NodeSelectorOpIn,
									Values:   []string{nodeName},
								},
							},
						},
					},
				},
			},
		}

		job := e2eutil.CreateJob(testCtx, &e2eutil.JobSpec{
			Namespace: testCtx.Namespace,
			Name:      jobName,
			Plugins: map[string][]string{
				"svc": {},
			},
			Tasks: []e2eutil.TaskSpec{
				{
					Img:      e2eutil.DefaultNginxImage,
					Req:      e2eutil.OneCPU,
					Min:      1,
					Rep:      1,
					Name:     taskName,
					Affinity: affinity,
				},
			},
		})

		err := e2eutil.WaitJobReady(testCtx, job)
		Expect(err).NotTo(HaveOccurred())

		pluginName := fmt.Sprintf("%s-svc", jobName)
		_, err = testCtx.Kubeclient.CoreV1().ConfigMaps(testCtx.Namespace).Get(context.TODO(),
			pluginName, v1.GetOptions{})
		Expect(err).NotTo(HaveOccurred())

		pod, err := testCtx.Kubeclient.CoreV1().Pods(testCtx.Namespace).Get(context.TODO(),
			fmt.Sprintf(helpers.PodNameFmt, jobName, taskName, 0), v1.GetOptions{})
		Expect(err).NotTo(HaveOccurred())
		for _, volume := range pod.Spec.Volumes {
			if volume.Name == pluginName {
				foundVolume = true
				break
			}
		}
		Expect(foundVolume).To(BeTrue())

		pods := e2eutil.GetTasksOfJob(testCtx, job)
		for _, pod := range pods {
			Expect(pod.Spec.NodeName).To(Equal(nodeName))
		}
	})

	It("Test SSh Plugin with Pod Affinity", func() {
		jobName := "job-with-ssh-plugin"
		taskName := "task"
		foundVolume := false

		_, rep := e2eutil.ComputeNode(testCtx, e2eutil.OneCPU)
		Expect(rep).NotTo(Equal(0))

		labels := map[string]string{"foo": "bar"}

		affinity := &cv1.Affinity{
			PodAffinity: &cv1.PodAffinity{
				RequiredDuringSchedulingIgnoredDuringExecution: []cv1.PodAffinityTerm{
					{
						LabelSelector: &v1.LabelSelector{
							MatchLabels: labels,
						},
						TopologyKey: "kubernetes.io/hostname",
					},
				},
			},
		}

		job := e2eutil.CreateJob(testCtx, &e2eutil.JobSpec{
			Namespace: testCtx.Namespace,
			Name:      jobName,
			Plugins: map[string][]string{
				"ssh": {"--no-root"},
			},
			Tasks: []e2eutil.TaskSpec{
				{
					Img:      e2eutil.DefaultNginxImage,
					Req:      e2eutil.HalfCPU,
					Min:      rep,
					Rep:      rep,
					Affinity: affinity,
					Labels:   labels,
					Name:     taskName,
				},
			},
		})

		err := e2eutil.WaitJobReady(testCtx, job)
		Expect(err).NotTo(HaveOccurred())

		secretName := genSSHSecretName(job)
		_, err = testCtx.Kubeclient.CoreV1().Secrets(testCtx.Namespace).Get(context.TODO(),
			secretName, v1.GetOptions{})
		Expect(err).NotTo(HaveOccurred())

		pod, err := testCtx.Kubeclient.CoreV1().Pods(testCtx.Namespace).Get(context.TODO(),
			fmt.Sprintf(helpers.PodNameFmt, jobName, taskName, 0), v1.GetOptions{})
		Expect(err).NotTo(HaveOccurred())
		for _, volume := range pod.Spec.Volumes {
			if volume.Name == secretName {
				foundVolume = true
				break
			}
		}
		Expect(foundVolume).To(BeTrue())

		pods := e2eutil.GetTasksOfJob(testCtx, job)
		// All pods should be scheduled to the same node.
		nodeName := pods[0].Spec.NodeName
		for _, pod := range pods {
			Expect(pod.Spec.NodeName).To(Equal(nodeName))
		}
	})

	It("Test SVC Plugin with disableNetworkPolicy", func() {
		jobName := "svc-with-disable-network-policy"
		taskName := "task"

		_, rep := e2eutil.ComputeNode(testCtx, e2eutil.OneCPU)
		Expect(rep).NotTo(Equal(0))

		job := e2eutil.CreateJob(testCtx, &e2eutil.JobSpec{
			Namespace: testCtx.Namespace,
			Name:      jobName,
			Plugins: map[string][]string{
				"svc": {"--disable-network-policy=true"},
			},
			Tasks: []e2eutil.TaskSpec{
				{
					Img:  e2eutil.DefaultNginxImage,
					Req:  e2eutil.OneCPU,
					Min:  1,
					Rep:  rep,
					Name: taskName,
				},
			},
		})

		err := e2eutil.WaitJobReady(testCtx, job)
		Expect(err).NotTo(HaveOccurred())

		// Check whether network policy is created with job name
		networkPolicyName := jobName
		_, err = testCtx.Kubeclient.NetworkingV1().NetworkPolicies(testCtx.Namespace).Get(context.TODO(), networkPolicyName, v1.GetOptions{})
		// Error will occur because there is no policy should be created
		Expect(err).To(HaveOccurred())
	})

	It("Test SVC Plugin with injectHostsEnv disabled", func() {
		jobName := "svc-without-hosts-env"
		taskName := "task"

		_, rep := e2eutil.ComputeNode(testCtx, e2eutil.OneCPU)
		Expect(rep).NotTo(Equal(0))

		job := e2eutil.CreateJob(testCtx, &e2eutil.JobSpec{
			Namespace: testCtx.Namespace,
			Name:      jobName,
			Plugins: map[string][]string{
				"svc": {"--inject-hosts-env=false"},
			},
			Tasks: []e2eutil.TaskSpec{
				{
					Img:  e2eutil.DefaultNginxImage,
					Req:  e2eutil.OneCPU,
					Min:  1,
					Rep:  rep,
					Name: taskName,
				},
			},
		})

		err := e2eutil.WaitJobReady(testCtx, job)
		Expect(err).NotTo(HaveOccurred())

		pluginName := fmt.Sprintf("%s-svc", jobName)
		_, err = testCtx.Kubeclient.CoreV1().ConfigMaps(testCtx.Namespace).Get(context.TODO(),
			pluginName, v1.GetOptions{})
		Expect(err).NotTo(HaveOccurred())

		pod, err := testCtx.Kubeclient.CoreV1().Pods(testCtx.Namespace).Get(context.TODO(),
			fmt.Sprintf(helpers.PodNameFmt, jobName, taskName, 0), v1.GetOptions{})
		Expect(err).NotTo(HaveOccurred())

		envNames := []string{
			fmt.Sprintf(svc.EnvTaskHostFmt, strings.ToUpper(taskName)),
			fmt.Sprintf(svc.EnvHostNumFmt, strings.ToUpper(taskName)),
		}

		containers := pod.Spec.Containers
		containers = append(containers, pod.Spec.InitContainers...)
		for _, container := range containers {
			for _, name := range envNames {
				for _, envi := range container.Env {
					Expect(envi.Name).NotTo(Equal(name),
						fmt.Sprintf("container: %s, env name: %s should not be injected", container.Name, name))
				}
			}
		}
	})

	It("Check Functionality of all plugins", func() {
		jobName := "job-with-all-plugin"
		taskName := "task"
		foundVolume := false
		foundEnv := false

		_, rep := e2eutil.ComputeNode(testCtx, e2eutil.OneCPU)
		Expect(rep).NotTo(Equal(0))

		job := e2eutil.CreateJob(testCtx, &e2eutil.JobSpec{
			Namespace: testCtx.Namespace,
			Name:      jobName,
			Plugins: map[string][]string{
				"ssh": {"--no-root"},
				"env": {},
				"svc": {},
			},
			Tasks: []e2eutil.TaskSpec{
				{
					Img:  e2eutil.DefaultNginxImage,
					Req:  e2eutil.OneCPU,
					Min:  1,
					Rep:  rep,
					Name: taskName,
				},
			},
		})

		err := e2eutil.WaitJobReady(testCtx, job)
		Expect(err).NotTo(HaveOccurred())

		secretName := genSSHSecretName(job)
		_, err = testCtx.Kubeclient.CoreV1().Secrets(testCtx.Namespace).Get(context.TODO(),
			secretName, v1.GetOptions{})
		Expect(err).NotTo(HaveOccurred())

		pod, err := testCtx.Kubeclient.CoreV1().Pods(testCtx.Namespace).Get(context.TODO(),
			fmt.Sprintf(helpers.PodNameFmt, jobName, taskName, 0), v1.GetOptions{})
		Expect(err).NotTo(HaveOccurred())
		for _, volume := range pod.Spec.Volumes {
			if volume.Name == secretName {
				foundVolume = true
				break
			}
		}
		Expect(foundVolume).To(BeTrue())

		// Check whether env exists in the containers and initContainers
		containers := pod.Spec.Containers
		containers = append(containers, pod.Spec.InitContainers...)
		envNames := []string{
			env.TaskVkIndex,
			env.TaskIndex,
			fmt.Sprintf(svc.EnvTaskHostFmt, strings.ToUpper(taskName)),
			fmt.Sprintf(svc.EnvHostNumFmt, strings.ToUpper(taskName)),
		}

		for _, container := range containers {
			for _, name := range envNames {
				foundEnv = false
				for _, envi := range container.Env {
					if envi.Name == name {
						foundEnv = true
						break
					}
				}

				Expect(foundEnv).To(BeTrue(),
					fmt.Sprintf("container: %s, env name: %s", container.Name, name))
			}
		}

		// Check whether service is created with job name
		_, err = testCtx.Kubeclient.CoreV1().Services(job.Namespace).Get(context.TODO(), job.Name, v1.GetOptions{})
		Expect(err).NotTo(HaveOccurred())
	})
})

func genSSHSecretName(job *batch.Job) string {
	return job.Name + "-ssh"
}
