package imports

import (
	"encoding/json"
	"fmt"
	"os"
	"strconv"

	"github.com/pingcap/errors"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/labels"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-api/pkg/k8s"
	"github.com/tidbcloud/infra-api/pkg/pingcap/label"
	"github.com/tidbcloud/infra-provider/util/config"
	batchv1 "k8s.io/api/batch/v1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/intstr"
	ctrl "sigs.k8s.io/controller-runtime"
)

type TableProgress struct {
	TotalWritten int64      `json:"w"`
	TotalSize    int64      `json:"z"`
	Status       taskStatus `json:"s"`
	Message      string     `json:"m,omitempty"`
}

type TaskProgress struct {
	Tables        map[string]*TableProgress `json:"t"`
	Status        taskStatus                `json:"s"`
	Message       string                    `json:"m,omitempty"`
	SizeWithIndex int64                     `json:"si"`
}

type TasksResponse struct {
	ID    int64  `json:"id"`
	Error string `json:"error"`
}

type LightningClient struct {
	k8s.HttpLink
}

func newLightningClient(link k8s.HttpLink) LightningClient {
	return LightningClient{HttpLink: link}
}

func (g *LightningClient) CreateTask(cfg string) (int64, error) {
	body, err := g.Post("/tasks", cfg)
	if err != nil {
		return 0, err
	}
	task := TasksResponse{}
	err = json.Unmarshal(body, &task)
	if err != nil {
		return 0, err
	}
	if task.ID == 0 {
		return 0, errors.New(task.Error)
	}
	return task.ID, nil
}

func (g *LightningClient) GetTaskProgress() (*TaskProgress, error) {
	data, err := g.Get("/progress/task")
	if err != nil {
		return nil, errors.AddStack(err)
	}

	task := new(TaskProgress)
	err = json.Unmarshal(data, task)
	if err != nil {
		return nil, errors.Annotatef(err, "failed to unmarshal response %s", data)
	}
	return task, nil
}

func newLightning(ci *ImportService, name string,
	cfg []byte, importCm *corev1.ConfigMap) (*batchv1.Job, *corev1.Secret) {
	ns := ci.env.ImportJobNamespace()
	secret := newLightningSecret(name, ns, cfg)
	volumes := []corev1.Volume{
		{
			Name: "config",
			VolumeSource: corev1.VolumeSource{
				Secret: &corev1.SecretVolumeSource{
					SecretName: name,
					Items: []corev1.KeyToPath{
						{Key: lightningConfigFileName, Path: "tidb-lightning.toml"},
					},
				},
			},
		},
	}
	mounts := []corev1.VolumeMount{
		{
			Name:      "config",
			MountPath: "/etc/tidb-lightning",
		},
	}
	secretName := ci.env.GetClusterCertSecretName()
	if secretName != "" {
		volumes = append(volumes, corev1.Volume{
			Name: "lightning-tls",
			VolumeSource: corev1.VolumeSource{
				Secret: &corev1.SecretVolumeSource{
					SecretName: secretName,
				},
			},
		})
		mounts = append(mounts, corev1.VolumeMount{
			Name:      "lightning-tls",
			MountPath: "/var/lib/lightning-tls",
		})
	}

	var backoffLimit int32 = 0
	if enableCheckpoint(importCm) {
		backoffLimit = 3
	}
	job := &batchv1.Job{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: ns,
			Labels:    lightningLabels(name),
		},
		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: lightningLabels(name),
				},
				Spec: corev1.PodSpec{
					DNSPolicy:     corev1.DNSClusterFirstWithHostNet,
					RestartPolicy: corev1.RestartPolicyNever,
					Containers: []corev1.Container{
						{
							Name:            "tidb-lightning",
							ImagePullPolicy: corev1.PullIfNotPresent,
							Command: []string{
								"/tidb-lightning",
								"--server-mode",
								"--config",
								"/etc/tidb-lightning/tidb-lightning.toml",
								"--log-file",
								"-",
							},
							VolumeMounts: mounts,
						},
					},
					Volumes: volumes,
				},
			},
		},
	}

	return job, secret
}

