package job

import (
	"cmp"
	"context"
	"fmt"
	"math"
	"os"
	"strconv"
	"strings"

	"github.com/aws/smithy-go/ptr"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/labels"
	batchv1 "k8s.io/api/batch/v1"
	corev1 "k8s.io/api/core/v1"
	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/yaml"
	k8s "sigs.k8s.io/controller-runtime/pkg/client"
)

const (
	scalerConfigKey  = "scaler-config"
	configFileKey    = "config-file"
	semConfigKey     = "sem-config"
	startupScriptKey = "startup-script"

	tidbWorkerContainerName = "tidb-worker-job"
	tidbWorkerScalerName    = "tidb-worker-scaler"

	labelTiDBWorker      = "tidb-worker"
	labelComponentWorker = "worker"
	serviceAccountName   = "access-s3-sa"

	prometheusIOScrape = "true"
	prometheusIOPath   = "/metrics"
	prometheusIOPort   = ":10080"

	metricsCollectionIntervalSecond = 15

	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"

	defaultComponentNode = "burstable"
)

func mkConfigKey(paths ...string) string {
	return strings.Join(paths, ".")
}

type Type string

const (
	TypeGC          Type = "gc"
	TypeGCV2        Type = "gcv2"
	TypeDDL         Type = "ddl"
	TypeImportInto  Type = "import-into"
	TypeBatch       Type = "batch"
	TypeRemoteQuery Type = "remote-query"
	TypeTTL         Type = "ttl"
	TypeShared      Type = "shared" // shared worker for ddl/batch tasks
	TypeAutoAnalyze Type = "auto-analyze"
)

type Builder struct {
	namespace string
	client    k8s.Client
}

func NewBuilder(namespace string, client k8s.Client) *Builder {
	return &Builder{
		namespace: namespace,
		client:    client,
	}
}

type Options struct {
	TiDBPool         string
	WorkerImage      string
	KeyspaceName     string
	JobType          Type
	JobName          string
	ExecID           string
	SafeToEvict      bool
	WorkerCertSecret string
	ComponentNode    string
}

func (f *Builder) Build(ctx context.Context, opts Options) (*batchv1.Job, error) {
	template, err := f.acquireWorkerConfig(ctx, opts.JobType, opts.TiDBPool, opts.WorkerImage)
	if err != nil {
		return nil, err
	}

	job := f.build(template, opts)

	return job, nil
}

func (f *Builder) acquireWorkerConfig(ctx context.Context, jobType Type, tidbPool, workerImage string) (*workerTemplate, error) {
	poolConfig := corev1.ConfigMap{}

	err := f.client.Get(ctx, k8s.ObjectKey{
		Namespace: tidbPool,
		Name:      types.ServerlessPoolConfig,
	}, &poolConfig)
	if err != nil {
		return nil, err
	}

	workerConfig := corev1.ConfigMap{}
	err = f.client.Get(ctx, k8s.ObjectKey{
		Namespace: f.namespace,
		Name:      scalerConfigKey,
	}, &workerConfig)
	if err != nil {
		return nil, err
	}

	resolvedResourceKey := findConfigKey(&workerConfig, tidbPool, string(jobType), types.TidbResourcesKey)
	if resolvedResourceKey == "" {
		return nil, fmt.Errorf("%s nof found in %s, pool: %s, jobType: %s",
			types.TidbResourcesKey, scalerConfigKey, tidbPool, string(jobType))
	}
	resources := corev1.ResourceRequirements{}
	err = yaml.Unmarshal(
		[]byte(workerConfig.Data[resolvedResourceKey]),
		&resources,
	)
	if err != nil {
		return nil, err
	}

	resolvedTiDBConfigKey := findConfigKey(&workerConfig, tidbPool, string(jobType), configFileKey)
	resolvedStartupScriptKey := findConfigKey(&workerConfig, tidbPool, string(jobType), startupScriptKey)
	resolvedSEMConfigKey := findConfigKey(&workerConfig, tidbPool, string(jobType), semConfigKey)

	if resolvedTiDBConfigKey == "" {
		return nil, fmt.Errorf("%s nof found in %s, pool: %s, jobType: %s",
			configFileKey, scalerConfigKey, tidbPool, string(jobType))
	}
	if resolvedStartupScriptKey == "" {
		return nil, fmt.Errorf("%s nof found in %s, pool: %s, jobType: %s",
			startupScriptKey, scalerConfigKey, tidbPool, string(jobType))
	}
	// sem-config is optional so we skip checking.

	// If workerImage not specified. Use the same image as the tidb pool.
	if workerImage == "" {
		workerImage = poolConfig.Data[types.TidbImageKey]
	}

	return &workerTemplate{
		tidbPool:         tidbPool,
		image:            workerImage,
		resources:        resources,
		tidbConfigKey:    resolvedTiDBConfigKey,
		startupScriptKey: resolvedStartupScriptKey,
		semConfigKey:     resolvedSEMConfigKey,
	}, nil
}

