package common

import (
	"testing"

	"github.com/samber/lo"
	"github.com/stretchr/testify/assert"

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

func TestGetCpuGenerationFilter(t *testing.T) {
	cfg := globalConf.Load()
	cfg.ServerConfig.FeatureGates = map[string]featgatesvc.FeatureGate{}
	featGateSvc := featgatesvc.NewFeatureGate(cfg.ServerConfig.FeatureGates, nil)

	t.Run("NewClusterReturnsNil", func(t *testing.T) {
		cluster := &entity.Cluster{}
		requestComp := &globaldedicatedv1.Cluster_ClusterComponent{}
		result := GetCpuGenerationFilter(featGateSvc, cluster, requestComp, true)
		assert.Nil(t, result)
	})

	t.Run("NonHomogeneousComponentIgnored", func(t *testing.T) {
		cluster := &entity.Cluster{
			Cluster: &globaldedicatedv1.Cluster{
				ClusterComponents: []*globaldedicatedv1.Cluster_ClusterComponent{
					{
						ComponentType: globaldedicatedv1.ComponentType_TIPROXY,
					},
				},
			},
		}
		requestComp := &globaldedicatedv1.Cluster_ClusterComponent{
			ComponentType: globaldedicatedv1.ComponentType_TIPROXY,
		}
		result := GetCpuGenerationFilter(featGateSvc, cluster, requestComp, false)
		assert.Nil(t, result)
	})

	t.Run("AwsArmClusterUsesExistingGen", func(t *testing.T) {
		cluster := &entity.Cluster{
			Cluster: &globaldedicatedv1.Cluster{
				ClusterComponents: []*globaldedicatedv1.Cluster_ClusterComponent{
					{
						ComponentType:         globaldedicatedv1.ComponentType_TIKV,
						InstanceCpuArch:       globaldedicatedv1.InstanceMu_ARM64,
						InstanceCpuGeneration: globaldedicatedv1.InstanceMu_GRAVITON2,
					},
				},
				CloudProviderId: globalcommonv1.CloudProvider_AWS,
			},
		}
		requestComp := &globaldedicatedv1.Cluster_ClusterComponent{
			ComponentType: globaldedicatedv1.ComponentType_TIKV,
		}
		result := GetCpuGenerationFilter(featGateSvc, cluster, requestComp, false)
		assert.NotNil(t, result)
		assert.Equal(t, globaldedicatedv1.InstanceMu_GRAVITON2, lo.FromPtr(result))
	})

	t.Run("AwsAmdClusterStaysUnspecified", func(t *testing.T) {
		cluster := &entity.Cluster{
			Cluster: &globaldedicatedv1.Cluster{
				ClusterComponents: []*globaldedicatedv1.Cluster_ClusterComponent{
					{
						ComponentType:         globaldedicatedv1.ComponentType_TIKV,
						InstanceCpuArch:       globaldedicatedv1.InstanceMu_AMD64,
						InstanceCpuGeneration: globaldedicatedv1.InstanceMu_CPU_GENERATION_UNSPECIFIED,
					},
				},
				CloudProviderId: globalcommonv1.CloudProvider_AWS,
			},
		}
		requestComp := &globaldedicatedv1.Cluster_ClusterComponent{
			ComponentType: globaldedicatedv1.ComponentType_TIKV,
		}
		result := GetCpuGenerationFilter(featGateSvc, cluster, requestComp, false)
		assert.Equal(t, globaldedicatedv1.InstanceMu_CPU_GENERATION_UNSPECIFIED, lo.FromPtr(result))
	})

	t.Run("AzureArmClusterStaysUnspecified", func(t *testing.T) {
		cluster := &entity.Cluster{
			Cluster: &globaldedicatedv1.Cluster{
				ClusterComponents: []*globaldedicatedv1.Cluster_ClusterComponent{
					{
						ComponentType:         globaldedicatedv1.ComponentType_TIKV,
						InstanceCpuArch:       globaldedicatedv1.InstanceMu_ARM64,
						InstanceCpuGeneration: globaldedicatedv1.InstanceMu_CPU_GENERATION_UNSPECIFIED,
					},
				},
				CloudProviderId: globalcommonv1.CloudProvider_AZURE,
			},
		}
		requestComp := &globaldedicatedv1.Cluster_ClusterComponent{
			ComponentType: globaldedicatedv1.ComponentType_TIKV,
		}
		result := GetCpuGenerationFilter(featGateSvc, cluster, requestComp, false)
		assert.Equal(t, globaldedicatedv1.InstanceMu_CPU_GENERATION_UNSPECIFIED, lo.FromPtr(result))
	})

	t.Run("GcpN4ClusterUsesExistingGen", func(t *testing.T) {
		cluster := &entity.Cluster{
			Cluster: &globaldedicatedv1.Cluster{
				ClusterComponents: []*globaldedicatedv1.Cluster_ClusterComponent{
					{
						ComponentType:         globaldedicatedv1.ComponentType_TIKV,
						InstanceCpuArch:       globaldedicatedv1.InstanceMu_AMD64,
						InstanceCpuGeneration: globaldedicatedv1.InstanceMu_GENERATION4_INTEL_XEON,
					},
				},
				CloudProviderId: globalcommonv1.CloudProvider_GCP,
			},
		}
		requestComp := &globaldedicatedv1.Cluster_ClusterComponent{
			ComponentType: globaldedicatedv1.ComponentType_TIKV,
		}
		result := GetCpuGenerationFilter(featGateSvc, cluster, requestComp, false)
		assert.NotNil(t, result)
		assert.Equal(t, globaldedicatedv1.InstanceMu_GENERATION4_INTEL_XEON, lo.FromPtr(result))
	})

	t.Run("GcpC4aClusterUsesExistingGen", func(t *testing.T) {
		cluster := &entity.Cluster{
			Cluster: &globaldedicatedv1.Cluster{
				ClusterComponents: []*globaldedicatedv1.Cluster_ClusterComponent{
					{
						ComponentType:         globaldedicatedv1.ComponentType_TIKV,
						InstanceCpuArch:       globaldedicatedv1.InstanceMu_ARM64,
						InstanceCpuGeneration: globaldedicatedv1.InstanceMu_GENERATION4_AXION,
					},
				},
				CloudProviderId: globalcommonv1.CloudProvider_GCP,
			},
		}
		requestComp := &globaldedicatedv1.Cluster_ClusterComponent{
			ComponentType: globaldedicatedv1.ComponentType_TIKV,
		}
		result := GetCpuGenerationFilter(featGateSvc, cluster, requestComp, false)
		assert.NotNil(t, result)
		assert.Equal(t, globaldedicatedv1.InstanceMu_GENERATION4_AXION, lo.FromPtr(result))
	})

	t.Run("GcpN2ClusterKeepsUnspecified", func(t *testing.T) {
		cluster := &entity.Cluster{
			Cluster: &globaldedicatedv1.Cluster{
				ClusterComponents: []*globaldedicatedv1.Cluster_ClusterComponent{
					{
						ComponentType:         globaldedicatedv1.ComponentType_TIKV,
						InstanceCpuArch:       globaldedicatedv1.InstanceMu_AMD64,
						InstanceCpuGeneration: globaldedicatedv1.InstanceMu_CPU_GENERATION_UNSPECIFIED,
					},
				},
				CloudProviderId: globalcommonv1.CloudProvider_GCP,
			},
		}
		requestComp := &globaldedicatedv1.Cluster_ClusterComponent{
			ComponentType: globaldedicatedv1.ComponentType_TIKV,
		}
		result := GetCpuGenerationFilter(featGateSvc, cluster, requestComp, false)
		assert.NotNil(t, result)
		assert.Equal(t, globaldedicatedv1.InstanceMu_CPU_GENERATION_UNSPECIFIED, lo.FromPtr(result))
	})
}