package tidbpool

import (
	"context"
	"os"
	"strconv"
	"strings"

	"github.com/pingcap/errors"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/labels"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/yaml"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

const (
	tidbServerPodTemplate = "tidb-server-pod-template"

	tidbMinMemoryEnvName = "MIN_MEMORY_LIMIT"
	tidbMaxMemoryEnvName = "MAX_MEMORY_LIMIT"
	tidbPodContainer     = "tidb"
	managerNsEnvName     = "MANAGER_NS"
	tidbWorkerNsEnvName  = "TIDB_WORKER_NS"

	defaultMemoryLimitRatio = 0.8
)

func (tp *tidbPool) decodeTidbServerPod(
	ctx context.Context, serverlessPoolConfigCm *corev1.ConfigMap,
) (*corev1.Pod, []string, error) {
	tidbServerPodTemplateCm := &corev1.ConfigMap{}
	if err := tp.client.Get(ctx,
		client.ObjectKey{
			Name:      types.TidbServerPodTemplateCm,
			Namespace: tp.poolNamespace,
		}, tidbServerPodTemplateCm); err != nil {
		return nil, nil, err
	}
	podTemplate := tidbServerPodTemplateCm.Data[tidbServerPodTemplate]
	if podTemplate == "" {
		return nil, nil, errors.New("tidbServerPodTemplate is empty")
	}
	tidbPod := &corev1.Pod{}
	if err := runtime.DecodeInto(tp.codecs.UniversalDecoder(), []byte(podTemplate), tidbPod); err != nil {
		return nil, nil, err
	}
	tidbPod.Namespace = tp.poolNamespace
	if tidbPod.Labels == nil {
		tidbPod.Labels = map[string]string{}
	}
	tidbPod.Labels[labels.NameLabelKey] = labels.TidbPodNameLabelValue
	tidbPod.Labels[labels.ComponentLabelKey] = labels.TidbPodComponentLabelValue
	useVPA := serverlessPoolConfigCm.Data[types.TidbUseVpaKey] == "true"
	if useVPA {
		tidbPod.Labels[types.UseVPALabelKey] = "true"
	}

	nodeSelector := "tidb"
	if serverlessPoolConfigCm.Data[types.TidbNodeSelector] != "" {
		nodeSelector = serverlessPoolConfigCm.Data[types.TidbNodeSelector]
	}
	tidbPod.Spec.NodeSelector =
		utils.UpdateMap(tidbPod.Spec.NodeSelector, labels.ServerlessComponentNodeLabelKey, nodeSelector)
	nodeGroups := updateNodeGroupSelector(tidbPod, serverlessPoolConfigCm.Data[types.NodeGroupSelector])
	tidbPod.Spec.Tolerations = append(tidbPod.Spec.Tolerations, corev1.Toleration{
		Key:      labels.ServerlessComponentNodeLabelKey,
		Operator: corev1.TolerationOpEqual,
		Effect:   corev1.TaintEffectNoSchedule,
		Value:    "tidb",
	})
	for i := range tidbPod.Spec.Containers {
		if tidbPod.Spec.Containers[i].Name == tidbPodContainer {
			tidbPod.Spec.Containers[i].Image = serverlessPoolConfigCm.Data[types.TidbImageKey]
			tidbPod.Spec.Containers[i].Env = append(tidbPod.Spec.Containers[i].Env, corev1.EnvVar{
				Name:  managerNsEnvName,
				Value: os.Getenv("NAMESPACE"),
			})
			tidbPod.Spec.Containers[i].Env = append(tidbPod.Spec.Containers[i].Env, corev1.EnvVar{
				Name:  tidbWorkerNsEnvName,
				Value: tp.cfg.TiDBWorkerNS,
			})
			if tidbResources := serverlessPoolConfigCm.Data[types.TidbResourcesKey]; tidbResources != "" {
				if err := yaml.Unmarshal([]byte(tidbResources), &tidbPod.Spec.Containers[i].Resources); err != nil {
					return nil, nil, err
				}
			}
			if useVPA {
				// set MIN_MEMORY_LIMIT and MAX_MEMORY_LIMIT to request memory and limit memory,
				// scaled by tidb-memory-limit-ratio to reserve headroom for Go GC and system overhead.
				var resources corev1.ResourceRequirements
				if err := yaml.Unmarshal([]byte(serverlessPoolConfigCm.Data[types.TidbResourcesKey]), &resources); err != nil {
					return nil, nil, err
				}
				ratio := defaultMemoryLimitRatio
				if v := serverlessPoolConfigCm.Data[types.TidbMemoryLimitRatioKey]; v != "" {
					if parsed, err := strconv.ParseFloat(v, 64); err == nil && parsed > 0 && parsed <= 1 {
						ratio = parsed
					}
				}
				requests, limits := resources.Requests[corev1.ResourceMemory], resources.Limits[corev1.ResourceMemory]
				minMemory := int64(float64(requests.Value()) * ratio)
				maxMemory := int64(float64(limits.Value()) * ratio)
				tidbPod.Spec.Containers[i].Env = append(tidbPod.Spec.Containers[i].Env, corev1.EnvVar{
					Name:  tidbMinMemoryEnvName,
					Value: strconv.FormatInt(minMemory, 10),
				}, corev1.EnvVar{
					Name:  tidbMaxMemoryEnvName,
					Value: strconv.FormatInt(maxMemory, 10),
				})
			}
		}
	}

	return tidbPod, nodeGroups, nil
}

func updateNodeGroupSelector(pod *corev1.Pod, nodeGroupSelector string) (selectors []string) {
	for _, selector := range strings.Split(nodeGroupSelector, ",") {
		if selector = strings.TrimSpace(selector); selector != "" {
			selectors = append(selectors, selector)
		}
	}
	if len(selectors) == 0 {
		return
	}

	if pod.Spec.Affinity == nil {
		pod.Spec.Affinity = &corev1.Affinity{}
	}
	if pod.Spec.Affinity.NodeAffinity == nil {
		pod.Spec.Affinity.NodeAffinity = &corev1.NodeAffinity{}
	}
	if pod.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution == nil {
		pod.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = &corev1.NodeSelector{}
	}
	req := pod.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution
	if len(req.NodeSelectorTerms) == 0 {
		req.NodeSelectorTerms = []corev1.NodeSelectorTerm{{}}
	}
	req.NodeSelectorTerms[0].MatchExpressions = append(req.NodeSelectorTerms[0].MatchExpressions, corev1.NodeSelectorRequirement{
		Key:      labels.ServerlessNodeGroupLabelKey,
		Operator: corev1.NodeSelectorOpIn,
		Values:   selectors,
	})
	return
}