func newLightningService(name string, ns string) *corev1.Service {
	selector := map[string]string{
		label.NameLabelKey: name,
	}

	svc := &corev1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: ns,
			Labels:    lightningLabels(name),
		},
		Spec: corev1.ServiceSpec{
			Selector: selector,
			Ports: []corev1.ServicePort{
				{
					Protocol:   corev1.ProtocolTCP,
					Port:       lightningPort,
					TargetPort: intstr.FromInt32(lightningPort),
				},
			},
		},
	}

	return svc
}

func newLightningSecret(name string, ns string, cfg []byte) *corev1.Secret {
	cm := &corev1.Secret{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: ns,
			Labels:    lightningLabels(name),
		},
		Data: map[string][]byte{
			lightningConfigFileName: cfg,
		},
	}

	return cm
}

func preCreateJobHook(e *ServerlessTierEnv, job *batchv1.Job, importCm *corev1.ConfigMap) error {
	if job.Spec.Template.Labels == nil {
		job.Spec.Template.Labels = map[string]string{}
	}

	job.Spec.Template.Labels[types.MTLSLabelKey] = fmt.Sprintf("%v", e.managerCfg.ImportEnableTLS)

	for i := 0; i < len(job.Spec.Template.Spec.Containers); i++ {
		if job.Spec.Template.Spec.Containers[i].Name == "tidb-lightning" {
			job.Spec.Template.Spec.Containers[i].Image = e.managerCfg.ServerlessLightningImage
			// set resource limit
			resources, err := getLightningResources(e, importCm)
			if err != nil {
				return err
			}
			job.Spec.Template.Spec.Containers[i].Resources = *resources
		}
	}

	applyImportScheduling(e, &job.Spec.Template.Spec, importCm)
	job.Spec.Template.Spec.ServiceAccountName = e.managerCfg.ImportSAName
	job.Spec.Template.Spec.HostNetwork = false
	return nil
}

func applyImportScheduling(
	e *ServerlessTierEnv,
	podSpec *corev1.PodSpec,
	importCm *corev1.ConfigMap,
) {
	defaultValue := "import-pool"
	selectedNodeSelectorKey := nodeSelectorKey
	selectedTolerationsKey := tolerationsKey

	switch {
	case e.importObj.Annotations[importSuperVIPAnno] == "true" ||
		e.clusterObj.Status.TidbPool == superVipTidbPool ||
		e.clusterObj.Status.TidbPool == canarySuperVipTidbPool:
		defaultValue = "dedicated"
		selectedNodeSelectorKey = superVipNodeSelectorKey
		selectedTolerationsKey = superVipTolerationsKey
	case e.clusterObj.Status.TidbPool == vipTidbPool ||
		e.clusterObj.Status.TidbPool == canaryVipTidbPool:
		selectedNodeSelectorKey = vipNodeSelectorKey
		selectedTolerationsKey = vipTolerationsKey
	}

	nodeSelector, nodeSelectorConfigured := importCm.Data[selectedNodeSelectorKey]
	nodeSelectorConfigured = nodeSelectorConfigured && nodeSelector != ""
	toleration, tolerationConfigured := importCm.Data[selectedTolerationsKey]
	tolerationConfigured = tolerationConfigured && toleration != ""

	// Preserve the existing scheduling behavior only when neither setting is
	// configured. Once one is configured, selector and toleration are independent.
	if !nodeSelectorConfigured && !tolerationConfigured {
		nodeSelector = defaultValue
		toleration = defaultValue
		nodeSelectorConfigured = true
		tolerationConfigured = true
	}

	if nodeSelectorConfigured {
		podSpec.NodeSelector = map[string]string{
			labels.ServerlessComponentNodeLabelKey: nodeSelector,
		}
		// The dedicated-cpu-num label only exists on nodes of the "dedicated"
		// pool, so it is attached only when scheduling onto that pool.
		if nodeSelector == "dedicated" {
			dedicatedNodeCpuNum := importCm.Data[labels.ServerlessDedicatedNodeCpuNumKey]
			if dedicatedNodeCpuNum == "" {
				dedicatedNodeCpuNum = "64"
			}
			podSpec.NodeSelector[labels.ServerlessDedicatedNodeCpuNumKey] = dedicatedNodeCpuNum
		}
	}
	if tolerationConfigured {
		podSpec.Tolerations = []corev1.Toleration{{
			Key:      labels.ServerlessComponentNodeLabelKey,
			Operator: corev1.TolerationOpEqual,
			Effect:   corev1.TaintEffectNoSchedule,
			Value:    toleration,
		}}
	}
}

