// Package common add by blacknc <blacknc@163.com> at 2024/08/28 17:17:20
package common

import (
	"context"
	"errors"
	"fmt"
	"strconv"
	"time"

	batchv1 "k8s.io/api/batch/v1"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"

	"github.com/tidbcloud/eks-provider/controllers/imports/task"
	"github.com/tidbcloud/eks-provider/controllers/utils"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/controllers/common"
	"github.com/tidbcloud/infra-provider/pkg/deploy"
	"github.com/tidbcloud/infra-provider/pkg/k8s"
	"github.com/tidbcloud/infra-provider/pkg/namer"
	"github.com/tidbcloud/infra-provider/pkg/util"
	"github.com/tidbcloud/infra-provider/util/collections"
)

const (
	lightningPort            = 8289
	lightningImageAnnotation = "import.tidbcloud.com/image"
)

var (
	lightningToleration = corev1.Toleration{
		Key:      "dedicated",
		Operator: "Equal",
		Value:    "tidb-" + task.LightningComponent,
		Effect:   corev1.TaintEffectNoSchedule,
	}
)

type Job struct {
	*task.BaseTask
}

func (t *Job) Name() string {
	return "Job"
}

func NewJob(ctx *task.Context) *Job {
	return &Job{
		BaseTask: &task.BaseTask{Context: ctx},
	}
}

func (t *Job) Create(ctx context.Context) deploy.Result {

	job, err := t.describeJob()
	if err != nil {
		return deploy.Error(err)
	}

	clusterCli := t.EnvProvider.GetClient()
	if _, err := controllerutil.CreateOrUpdate(ctx, clusterCli, job, func() error { return nil }); err != nil {
		return deploy.Error(err)
	}
	return deploy.Done()
}

func (t *Job) Delete(ctx context.Context) deploy.Result {
	if !t.IsClusterExist() {
		return deploy.Done()
	}

	job := &batchv1.Job{
		ObjectMeta: metav1.ObjectMeta{
			Name:      t.GetJobName(),
			Namespace: t.GetClusterNameSpace(),
		},
	}

	// Delete job
	clusterCli := t.EnvProvider.GetClient()

	err := clusterCli.Get(ctx, client.ObjectKey{Name: t.GetJobName(), Namespace: t.GetClusterNameSpace()}, job)
	if err != nil {
		if !apierrors.IsNotFound(err) {
			t.Logger.Error(err, "Failed to get job")
			return deploy.Error(err)
		}
		// If job not found, we're done
		return deploy.Done()
	}

	if err := k8s.DeleteObject(ctx, clusterCli, job); err != nil {
		return deploy.Error(err)
	}

	// delete pods
	podList := &corev1.PodList{}
	if err := clusterCli.List(ctx, podList, client.InNamespace(t.GetClusterNameSpace()),
		client.MatchingLabels(map[string]string{"job-name": t.GetJobName()})); err == nil {
		for _, pod := range podList.Items {
			if err := k8s.DeleteObject(ctx, clusterCli, &pod); err != nil {
				return deploy.Error(err)
			}
		}
	}

	return deploy.Done()
}

func (t *Job) PostDelete(ctx context.Context) deploy.Result {
	// If cluster doesn't exist, we can't check for job/pod deletion completion
	if !t.IsClusterExist() {
		return deploy.Done()
	}

	clusterCli := t.EnvProvider.GetClient()

	// Check if any pods associated with this job still exist
	podList := &corev1.PodList{}
	if err := clusterCli.List(ctx, podList, client.InNamespace(t.GetClusterNameSpace()),
		client.MatchingLabels(map[string]string{"job-name": t.GetJobName()})); err == nil {
		if len(podList.Items) > 0 {
			t.Logger.Info("Waiting for pods to be deleted", "remaining", len(podList.Items))
			return deploy.Error(fmt.Errorf("waiting for %d pods to be deleted", len(podList.Items))).WithRequeuAfter(time.Second * 20)
		}
	} else {
		t.Logger.Error(err, "Failed to list pods for job")
	}

	// Check if job is deleted
	job := &batchv1.Job{
		ObjectMeta: metav1.ObjectMeta{
			Name:      t.GetJobName(),
			Namespace: t.GetClusterNameSpace(),
		},
	}
	return deploy.ErrorOrDone(k8s.DeleteObjectCompleted(ctx, clusterCli, job))
}

