package resource

import (
	"testing"

	"github.com/stretchr/testify/require"
	corev1 "k8s.io/api/core/v1"
	kresource "k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

var (
	firefoxName  = "firefox"
	firefox      = "firefox"
	firefoxImage = "docker.io/linuxserver/firefox:latest"
)

func TestAddUIRequestResourcesToPod(t *testing.T) {
	tcs := []struct {
		Pod         *corev1.Pod
		ExpectedPod *corev1.Pod
	}{
		{
			// test case 1: pod requesting UI has label, volume and envvar added
			Pod: &corev1.Pod{
				ObjectMeta: metav1.ObjectMeta{
					Name:      firefoxName,
					Namespace: firefoxName,
				},
				Spec: corev1.PodSpec{
					Containers: []corev1.Container{
						{
							Name:  firefox,
							Image: firefoxImage,
							Resources: corev1.ResourceRequirements{
								Requests: corev1.ResourceList{
									UIRequestResource.ResourceName(): kresource.MustParse("1"),
								},
							},
						},
					},
				},
			},
			ExpectedPod: &corev1.Pod{
				ObjectMeta: metav1.ObjectMeta{
					Name:      firefoxName,
					Namespace: firefoxName,
					Labels:    map[string]string{UIRequestResource.String(): "true"},
				},
				Spec: corev1.PodSpec{
					Containers: []corev1.Container{
						{
							Name:  firefox,
							Image: firefoxImage,
							Resources: corev1.ResourceRequirements{
								Requests: corev1.ResourceList{
									UIRequestResource.ResourceName(): kresource.MustParse("1"),
								},
							},
							Env: []corev1.EnvVar{
								{Name: displayEnvVar, Value: displayZero},
							},
							VolumeMounts: []corev1.VolumeMount{
								{
									Name:      x11Volume,
									MountPath: x11DirPath,
									ReadOnly:  false,
								},
							},
						},
					},
					Volumes: []corev1.Volume{
						{
							Name: x11Volume,
							VolumeSource: corev1.VolumeSource{
								HostPath: &corev1.HostPathVolumeSource{Path: x11DirPath, Type: &hostPathDir},
							},
						},
					},
				},
			},
		},
		{
			// test case 2: pod's existing x11 volume and envvar are not changed
			Pod: &corev1.Pod{
				ObjectMeta: metav1.ObjectMeta{
					Name:      firefoxName,
					Namespace: firefoxName,
					Labels:    map[string]string{UIRequestResource.String(): labelValue},
				},
				Spec: corev1.PodSpec{
					Containers: []corev1.Container{
						{
							Name:  firefox,
							Image: firefoxImage,
							Resources: corev1.ResourceRequirements{
								Requests: corev1.ResourceList{
									UIRequestResource.ResourceName(): kresource.MustParse("1"),
								},
							},
							Env: []corev1.EnvVar{
								{Name: displayEnvVar, Value: ":1.0"},
							},
							VolumeMounts: []corev1.VolumeMount{
								{
									Name:      x11Volume,
									MountPath: x11DirPath,
									ReadOnly:  true,
								},
							},
						},
					},
					Volumes: []corev1.Volume{
						{
							Name: x11Volume,
							VolumeSource: corev1.VolumeSource{
								HostPath: &corev1.HostPathVolumeSource{Path: "/tmp/.X12-unix", Type: &hostPathDir},
							},
						},
					},
				},
			},
			ExpectedPod: &corev1.Pod{
				ObjectMeta: metav1.ObjectMeta{
					Name:      firefoxName,
					Namespace: firefoxName,
					Labels:    map[string]string{UIRequestResource.String(): labelValue},
				},
				Spec: corev1.PodSpec{
					Containers: []corev1.Container{
						{
							Name:  firefox,
							Image: firefoxImage,
							Resources: corev1.ResourceRequirements{
								Requests: corev1.ResourceList{
									UIRequestResource.ResourceName(): kresource.MustParse("1"),
								},
							},
							Env: []corev1.EnvVar{
								{Name: displayEnvVar, Value: ":1.0"},
							},
							VolumeMounts: []corev1.VolumeMount{
								{
									Name:      x11Volume,
									MountPath: x11DirPath,
									ReadOnly:  true,
								},
							},
						},
					},
					Volumes: []corev1.Volume{
						{
							Name: x11Volume,
							VolumeSource: corev1.VolumeSource{
								HostPath: &corev1.HostPathVolumeSource{Path: "/tmp/.X12-unix", Type: &hostPathDir},
							},
						},
					},
				},
			},
		},
		{
			// test case 3: multiple containers, only containers with requests mutated, existing fields ignored
			Pod: &corev1.Pod{
				ObjectMeta: metav1.ObjectMeta{
					Name:      firefoxName,
					Namespace: firefoxName,
					Labels: map[string]string{
						UIRequestResource.String(): labelValue,
						"app.kubernetes.io/name":   firefoxName,
					},
				},
				Spec: corev1.PodSpec{
					Containers: []corev1.Container{
						{
							Name:  firefox,
							Image: firefoxImage,
							Resources: corev1.ResourceRequirements{
								Requests: corev1.ResourceList{
									UIRequestResource.ResourceName(): kresource.MustParse("1"),
								},
							},
							Env: []corev1.EnvVar{
								{Name: displayEnvVar, Value: ":1.0"},
							},
						},
						{
							Name:  "chromium",
							Image: "docker.io/linuxserver/chromium",
							Resources: corev1.ResourceRequirements{
								Requests: corev1.ResourceList{
									UIRequestResource.ResourceName(): kresource.MustParse("1"),
								},
							},
							VolumeMounts: []corev1.VolumeMount{
								{
									Name:      x11Volume,
									MountPath: x11DirPath,
									ReadOnly:  true,
								},
							},
						},
						{
							Name:  "ubuntu",
							Image: "docker.io/ubuntu",
						},
					},
				},
			},
			ExpectedPod: &corev1.Pod{
				ObjectMeta: metav1.ObjectMeta{
					Name:      firefoxName,
					Namespace: firefoxName,
					Labels: map[string]string{
						UIRequestResource.String(): labelValue,
						"app.kubernetes.io/name":   firefoxName,
					},
				},
				Spec: corev1.PodSpec{
					Containers: []corev1.Container{
						{
							Name:  firefox,
							Image: firefoxImage,
							Resources: corev1.ResourceRequirements{
								Requests: corev1.ResourceList{
									UIRequestResource.ResourceName(): kresource.MustParse("1"),
								},
							},
							Env: []corev1.EnvVar{
								{Name: displayEnvVar, Value: ":1.0"},
							},
							VolumeMounts: []corev1.VolumeMount{
								{
									Name:      x11Volume,
									MountPath: x11DirPath,
									ReadOnly:  false,
								},
							},
						},
						{
							Name:  "chromium",
							Image: "docker.io/linuxserver/chromium",
							Resources: corev1.ResourceRequirements{
								Requests: corev1.ResourceList{
									UIRequestResource.ResourceName(): kresource.MustParse("1"),
								},
							},
							Env: []corev1.EnvVar{
								{Name: displayEnvVar, Value: displayZero},
							},
							VolumeMounts: []corev1.VolumeMount{
								{
									Name:      x11Volume,
									MountPath: x11DirPath,
									ReadOnly:  true,
								},
							},
						},
						{
							Name:  "ubuntu",
							Image: "docker.io/ubuntu",
						},
					},
					Volumes: []corev1.Volume{
						{
							Name: x11Volume,
							VolumeSource: corev1.VolumeSource{
								HostPath: &corev1.HostPathVolumeSource{Path: x11DirPath, Type: &hostPathDir},
							},
						},
					},
				},
			},
		},
	}
	for idx, tc := range tcs {
		pod := tc.Pod.DeepCopy()
		podTemplate := corev1.PodTemplateSpec{Spec: pod.Spec}
		podTemplate = InjectResourceIntoPod(podTemplate, UIRequestResource)
		require.Equal(t, tc.ExpectedPod.Spec, podTemplate.Spec, "test case %d failed: pod not as expected", idx+1)
	}
}