func getLightningResources(
	e *ServerlessTierEnv,
	importCm *corev1.ConfigMap,
) (*corev1.ResourceRequirements, error) {
	cpuRequestKey := cpuResourceRequestKey
	cpuLimitKey := cpuResourceLimitKey
	memoryRequestKey := memoryResourceRequestKey
	memoryLimitKey := memoryResourceLimitKey
	switch e.clusterObj.Status.TidbPool {
	case vipTidbPool, canaryVipTidbPool:
		cpuRequestKey = vipCpuResourceRequestKey
		cpuLimitKey = vipCpuResourceLimitKey
		memoryRequestKey = vipMemoryResourceRequestKey
		memoryLimitKey = vipMemoryResourceLimitKey
	case superVipTidbPool, canarySuperVipTidbPool:
		cpuRequestKey = superVipCpuResourceRequestKey
		cpuLimitKey = superVipCpuResourceLimitKey
		memoryRequestKey = superVipMemoryResourceRequestKey
		memoryLimitKey = superVipMemoryResourceLimitKey
	}

	if e.importObj.Annotations[importSuperVIPAnno] == "true" {
		cpuRequestKey = superVipCpuResourceRequestKey
		cpuLimitKey = superVipCpuResourceLimitKey
		memoryRequestKey = superVipMemoryResourceRequestKey
		memoryLimitKey = superVipMemoryResourceLimitKey
	}

	// default resource
	cpuResourceRequest := defaultCpuResourceRequest
	cpuResourceLimit := defaultCpuResourceLimit
	memoryResourceRequest := defaultMemoryResourceRequest
	memoryResourceLimit := defaultMemoryResourceLimit

	if val, ok := importCm.Data[cpuRequestKey]; ok {
		cpuResourceRequest = val
	}
	if val, ok := importCm.Data[cpuLimitKey]; ok {
		cpuResourceLimit = val
	}
	if val, ok := importCm.Data[memoryRequestKey]; ok {
		memoryResourceRequest = val
	}
	if val, ok := importCm.Data[memoryLimitKey]; ok {
		memoryResourceLimit = val
	}

	parseQuantity := func(key, value string) (resource.Quantity, error) {
		quantity, err := resource.ParseQuantity(value)
		if err != nil {
			return resource.Quantity{}, fmt.Errorf("invalid resource quantity for %q: %q: %w", key, value, err)
		}
		return quantity, nil
	}

	cpuLimit, err := parseQuantity(cpuLimitKey, cpuResourceLimit)
	if err != nil {
		return nil, err
	}
	memoryLimit, err := parseQuantity(memoryLimitKey, memoryResourceLimit)
	if err != nil {
		return nil, err
	}
	cpuRequest, err := parseQuantity(cpuRequestKey, cpuResourceRequest)
	if err != nil {
		return nil, err
	}
	memoryRequest, err := parseQuantity(memoryRequestKey, memoryResourceRequest)
	if err != nil {
		return nil, err
	}

	resources := &corev1.ResourceRequirements{
		Limits: corev1.ResourceList{
			corev1.ResourceCPU:    cpuLimit,
			corev1.ResourceMemory: memoryLimit,
		},
		Requests: corev1.ResourceList{
			corev1.ResourceCPU:    cpuRequest,
			corev1.ResourceMemory: memoryRequest,
		},
	}

	return resources, nil
}