func (t *Job) describeJob() (*batchv1.Job, error) {

	ns := t.GetClusterNameSpace()
	name := t.GetJobName()
	tcName := t.Namer.GetTiDBClusterName()

	volumes := []corev1.Volume{
		{
			Name: "config",
			VolumeSource: corev1.VolumeSource{
				Secret: &corev1.SecretVolumeSource{
					SecretName: name,
					Items: []corev1.KeyToPath{
						{Key: lightningConfigFileName, Path: lightningConfigFileName},
					},
				},
			},
		},
	}
	mounts := []corev1.VolumeMount{
		{
			Name:      "config",
			MountPath: "/etc/tidb-lightning",
		},
	}

	clusterCertSecretName := getClusterCertSecretName(tcName)
	tidbServerCertSecretName := getTiDBServerCertSecretName(tcName)
	volumes = append(volumes, corev1.Volume{
		Name: "lightning-tls",
		VolumeSource: corev1.VolumeSource{
			Secret: &corev1.SecretVolumeSource{
				SecretName: clusterCertSecretName,
			},
		},
	}, corev1.Volume{
		Name: "tidb-server-cert",
		VolumeSource: corev1.VolumeSource{
			Secret: &corev1.SecretVolumeSource{
				SecretName: tidbServerCertSecretName,
			},
		},
	})
	mounts = append(mounts, corev1.VolumeMount{
		Name:      "lightning-tls",
		MountPath: task.ClusterCertMountDir,
	}, corev1.VolumeMount{
		Name:      "tidb-server-cert",
		MountPath: task.TiDBServerCertMountDir,
	})

	volumes = append(volumes, t.LightningVolumes.Volumes...)
	mounts = append(mounts, t.LightningVolumes.Mounts...)
	envs := t.LightningVolumes.Envs

	volumes, mounts = t.mountPVCVolume(volumes, mounts, t.Import.Spec.StorageRequest == "")

	image := util.GetAnnotation(t.Import, lightningImageAnnotation)
	if image == "" {
		return nil, errors.New("image information not found in the annotation of import cr")
	}
	t.Logger.Info("lightning use image", "image", image)

	var backoffLimit int32 = 0
	job := &batchv1.Job{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: ns,
			Labels:    t.LightningLabels,
		},
		Spec: batchv1.JobSpec{
			BackoffLimit: &backoffLimit,
			Template: corev1.PodTemplateSpec{
				ObjectMeta: metav1.ObjectMeta{
					Annotations: map[string]string{
						"prometheus.io/scrape": "true",
						"prometheus.io/port":   strconv.Itoa(lightningPort),
						"prometheus.io/path":   "/metrics",
					},
					Labels: t.LightningLabels,
				},
				Spec: corev1.PodSpec{
					HostNetwork:   true,
					DNSPolicy:     corev1.DNSClusterFirstWithHostNet,
					RestartPolicy: corev1.RestartPolicyNever,
					Tolerations: []corev1.Toleration{
						lightningToleration,
					},
					// sometimes after start lightning will meet issue about dns(do not know the detail reason):
					// dial tcp: lookup db-tidb on 100.64.0.10:53: no such host
					//
					// this just check that we can resolve the host before start lightning.

					// 2023/04/23(okjiang)
					// As the time elapsed, we can't reproduce this issue now. Let's remove this piece of code for now.
					// If we happen to encounter this issue again in the near future, we will investigate its root cause.
					// TODO: If we don't encounter it within two months, please remove these comments.
					Containers: []corev1.Container{
						{
							Name:            "tidb-lightning",
							Image:           image,
							ImagePullPolicy: corev1.PullIfNotPresent,
							Command: []string{
								"/tidb-lightning",
								"--server-mode",
								"--config",
								"/etc/tidb-lightning/tidb-lightning.toml",
								"--log-file",
								"-",
							},
							VolumeMounts: mounts,
							Env:          envs,
						},
					},
					Volumes:      volumes,
					NodeSelector: GetNodeLabels(t.Namer, t.Import),
				},
			},
		},
	}

	labels := utils.CopyLabelsForWorkloads(t.Import.Labels)
	job.Labels = collections.MergeMaps(labels, job.Labels)
	job.Spec.Template.Labels = collections.MergeMaps(labels, job.Spec.Template.Labels)
	return job, nil
}

// getClusterCertSecretName see also infra-api/controllers/imports/env.go:Env.GetClusterCertSecretName
func getClusterCertSecretName(tcName string) string {
	return fmt.Sprintf("%s-cluster-client-secret", tcName)
}

// getTiDBServerCertSecretName see also infra-api/controllers/imports/env.go:Env.GetTiDBServerCertSecretName
func getTiDBServerCertSecretName(tcName string) string {
	return fmt.Sprintf("%s-tidb-server-secret", tcName)
}

func (t *Job) mountPVCVolume(volumes []corev1.Volume, mounts []corev1.VolumeMount, isEmpty bool) ([]corev1.Volume, []corev1.VolumeMount) {
	volume := corev1.Volume{Name: t.Import.Name}
	if isEmpty {
		volume.VolumeSource = corev1.VolumeSource{
			EmptyDir: &corev1.EmptyDirVolumeSource{},
		}
	} else {
		volume.VolumeSource = corev1.VolumeSource{
			PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{
				ClaimName: t.Import.Name,
				ReadOnly:  false,
			},
		}
	}

	mount := corev1.VolumeMount{
		Name:      t.Import.Name,
		MountPath: "/etc/endpoint",
	}

	volumes = append(volumes, volume)
	mounts = append(mounts, mount)
	return volumes, mounts
}

func GetNodeLabels(nm *namer.Namer, i *cloudv1.Import) map[string]string {
	tags := cloudv1.TagsFromLabels(i.Labels)
	tags = common.ConcatenateTags(tags, map[string]string{
		common.UsedBy:    common.CustomerTiDB,
		common.Component: task.LightningComponent,
	})
	labels := nm.WithComponent(task.LightningComponent).
		WithNodeGroupSetOwner(util.FnvHash(i.ClusterScopeName())).
		Labels()
	tags = common.ConcatenateTags(tags, labels)
	return tags
}