// findConfigKey retrieves the key for the specific type of config in a hierarchical manner.
// The hierarchy is as follows:
// 1. look for the most specific configuration: {pool}.{worker_type}.{config_type}
// 2. If not found, it will search for: {worker_type}.{config_type}
// 3. If still not found, it will search for: {config_type}
// An empty string is returned if none of the above is found.
func findConfigKey(cm *corev1.ConfigMap, pool, workerType, configType string) string {
	keys := []string{
		mkConfigKey(pool, workerType, configType),
		mkConfigKey(workerType, configType),
		configType,
	}
	for _, key := range keys {
		if _, exists := cm.Data[key]; exists {
			return key
		}
	}
	return ""
}

//nolint:funlen // not need
func (f *Builder) build(
	template *workerTemplate,
	opts Options,
) *batchv1.Job {
	componentNode := cmp.Or(opts.ComponentNode, defaultComponentNode)

	container := corev1.Container{
		Name:      tidbWorkerContainerName,
		Image:     template.image,
		Resources: template.resources,
		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"},
		// Setup keyspace name env variable.
		Env: []corev1.EnvVar{
			{
				Name:  "KEYSPACE_NAME",
				Value: opts.KeyspaceName,
			},
			{
				Name:  "TIDB_WORKER_NS",
				Value: os.Getenv("NAMESPACE"),
			},
			{
				Name:  "EXEC_ID",
				Value: opts.ExecID,
			},
			{
				Name:  "TIDB_POOL",
				Value: opts.TiDBPool,
			},
			{
				Name: "POD_IP",
				ValueFrom: &corev1.EnvVarSource{
					FieldRef: &corev1.ObjectFieldSelector{
						FieldPath: "status.podIP",
					},
				},
			},
			{
				Name: "NAMESPACE",
				ValueFrom: &corev1.EnvVarSource{
					FieldRef: &corev1.ObjectFieldSelector{
						FieldPath: "metadata.namespace",
					},
				},
			},
		},
		VolumeMounts: []corev1.VolumeMount{
			{
				Name:      "config",
				ReadOnly:  true,
				MountPath: "/etc/tidb",
			},
			{
				Name:      "startup-script",
				ReadOnly:  true,
				MountPath: "/usr/local/bin",
			},
		},
		Lifecycle: &corev1.Lifecycle{
			// Sleep for 60 seconds before exiting.
			PreStop: &corev1.LifecycleHandler{
				Exec: &corev1.ExecAction{
					Command: []string{"bin/sh", "-c", "sleep " + strconv.Itoa(metricsCollectionIntervalSecond)},
				},
			},
		},
	}
	if opts.WorkerCertSecret != "" {
		addCertMount(&container)
	}

	configVol := corev1.Volume{
		Name: "config",
		VolumeSource: corev1.VolumeSource{
			ConfigMap: &corev1.ConfigMapVolumeSource{
				LocalObjectReference: corev1.LocalObjectReference{
					Name: scalerConfigKey,
				},
				Items: []corev1.KeyToPath{
					{
						Key:  template.tidbConfigKey,
						Path: "tidb.toml",
					},
				},
			},
		},
	}
	if template.semConfigKey != "" {
		configVol.VolumeSource.ConfigMap.Items = append(
			configVol.VolumeSource.ConfigMap.Items,
			corev1.KeyToPath{
				Key:  semConfigKey,
				Path: "sem-config.json",
			})
	}

	startupScriptVol := corev1.Volume{
		Name: "startup-script",
		VolumeSource: corev1.VolumeSource{
			ConfigMap: &corev1.ConfigMapVolumeSource{
				LocalObjectReference: corev1.LocalObjectReference{
					Name: scalerConfigKey,
				},
				Items: []corev1.KeyToPath{
					{
						Key:  template.startupScriptKey,
						Path: "tidb_start_script.sh",
					},
				},
			},
		},
	}

	jobLabels := map[string]string{
		labels.NameLabelKey:        labelTiDBWorker,
		labels.ComponentLabelKey:   labelComponentWorker,
		labels.InstanceLabelKey:    opts.TiDBPool + "-" + opts.JobName,
		types.KeyspaceNameLabelKey: opts.KeyspaceName,
		types.TiDBPoolLabelKey:     template.tidbPool,
		types.JobTypeLabelKey:      string(opts.JobType),
		labels.ManagedByKey:        tidbWorkerScalerName,
	}
	jobAnnotations := map[string]string{
		labels.PrometheusIOScrapeLabel: prometheusIOScrape,
		labels.PrometheusIOPathLabel:   prometheusIOPath,
		labels.PrometheusIOPortLabel:   prometheusIOPort,
	}
	if !opts.SafeToEvict {
		jobAnnotations[types.CASafeToEvict] = "false"
	}
	job := &batchv1.Job{
		ObjectMeta: v1.ObjectMeta{
			Name:        opts.JobName,
			Namespace:   f.namespace,
			Labels:      jobLabels,
			Annotations: jobAnnotations,
		},
		Spec: batchv1.JobSpec{
			Template: corev1.PodTemplateSpec{
				ObjectMeta: v1.ObjectMeta{
					Labels:      jobLabels,
					Annotations: jobAnnotations,
				},
				Spec: corev1.PodSpec{
					Tolerations: []corev1.Toleration{{
						Key:      labels.ServerlessComponentNodeLabelKey,
						Operator: corev1.TolerationOpEqual,
						Effect:   corev1.TaintEffectNoSchedule,
						Value:    componentNode,
					}},
					NodeSelector:                  map[string]string{labels.ServerlessComponentNodeLabelKey: componentNode},
					ServiceAccountName:            serviceAccountName,
					Containers:                    []corev1.Container{container},
					RestartPolicy:                 "OnFailure",
					TerminationGracePeriodSeconds: ptr.Int64(metricsCollectionIntervalSecond),
					Volumes:                       []corev1.Volume{configVol, startupScriptVol},
				},
			},
			BackoffLimit:            ptr.Int32(10),
			TTLSecondsAfterFinished: ptr.Int32(0),
		},
	}
	if opts.WorkerCertSecret != "" {
		addCertVol(job, opts.WorkerCertSecret)
	}
	return job
}

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(job *batchv1.Job, secret string) {
	certVol := corev1.Volume{
		Name: certVolName,
		VolumeSource: corev1.VolumeSource{
			Secret: &corev1.SecretVolumeSource{
				SecretName: secret,
			},
		},
	}
	job.Spec.Template.Spec.Volumes = append(job.Spec.Template.Spec.Volumes, certVol)
}