func preCreateSecretHook(
	secret *corev1.Secret,
	e *ServerlessTierEnv,
	importCm *corev1.ConfigMap,
) error {
	tidbUserName := e.clusterObj.Status.KeyspaceToken + "." + e.importObj.Spec.Source.UserName
	tidbHost := "tidb-gateway-inner." + os.Getenv("NAMESPACE") + ".svc.cluster.local"
	pdAddr := "serverless-cluster-pd.tidb-serverless.svc.cluster.local:2379"

	cfgBytes := secret.Data[lightningConfigFileName]
	impConfigTOML := new(cloudv1.TomlConfig)
	if err := impConfigTOML.UnmarshalTOML(cfgBytes); err != nil {
		ctrl.Log.Error(err, "unmarshal lightning config",
			"KeyspaceToken", e.clusterObj.Status.KeyspaceToken)
		return err
	}
	cfg := config.New(impConfigTOML.MP)

	cfg.Set("tidb.user", tidbUserName)
	cfg.Set("tidb.pd-addr", pdAddr)
	cfg.Set("tidb.host", tidbHost)
	cfg.Set("tidb.port", tidbPort)
	cfg.Set("tikv-importer.addr", e.buildUrl("tikv-api-service.tidb-serverless.svc.cluster.local:19500"))
	// If the manager doesn't enable TLS, we will not mount cert secret to import job.
	if e.GetClusterCertSecretName() == "" {
		cfg.Del("security")
	} else {
		// Otherwise, use mount cert secret to import job.
		cfg.Set("security.ca-path", "/var/lib/lightning-tls/ca.crt")
		cfg.Set("security.cert-path", "/var/lib/lightning-tls/tls.crt")
		cfg.Set("security.key-path", "/var/lib/lightning-tls/tls.key")
		cfg.Set("tidb.tls", "")
	}
	id, err := e.getImportTaskID()
	// TODO The v1beta1 API import id is string type, we need to creating a mapping later.
	if err == nil {
		cfg.Set("id", id)
	}

	// config remote backend
	cfg.Set("tikv-importer.backend", "remote")
	cfg.Set("tikv-importer.keyspace-name", e.clusterObj.Status.KeyspaceToken)
	cfg.Set("tikv-importer.duplicate-resolution", "remove")
	cfg.Del("tikv-importer.on-duplicate")

	ctrl.Log.Info("use remote backend", "keyspaceToken", e.clusterObj.Status.KeyspaceToken)

	err = applyLightningConfigFromCm(cfg, extraLightningConfigKey, e.clusterObj.Status.KeyspaceToken, importCm)
	if err != nil {
		return err
	}

	switch e.clusterObj.Status.TidbPool {
	case vipTidbPool, canaryVipTidbPool:
		err := applyLightningConfigFromCm(cfg, extraVipLightningConfigKey, e.clusterObj.Status.KeyspaceToken, importCm)
		if err != nil {
			return err
		}
	case superVipTidbPool, canarySuperVipTidbPool:
		err := applyLightningConfigFromCm(cfg, extraSuperVipLightningConfigKey,
			e.clusterObj.Status.KeyspaceToken, importCm)
		if err != nil {
			return err
		}
	}

	if e.importObj.Annotations[importSuperVIPAnno] == "true" {
		err := applyLightningConfigFromCm(cfg, extraSuperVipLightningConfigKey,
			e.clusterObj.Status.KeyspaceToken, importCm)
		if err != nil {
			return err
		}
	}

	cfgBytes, err = cfg.MarshalTOML()
	if err != nil {
		return err
	}
	secret.Data[lightningConfigFileName] = cfgBytes
	return nil
}

func applyLightningConfigFromCm(cfg *config.GenericConfig,
	cfgName, keyspaceName string, importCm *corev1.ConfigMap) error {
	// get extra lightning config from `import-cm`
	extraCfgStr, ok := importCm.Data[cfgName]
	if ok {
		ctrl.Log.Info("add extra lightning config", "keyspace", keyspaceName, "config", cfgName)
		extraImpCfgTOML := new(cloudv1.TomlConfig)
		if err := extraImpCfgTOML.UnmarshalTOML([]byte(extraCfgStr)); err != nil {
			ctrl.Log.Error(err, "unmarshal extra lightning config",
				"keyspace", keyspaceName, "config", cfgName)
			return err
		}
		// merge extra lightning config into lightning config
		extraCfg := config.New(extraImpCfgTOML.MP)
		for k, v := range extraCfg.FlatMP() {
			cfg.Set(k, v)
		}
	}
	return nil
}

func lightningLabels(name string) map[string]string {
	labels := map[string]string{
		label.ComponentLabelKey: lightningName,
		label.NameLabelKey:      name,
	}
	return labels
}