package operator

import (
	"fmt"
	"github.com/pingcap/tep/db/entity"
	"github.com/pingcap/tep/pkg/apis/apiserver"
	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	"github.com/pingcap/tidb-operator/pkg/label"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/utils/pointer"
)

func setComponent(component string, values ComponentValues, instance *apiserver.TiDBCluster) error {
	var scheduleNodeLabel string
	var componentSpec apiserver.ComponentSpec
	var pdConfig *v1alpha1.PDConfig
	var tidbConfig = TiDBDefaultConfig
	var tikvConfig = TiKVDefaultConfig

	switch component {
	case label.PDLabelVal:
		scheduleNodeLabel = TEPStorageNodeLabel
		componentSpec = instance.PD
		pdConfig = PDDefaultConfig
		if instance.ReplicationFactor != 0 {
			replica := uint64(instance.ReplicationFactor)
			pdConfig.Replication.MaxReplicas = &replica
		}
	case label.TiKVLabelVal:
		scheduleNodeLabel = TEPStorageNodeLabel
		componentSpec = instance.TiKV
		tikvConfig = TiKVDefaultConfig
	case label.TiDBLabelVal:
		componentSpec = instance.TiDB
		scheduleNodeLabel = TEPComputeNodeLabel
		tidbConfig = TiDBDefaultConfig
	default:
		return fmt.Errorf("unknown component type[%s]", component)
	}

	if componentSpec.GetReplicas() != 0 {
		values.SetReplicas(componentSpec.GetReplicas())
	}

	componentSpecCPULimit := componentSpec.GetResourceTemplate().CPU
	if componentSpecCPULimit != 0 {
		cpuLimit := resource.NewMilliQuantity(int64(componentSpecCPULimit)*1000, resource.DecimalSI)
		values.GetResources().Limits[corev1.ResourceCPU] = *cpuLimit
		values.GetResources().Requests[corev1.ResourceCPU] = *cpuLimit
		if component == label.TiDBLabelVal {
			maxProcs := uint(componentSpecCPULimit)
			tidbConfig.Performance.MaxProcs = &maxProcs
		}
		if component == label.TiKVLabelVal && componentSpecCPULimit > 8 {
			coprocessorConcurrency := int64(float64(componentSpecCPULimit) * 0.8)
			coprocessorConfig := &v1alpha1.TiKVCoprocessorReadPoolConfig{
				HighConcurrency:   pointer.Int64Ptr(coprocessorConcurrency),
				NormalConcurrency: pointer.Int64Ptr(coprocessorConcurrency),
				LowConcurrency:    pointer.Int64Ptr(coprocessorConcurrency),
			}
			tikvConfig.ReadPool.Coprocessor = coprocessorConfig
		}
	}

	componentSpecMemoryLimit := componentSpec.GetResourceTemplate().Memory
	if componentSpecMemoryLimit != 0 {
		memoryLimit := resource.NewQuantity(int64(componentSpecMemoryLimit)*1024*1024*1024, resource.BinarySI)
		values.GetResources().Limits[corev1.ResourceMemory] = *memoryLimit
		values.GetResources().Requests[corev1.ResourceMemory] = *memoryLimit
		if component == label.TiKVLabelVal && componentSpecMemoryLimit > 4 {
			capacity := uint32(float64(componentSpecMemoryLimit) * 0.4)
			var t = true
			blockCacheConfig := v1alpha1.TiKVBlockCacheConfig{
				Shared:   &t,
				Capacity: fmt.Sprintf("%dGB", capacity),
			}
			tikvConfig.Storage.BlockCache = &blockCacheConfig
		}
	}

	affinity := corev1.Affinity{}
	nodeAffinity := corev1.NodeAffinity{}
	nodeSelector := &corev1.NodeSelector{
		NodeSelectorTerms: []corev1.NodeSelectorTerm{
			{
				MatchExpressions: []corev1.NodeSelectorRequirement{
					{
						Key:      TEPSchedulableLabel,
						Operator: corev1.NodeSelectorOpIn,
						Values:   []string{entity.MachineOnline.String()},
					},
					{
						Key:      scheduleNodeLabel,
						Operator: corev1.NodeSelectorOpExists,
					},
				},
			},
		},
	}

	azCount := len(instance.AZs)
	if azCount != 0 {
		zoneNodeSelectorRequirement := corev1.NodeSelectorRequirement{
			Key:      "zone",
			Operator: corev1.NodeSelectorOpIn,
			Values:   instance.AZs,
		}
		nodeSelector.NodeSelectorTerms[0].MatchExpressions = append(nodeSelector.NodeSelectorTerms[0].MatchExpressions, zoneNodeSelectorRequirement)

		if azCount > 1 && component == label.PDLabelVal {
			pdConfig.Replication.LocationLabels = v1alpha1.StringSlice{"zone", "host"}
		}
	}
	nodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = nodeSelector
	affinity.NodeAffinity = &nodeAffinity
	affinity.PodAntiAffinity = getPodAntiAffinity(component, instance.Name, azCount)
	values.SetAffinity(affinity)

	return nil
}

func getPodAntiAffinity(component string, instanceName string, azCount int) *corev1.PodAntiAffinity {
	podAntiAffinity := &corev1.PodAntiAffinity{
		PreferredDuringSchedulingIgnoredDuringExecution: []corev1.WeightedPodAffinityTerm{
			{
				Weight: 80,
				PodAffinityTerm: corev1.PodAffinityTerm{
					LabelSelector: &metav1.LabelSelector{
						MatchLabels: map[string]string{
							label.InstanceLabelKey:  instanceName,
							label.ComponentLabelKey: component,
						},
					},
					TopologyKey: "kubernetes.io/hostname",
				},
			},
		},
	}

	if azCount > 1 {
		zoneWeightedPodAffinityTerm := corev1.WeightedPodAffinityTerm{
			Weight: 20,
			PodAffinityTerm: corev1.PodAffinityTerm{
				LabelSelector: &metav1.LabelSelector{
					MatchLabels: map[string]string{
						label.InstanceLabelKey:  instanceName,
						label.ComponentLabelKey: component,
					},
				},
				TopologyKey: "zone",
			},
		}
		podAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution = append(podAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution, zoneWeightedPodAffinityTerm)
	}

	return podAntiAffinity
}