type workerTemplate struct {
	tidbPool  string
	image     string
	resources corev1.ResourceRequirements

	tidbConfigKey    string
	startupScriptKey string
	semConfigKey     string
}

// countRunningJobs returns the number of running jobs matching the given label.
func countRunningJobs(clusterCli k8s.Client, namespace string, labels map[string]string) (int, error) {
	jobs := batchv1.JobList{}
	err := clusterCli.List(context.Background(), &jobs,
		k8s.InNamespace(namespace),
		k8s.MatchingLabels(labels),
	)
	if err != nil {
		return 0, err
	}

	count := 0
	for _, job := range jobs.Items {
		if job.Status.Active > 0 {
			count++
		}
	}
	return count, nil
}

// CheckJobQuota checks how many jobs can be created for the task with target label key and values.
// It returns math.MaxInt when given limit of 0.
func CheckJobQuota(clusterCli k8s.Client, namespace, labelKey, labelValue string, limit int) (int, error) {
	if limit <= 0 {
		return math.MaxInt, nil
	}
	currentRunningJobs, err := countRunningJobs(clusterCli, namespace, map[string]string{labelKey: labelValue})
	if err != nil {
		return 0, err
	}
	return limit - currentRunningJobs, nil
}

// CheckJobExists check if the given keyspace already have a job running.
func CheckJobExists(clusterCli k8s.Client, namespace, keyspaceName, jobType string) (bool, error) {
	labelSelector := map[string]string{
		types.KeyspaceNameLabelKey: keyspaceName,
		types.JobTypeLabelKey:      jobType,
	}
	currentRunningJobs, err := countRunningJobs(clusterCli, namespace, labelSelector)
	if err != nil {
		return false, err
	}
	return currentRunningJobs > 0, nil
}