package common

import (
	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/gitops"
)

func getInstanceArch(provider gitops.Provider, instanceType string) string {
	if instanceType == "" || provider == nil {
		return ""
	}
	it, err := provider.FindInstanceType(instanceType)
	if err != nil {
		return ""
	}
	return it.Arch
}

func (t *TiDBClusterV2) suspendCompute() bool {
	suspendCompute := false
	if t.Cluster.Spec.Cluster.SuspendAction != nil {
		suspendCompute = t.Cluster.Spec.Cluster.SuspendAction.SuspendCompute
	}
	return suspendCompute
}

func (t *TiDBClusterV2) buildNodeSelectorAndTolerations(comp cloudv1.ComponentType, instanceType string) (ns map[string]string, to []corev1.Toleration) {
	ns = t.Namer.WithComponent(string(comp)).Labels()
	to = utils.Toleration(comp)

	if t.featureGates.Enabled(feature.UseClusterAutoscaler) {
		// Recreate the node selectors to remove `tidb-component`, `tidb-namespace`, and `tidbcloud-cluster`
		ns = map[string]string{}
		to = append(to, utils.TolerationWithComponentShared()...)
	}

	ns[consts.InstanceTypeLabelKey] = instanceType

	if t.featureGates.Enabled(feature.AddNodePoolDataPlaneLabel) {
		ns[consts.LabelKeyNodeGroupType] = consts.LabelValNodeGroupTypeDataPlane
	}
	if t.featureGates.Enabled(feature.AddK8sVersionToNodeSelector) {
		ns[consts.K8sVersionLabelKey] = t.k8sVersion
	}
	if t.featureGates.Enabled(feature.Add3LayerNodeSelectorAndTolerations) {
		ns[consts.LabelKeyFor3LayerArch] = consts.LabelValTrue
		to = append(to, *utils.TolerationWith3LayerArch())
	}

	// Add ARM64 toleration for GKE ARM instances (e.g. c4a-*)
	if t.Config.Provider == consts.ProviderTypeGKE && getInstanceArch(t.Provider, instanceType) == consts.ARM64ArchName {
		to = append(to, *utils.TolerationWithArch(consts.ARM64ArchName))
	}

	return ns, to
}

func (t *TiDBClusterV2) buildTopologySpreadConstraints() []corev1.TopologySpreadConstraint {
	var constraints []corev1.TopologySpreadConstraint
	if t.featureGates.Enabled(feature.UseClusterAutoscaler) {
		constraint := corev1.TopologySpreadConstraint{
			MaxSkew:            1,
			TopologyKey:        consts.ZoneLabelKey,
			MinDomains:         ptr.To(int32(len(t.Zones()))),
			NodeAffinityPolicy: ptr.To(corev1.NodeInclusionPolicyIgnore),
			WhenUnsatisfiable:  corev1.DoNotSchedule,
		}
		constraints = append(constraints, constraint)
	}
	return constraints
}

// buildStoreContainer builds the tikv or tiflash container.
func (t *TiDBClusterV2) buildStoreContainer(containerName string, res ComponentResources) corev1.Container {
	container := corev1.Container{
		Name: containerName,
		Resources: corev1.ResourceRequirements{
			Requests: corev1.ResourceList{
				corev1.ResourceCPU:    *res.Requests.CPU,
				corev1.ResourceMemory: *res.Requests.Memory,
			},
			Limits: corev1.ResourceList{
				corev1.ResourceMemory: *res.Limits.Memory,
			},
		},
		Env: []corev1.EnvVar{
			{
				Name:  "TZ",
				Value: "UTC",
			},
		},
		SecurityContext: &corev1.SecurityContext{
			Privileged: ptr.To(false),
		},
	}

	if t.featureGates.Enabled(feature.AddKmsCredentialsEnv) {
		container.Env = append(container.Env,
			corev1.EnvVar{
				Name: consts.AWS_ACCESS_KEY_ENV_VAR,
				ValueFrom: &corev1.EnvVarSource{
					SecretKeyRef: &corev1.SecretKeySelector{
						LocalObjectReference: corev1.LocalObjectReference{Name: gcpKMSSecretName},
						Key:                  consts.AWS_ACCESS_KEY_ENV_VAR,
					},
				},
			},
			corev1.EnvVar{
				Name: consts.AWS_SECRET_KEY_ENV_VAR,
				ValueFrom: &corev1.EnvVarSource{
					SecretKeyRef: &corev1.SecretKeySelector{
						LocalObjectReference: corev1.LocalObjectReference{Name: gcpKMSSecretName},
						Key:                  consts.AWS_SECRET_KEY_ENV_VAR,
					},
				},
			},
		)
	}

	return container
}