package common

import (
	"github.com/samber/lo"

	globaldedicatedv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/dedicated/v1"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/entity"
	"github.com/tidbcloud/tidb-management-service/utilities/featuregate"
)

func GetCpuArch(
	featGateCli featuregate.Service,
	cluster *entity.Cluster,
	requestComp *globaldedicatedv1.Cluster_ClusterComponent,
	newCluster bool,
) *globaldedicatedv1.InstanceMu_CpuArch {
	// If the component is not a CPU homogeneous component, we skip the filter.
	_, isCpuHomogeneousComponent := entity.ComponentsCpuHomogeneous[requestComp.ComponentType]
	if !isCpuHomogeneousComponent {
		return nil
	}

	// Get featuregate cpu arch preference.
	var cpuArch *globaldedicatedv1.InstanceMu_CpuArch
	if featGateCli.CheckFeatureEnabled(
		featuregate.PreferArm,
		cluster.OrganizationId,
		cluster.ProjectId,
		cluster.GetCloudProviderId().String(),
		cluster.RegionId.String(),
	) {
		cpuArch = lo.ToPtr(globaldedicatedv1.InstanceMu_ARM64)
	} else {
		cpuArch = lo.ToPtr(globaldedicatedv1.InstanceMu_AMD64)
	}

	if newCluster {
		// If it's a new cluster, we can directly return the featuregate cpu arch.
		return cpuArch
	}

	// If it's an existing cluster, we need to override the cpu arch from existing homogeneous components.
	for _, comp := range cluster.GetClusterComponents() {
		_, isCpuHomogeneousComponent = entity.ComponentsCpuHomogeneous[comp.ComponentType]
		if isCpuHomogeneousComponent {
			cpuArch = lo.ToPtr(comp.GetInstanceCpuArch())
			break
		}
	}

	// If no existing homogeneous component found, we return the featuregate cpu arch.

	return cpuArch
}

// GetCpuGenerationFilter returns the CPU generation filter for the given cluster and request component.
// For existing AWS ARM64 clusters, we need to make sure the CPU generation of instances is consistent across homogeneous
// components when scaling.
// For existing GCP clusters, we also need to make sure the CPU generation of instances is consistent across homogeneous components.
func GetCpuGenerationFilter(
	featGateCli featuregate.Service,
	cluster *entity.Cluster,
	requestComp *globaldedicatedv1.Cluster_ClusterComponent,
	newCluster bool,
) *globaldedicatedv1.InstanceMu_CpuGeneration {
	// If the component is not a CPU homogeneous component, we skip the filter.
	_, isCpuHomogeneousComponent := entity.ComponentsCpuHomogeneous[requestComp.ComponentType]
	if !isCpuHomogeneousComponent {
		return nil
	}

	if featGateCli.CheckFeatureEnabled(
		featuregate.AllowHeterogeneousCpuGeneration,
		cluster.OrganizationId,
		cluster.ProjectId,
		cluster.GetCloudProviderId().String(),
		cluster.RegionId.String(),
	) {
		return nil
	}

	// We only check cpu generation for existing clusters.
	if newCluster {
		return nil
	}

	var cpuGeneration *globaldedicatedv1.InstanceMu_CpuGeneration
	for _, comp := range cluster.GetClusterComponents() {
		_, isCpuHomogeneousComponent = entity.ComponentsCpuHomogeneous[comp.ComponentType]
		if isCpuHomogeneousComponent {
			cpuGeneration = lo.ToPtr(comp.GetInstanceCpuGeneration())
			break
		}
	}

	// If no existing homogeneous component found, we skip the filter.

	return cpuGeneration
}