package factory

import (
	"os"

	rm "github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/resourcemanager/types"
	providertypes "github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/labels"
	corev1 "k8s.io/api/core/v1"
	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

const (
	tidbWorkerContainerName = "tidb-worker-pool"
	tidbServerPort          = 4000
	tidbStatusPort          = 10080
	tidbServerPortName      = "server"
	tidbStatusPortName      = "status"

	certVolName = "cert-volume"
	// The following environment variables are referenced by worker startup script.
	clusterCAPathEnvName   = "CLUSTER_CA"
	clusterCertPathEnvName = "CLUSTER_CERT"
	clusterKeyPathEnvName  = "CLUSTER_KEY"

	// The following mount paths and filenames are referenced by EnvoyFilter.
	certMountPath   = "/server-cert"
	clusterCAFile   = certMountPath + "/ca.crt"
	clusterCertFile = certMountPath + "/tls.crt"
	clusterKeyFile  = certMountPath + "/tls.key"

	// Annotations for prometheus scrape.
	prometheusIOScrape = "true"
	prometheusIOPath   = "/metrics"
	prometheusIOPort   = ":10080"

	// labels
	labelTiDBWorker      = "tidb-worker"
	labelComponentWorker = "worker"
	labelPoolManager     = "pool-manager"
)

// WorkerLabels returns the shared labels for every worker pods and services.
func WorkerLabels() map[string]string {
	return map[string]string{
		labels.NameLabelKey:      labelTiDBWorker,
		labels.ComponentLabelKey: labelComponentWorker,
		labels.ManagedByKey:      labelPoolManager,
	}
}

// WorkerPod constructs the worker pod from the given spec.
func WorkerPod(
	worker rm.Worker,
	pool rm.WorkerPool,
	namespace string,
	revision string,
	spec *corev1.PodSpec,
) *corev1.Pod {
	podLabels := WorkerLabels()
	podLabels[providertypes.RevisionLabelKey] = revision
	podLabels[providertypes.WorkerLabelKey] = string(worker)
	podLabels[providertypes.WorkerPoolLabelKey] = string(pool)
	return &corev1.Pod{
		ObjectMeta: v1.ObjectMeta{
			Name:      workerPodName(worker, revision),
			Namespace: namespace,
			Labels:    podLabels,
			Annotations: map[string]string{
				labels.PrometheusIOScrapeLabel: prometheusIOScrape,
				labels.PrometheusIOPathLabel:   prometheusIOPath,
				labels.PrometheusIOPortLabel:   prometheusIOPort,
			},
		},
		Spec: *spec,
	}
}

// WorkerSvc returns the headless service that exposes the worker with given worker handle.
func WorkerSvc(worker rm.Worker, pool rm.WorkerPool, namespace string) *corev1.Service {
	svcLabels := WorkerLabels()
	svcLabels[providertypes.WorkerLabelKey] = string(worker)
	svcLabels[providertypes.WorkerPoolLabelKey] = string(pool)
	return &corev1.Service{
		ObjectMeta: v1.ObjectMeta{
			Name:      svcName(worker),
			Namespace: namespace,
			Labels:    svcLabels,
		},
		Spec: corev1.ServiceSpec{
			ClusterIP: corev1.ClusterIPNone,
			Selector: map[string]string{
				providertypes.WorkerLabelKey:     string(worker),
				providertypes.WorkerPoolLabelKey: string(pool),
			},
			Ports: []corev1.ServicePort{
				{Port: tidbServerPort, Name: tidbServerPortName, Protocol: corev1.ProtocolTCP},
				{Port: tidbStatusPort, Name: tidbStatusPortName, Protocol: corev1.ProtocolTCP},
			},
		},
	}
}

func WorkerSpec(image string, resources *corev1.ResourceRequirements, configMap, secret string) *corev1.PodSpec {
	podSpec := corev1.PodSpec{
		Tolerations: []corev1.Toleration{
			{
				Key:      labels.ServerlessComponentNodeLabelKey,
				Operator: corev1.TolerationOpEqual,
				Effect:   corev1.TaintEffectNoSchedule,
				Value:    "burstable",
			},
		},
		Affinity: &corev1.Affinity{
			NodeAffinity: &corev1.NodeAffinity{
				RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{
					NodeSelectorTerms: []corev1.NodeSelectorTerm{
						{
							MatchExpressions: []corev1.NodeSelectorRequirement{
								{
									Key:      labels.ServerlessComponentNodeLabelKey,
									Operator: corev1.NodeSelectorOpIn,
									Values:   []string{"burstable"},
								},
							},
						},
					},
				},
			},
		},
		Containers: []corev1.Container{
			{
				Name:  tidbWorkerContainerName,
				Image: image,
				Ports: []corev1.ContainerPort{
					{ContainerPort: tidbServerPort, Name: tidbServerPortName, Protocol: corev1.ProtocolTCP},
					{ContainerPort: tidbStatusPort, Name: tidbStatusPortName, Protocol: corev1.ProtocolTCP},
				},
				Command: []string{"/bin/sh", "/usr/local/bin/tidb_start_script.sh"},
				Env: []corev1.EnvVar{
					{
						Name:  "TIDB_WORKER_NS",
						Value: os.Getenv("NAMESPACE"),
					},
				},
				VolumeMounts: []corev1.VolumeMount{
					{
						Name:      "config",
						ReadOnly:  true,
						MountPath: "/etc/tidb",
					},
					{
						Name:      "startup-script",
						ReadOnly:  true,
						MountPath: "/usr/local/bin",
					},
				},
			},
		},
		Volumes: []corev1.Volume{
			{
				Name: "config",
				VolumeSource: corev1.VolumeSource{
					ConfigMap: &corev1.ConfigMapVolumeSource{
						LocalObjectReference: corev1.LocalObjectReference{
							Name: configMap,
						},
						Items: []corev1.KeyToPath{
							{
								Key:  "tidb-config",
								Path: "tidb.toml",
							},
							{
								Key:  "sem-config",
								Path: "sem-config.json",
							},
						},
					},
				},
			},
			{
				Name: "startup-script",
				VolumeSource: corev1.VolumeSource{
					ConfigMap: &corev1.ConfigMapVolumeSource{
						LocalObjectReference: corev1.LocalObjectReference{
							Name: configMap,
						},
						Items: []corev1.KeyToPath{
							{
								Key:  "startup-script",
								Path: "tidb_start_script.sh",
							},
						},
					},
				},
			},
		},
	}

	// Set resource if specified.
	if resources != nil {
		podSpec.Containers[0].Resources = *resources
	}

	// Set secret if specified.
	if secret != "" {
		addCertMount(&podSpec.Containers[0])
		addCertVol(&podSpec, secret)
	}

	return &podSpec
}

func addCertMount(container *corev1.Container) {
	certMount := corev1.VolumeMount{
		Name:      certVolName,
		ReadOnly:  true,
		MountPath: certMountPath,
	}
	container.VolumeMounts = append(container.VolumeMounts, certMount)
	envVars := []corev1.EnvVar{
		{
			Name:  clusterCAPathEnvName,
			Value: clusterCAFile,
		},
		{
			Name:  clusterCertPathEnvName,
			Value: clusterCertFile,
		},
		{
			Name:  clusterKeyPathEnvName,
			Value: clusterKeyFile,
		},
	}
	container.Env = append(container.Env, envVars...)
}

func addCertVol(podSpec *corev1.PodSpec, secret string) {
	certVol := corev1.Volume{
		Name: certVolName,
		VolumeSource: corev1.VolumeSource{
			Secret: &corev1.SecretVolumeSource{
				SecretName: secret,
			},
		},
	}
	podSpec.Volumes = append(podSpec.Volumes, certVol)
}