package builder

import (
	opv2 "github.com/pingcap/tidb-operator/api/v2/core/v1alpha1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/utils/ptr"

	"github.com/tidbcloud/eks-provider/controllers/cluster/task/common/feature"
	"github.com/tidbcloud/eks-provider/controllers/utils"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/namer"
)

func buildSchedulePolicies(zones []string) []opv2.SchedulePolicy {
	topologies := make([]opv2.ScheduleTopology, len(zones))
	for i, zone := range zones {
		topologies[i] = opv2.ScheduleTopology{Topology: map[string]string{consts.ZoneLabelKey: zone}}
	}
	return []opv2.SchedulePolicy{
		{
			Type: opv2.SchedulePolicyTypeEvenlySpread,
			EvenlySpread: &opv2.SchedulePolicyEvenlySpread{
				Topologies: topologies,
			},
		},
	}
}

func buildNodeSelector(gates feature.Gates, ns, cluster, component, instanceType, k8sVersion string) map[string]string {
	s := map[string]string{
		namer.ClusterSelectorKey:   cluster,
		namer.ComponentSelectorKey: component,
		namer.NamespaceSelectorKey: ns,
	}
	if gates.Enabled(feature.UseClusterAutoscaler) {
		// Recreate the node selectors to remove `tidb-component`, `tidb-namespace`, and `tidbcloud-cluster`
		s = map[string]string{}
	}
	s[consts.InstanceTypeLabelKey] = instanceType

	if gates.Enabled(feature.AddNodePoolDataPlaneLabel) {
		s[consts.LabelKeyNodeGroupType] = consts.LabelValNodeGroupTypeDataPlane
	}
	if gates.Enabled(feature.AddK8sVersionToNodeSelector) {
		s[consts.K8sVersionLabelKey] = k8sVersion
	}
	if gates.Enabled(feature.Add3LayerNodeSelectorAndTolerations) {
		s[consts.LabelKeyFor3LayerArch] = consts.LabelValTrue
	}

	return s
}

func buildTolerations(gates feature.Gates, component, arch, provider string) []corev1.Toleration {
	to := utils.Toleration(cloudv1.ComponentType(component))
	if gates.Enabled(feature.UseClusterAutoscaler) {
		to = append(to, utils.TolerationWithComponentShared()...)
	}

	if gates.Enabled(feature.Add3LayerNodeSelectorAndTolerations) {
		to = append(to, *utils.TolerationWith3LayerArch())
	}

	if provider == consts.ProviderTypeGKE && arch == consts.ARM64ArchName {
		to = append(to, *utils.TolerationWithArch(consts.ARM64ArchName))
	}

	return to
}

func buildTopologySpreadConstraints(gates feature.Gates, zones []string) []corev1.TopologySpreadConstraint {
	var constraints []corev1.TopologySpreadConstraint
	if gates.Enabled(feature.UseClusterAutoscaler) {
		constraint := corev1.TopologySpreadConstraint{
			MaxSkew:            1,
			TopologyKey:        consts.ZoneLabelKey,
			MinDomains:         ptr.To(int32(len(zones))),
			NodeAffinityPolicy: ptr.To(corev1.NodeInclusionPolicyIgnore),
			// It's a required field.
			WhenUnsatisfiable: corev1.DoNotSchedule,
		}
		constraints = append(constraints, constraint)
	}
	return constraints
}