package spu

import (
	"testing"

	"github.com/samber/lo"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/mock/gomock"

	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"
	spurepo "github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/repository/spu"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/repository/spu/mock_spu"
)

func TestGetTiKVSku(t *testing.T) {
	// Main test setup
	t.Run("test select tikv sku", func(t *testing.T) {
		ctx := t.Context()
		ctrl := gomock.NewController(t, gomock.WithOverridableExpectations())
		defer ctrl.Finish()

		spuRepo := mock_spu.NewMockRepoSpu(ctrl)
		skuRepo := mock_spu.NewMockSkuRepo(ctrl)

		svc := NewSkuService(spuRepo, skuRepo, nil, nil)

		parameter := &SelectSkusRequest{
			TenantId:  111,
			ProjectId: 222,
			Provider:  globalcommonv1.CloudProvider_AWS,
			Region:    globalcommonv1.Region_AWS_US_EAST_1,
			SpuName:   lo.ToPtr("8C16G"),
		}
		spu := &entity.Spu{
			Id:            1,
			Key:           "8C16G",
			ComponentType: globaldedicatedv1.ComponentType_TIKV,
			CpuCore:       8,
			MemorySizeGi:  16,
		}
		sku1 := &entity.Sku{
			Id:             1,
			Key:            "4C",
			CloudProvider:  globalcommonv1.CloudProvider_AWS,
			ComponentType:  globaldedicatedv1.ComponentType_TIKV,
			SpuId:          1,
			SpuName:        "",
			SpuKey:         "8C16G",
			SpuDisplayName: "",
			Priority:       1,
			InstanceMuId:   1,
			CpuArch:        globaldedicatedv1.InstanceMu_ARM64,
			InstanceType:   "c5.xlarge",
			CpuCore:        8,
			MemoryGi:       16,
		}
		sku2 := &entity.Sku{
			Id:             2,
			Key:            "8C",
			CloudProvider:  globalcommonv1.CloudProvider_AWS,
			ComponentType:  globaldedicatedv1.ComponentType_TIKV,
			SpuId:          1,
			SpuName:        "",
			SpuKey:         "8C16G",
			SpuDisplayName: "",
			Priority:       1,
			InstanceMuId:   2,
			CpuArch:        globaldedicatedv1.InstanceMu_ARM64,
			InstanceType:   "c5.xlarge",
			CpuCore:        8,
			MemoryGi:       16,
		}
		spuMappings := map[uint64][]*entity.SpuMapping{
			spu.Id: {
				{
					SpuId:    spu.Id,
					SkuKey:   sku2.Key,
					Provider: globalcommonv1.CloudProvider_AWS,
					Regions: map[globalcommonv1.Region_Name]bool{
						globalcommonv1.Region_AWS_US_WEST_2: true,
						globalcommonv1.Region_AWS_US_EAST_1: true,
					},
					TenantIDs: map[uint64]bool{
						222: true,
					},
				},
				{
					SpuId:    spu.Id,
					SkuKey:   sku1.Key,
					Provider: globalcommonv1.CloudProvider_AWS,
					Regions: map[globalcommonv1.Region_Name]bool{
						globalcommonv1.Region_AWS_US_WEST_2: true,
						globalcommonv1.Region_AWS_US_EAST_1: true,
					},
					TenantIDs: map[uint64]bool{
						111: true,
						222: true,
						333: true,
						444: true,
					},
				},
				{
					SpuId:    spu.Id,
					SkuKey:   sku2.Key,
					Provider: globalcommonv1.CloudProvider_AWS,
					Regions: map[globalcommonv1.Region_Name]bool{
						globalcommonv1.Region_AWS_US_WEST_2: true,
						globalcommonv1.Region_AWS_US_EAST_1: true,
						globalcommonv1.Region_AWS_AP_EAST_1: true,
					},
					TenantIDs: map[uint64]bool{
						222: true,
						333: true,
					},
				},
			},
		}

		t.Run("should succeed", func(t *testing.T) {
			t.Run("select sku1", func(t *testing.T) {
				parameter.TenantId = 111
				parameter.Region = globalcommonv1.Region_AWS_US_WEST_2
				parameter.CpuArch = lo.ToPtr(globaldedicatedv1.InstanceMu_ARM64)

				spuRepo.EXPECT().Query(gomock.Any(), &spurepo.QuerySpuFilter{
					Name:           "8C16G",
					ComponentTypes: []globaldedicatedv1.ComponentType{globaldedicatedv1.ComponentType_TIKV},
				}).Return([]*entity.Spu{spu}, nil).AnyTimes()
				skuRepo.EXPECT().Query(gomock.Any(), &spurepo.SkuFilter{
					Component: globaldedicatedv1.ComponentType_TIKV,
					SpuId:     lo.ToPtr(spu.Id),
				}).Return([]*entity.Sku{sku1, sku2}, nil).AnyTimes()
				skuRepo.EXPECT().QuerySpuMappings(gomock.Any(), []uint64{spu.Id}).Return(spuMappings, nil).AnyTimes()

				skus, err := svc.SelectSkusByComponentType(ctx, globaldedicatedv1.ComponentType_TIKV, parameter)
				require.NoError(t, err)
				sku := skus[0]
				assert.Equal(t, sku1.Id, sku.Id)
				assert.Equal(t, sku1.Key, sku.Key)
				assert.Equal(t, globaldedicatedv1.InstanceMu_ARM64, sku.GetCpuArch())
				assert.Equal(t, globaldedicatedv1.ComponentType_TIKV, sku.GetComponentType())
			})

			t.Run("select sku2", func(t *testing.T) {
				parameter.TenantId = 222
				parameter.Region = globalcommonv1.Region_AWS_US_WEST_2
				parameter.CpuArch = lo.ToPtr(globaldedicatedv1.InstanceMu_ARM64)

				spuRepo.EXPECT().Query(gomock.Any(), &spurepo.QuerySpuFilter{
					Name:           "8C16G",
					ComponentTypes: []globaldedicatedv1.ComponentType{globaldedicatedv1.ComponentType_TIKV},
				}).Return([]*entity.Spu{spu}, nil).AnyTimes()
				skuRepo.EXPECT().Query(gomock.Any(), &spurepo.SkuFilter{
					Component: globaldedicatedv1.ComponentType_TIKV,
					SpuId:     lo.ToPtr(spu.Id),
				}).Return([]*entity.Sku{sku1, sku2}, nil).AnyTimes()
				skuRepo.EXPECT().QuerySpuMappings(gomock.Any(), []uint64{spu.Id}).Return(spuMappings, nil).AnyTimes()

				tierUnits, err := svc.SelectSkusByComponentType(ctx, globaldedicatedv1.ComponentType_TIKV, parameter)
				require.NoError(t, err)
				tierUnit := tierUnits[0]
				assert.Equal(t, sku2.Id, tierUnit.Id)
				assert.Equal(t, sku2.Key, tierUnit.Key)
				assert.Equal(t, globaldedicatedv1.InstanceMu_ARM64, tierUnit.GetCpuArch())
				assert.Equal(t, globaldedicatedv1.ComponentType_TIKV, tierUnit.GetComponentType())
			})

			t.Run("select sku2 with other region", func(t *testing.T) {
				parameter.TenantId = 333
				parameter.Region = globalcommonv1.Region_AWS_AP_EAST_1
				parameter.CpuArch = lo.ToPtr(globaldedicatedv1.InstanceMu_ARM64)

				spuRepo.EXPECT().Query(gomock.Any(), &spurepo.QuerySpuFilter{
					Name:           "8C16G",
					ComponentTypes: []globaldedicatedv1.ComponentType{globaldedicatedv1.ComponentType_TIKV},
				}).Return([]*entity.Spu{spu}, nil).AnyTimes()
				skuRepo.EXPECT().Query(gomock.Any(), &spurepo.SkuFilter{
					Component: globaldedicatedv1.ComponentType_TIKV,
					SpuId:     lo.ToPtr(spu.Id),
				}).Return([]*entity.Sku{sku1, sku2}, nil).AnyTimes()
				skuRepo.EXPECT().QuerySpuMappings(gomock.Any(), []uint64{spu.Id}).Return(spuMappings, nil).AnyTimes()

				tierUnits, err := svc.SelectSkusByComponentType(ctx, globaldedicatedv1.ComponentType_TIKV, parameter)
				require.NoError(t, err)
				tierUnit := tierUnits[0]
				assert.Equal(t, sku2.Id, tierUnit.Id)
				assert.Equal(t, sku2.Key, tierUnit.Key)
				assert.Equal(t, globaldedicatedv1.InstanceMu_ARM64, tierUnit.GetCpuArch())
				assert.Equal(t, globaldedicatedv1.ComponentType_TIKV, tierUnit.GetComponentType())
			})
		})

		t.Run("should fail", func(t *testing.T) {
			t.Run("empty tenantId", func(t *testing.T) {
				tmp := *parameter
				tmp.TenantId = 0
				_, err := svc.SelectSkusByComponentType(ctx, globaldedicatedv1.ComponentType_TIKV, &tmp)
				assert.Error(t, err)
			})

			t.Run("wrong tenantId", func(t *testing.T) {
				parameter.TenantId = 9991
				parameter.Region = globalcommonv1.Region_AWS_US_WEST_2
				parameter.CpuArch = lo.ToPtr(globaldedicatedv1.InstanceMu_ARM64)

				spuRepo.EXPECT().Query(gomock.Any(), &spurepo.QuerySpuFilter{
					Name:           "8C16G",
					ComponentTypes: []globaldedicatedv1.ComponentType{globaldedicatedv1.ComponentType_TIKV},
				}).Return([]*entity.Spu{spu}, nil).AnyTimes()
				skuRepo.EXPECT().Query(gomock.Any(), &spurepo.SkuFilter{
					Component: globaldedicatedv1.ComponentType_TIKV,
					SpuId:     lo.ToPtr(spu.Id),
				}).Return([]*entity.Sku{sku1, sku2}, nil).AnyTimes()
				skuRepo.EXPECT().QuerySpuMappings(gomock.Any(), []uint64{spu.Id}).Return(spuMappings, nil).AnyTimes()

				tierUnits, err := svc.SelectSkusByComponentType(ctx, globaldedicatedv1.ComponentType_TIKV, parameter)
				require.Error(t, err)
				assert.Contains(t, err.ErrMsg(), "unable to select sku")
				assert.Empty(t, tierUnits)
			})

			t.Run("wrong region", func(t *testing.T) {
				parameter.TenantId = 111
				parameter.Region = globalcommonv1.Region_AWS_AF_SOUTH_1
				parameter.CpuArch = lo.ToPtr(globaldedicatedv1.InstanceMu_ARM64)

				spuRepo.EXPECT().Query(gomock.Any(), &spurepo.QuerySpuFilter{
					Name:           "8C16G",
					ComponentTypes: []globaldedicatedv1.ComponentType{globaldedicatedv1.ComponentType_TIKV},
				}).Return([]*entity.Spu{spu}, nil).AnyTimes()
				skuRepo.EXPECT().Query(gomock.Any(), &spurepo.SkuFilter{
					Component: globaldedicatedv1.ComponentType_TIKV,
					SpuId:     lo.ToPtr(spu.Id),
				}).Return([]*entity.Sku{sku1, sku2}, nil).AnyTimes()
				skuRepo.EXPECT().QuerySpuMappings(gomock.Any(), []uint64{spu.Id}).Return(spuMappings, nil).AnyTimes()

				tierUnits, err := svc.SelectSkusByComponentType(ctx, globaldedicatedv1.ComponentType_TIKV, parameter)
				require.Error(t, err)
				assert.Contains(t, err.ErrMsg(), "unable to select sku")
				assert.Empty(t, tierUnits)
			})

			t.Run("empty provider", func(t *testing.T) {
				tmp := *parameter
				tmp.Provider = globalcommonv1.CloudProvider_NAME_UNSPECIFIED
				_, err := svc.SelectSkusByComponentType(ctx, globaldedicatedv1.ComponentType_TIKV, &tmp)
				assert.Error(t, err)
			})

			t.Run("empty spu", func(t *testing.T) {
				spuRepo.EXPECT().Query(gomock.Any(), &spurepo.QuerySpuFilter{
					Name:           "8C16G",
					ComponentTypes: []globaldedicatedv1.ComponentType{globaldedicatedv1.ComponentType_TIKV},
				}).Return(nil, nil)

				_, err := svc.SelectSkusByComponentType(ctx, globaldedicatedv1.ComponentType_TIKV, parameter)
				assert.Error(t, err)
			})

			t.Run("wrong cpu arch", func(t *testing.T) {
				parameter.TenantId = 111
				parameter.Region = globalcommonv1.Region_AWS_US_WEST_2
				parameter.CpuArch = lo.ToPtr(globaldedicatedv1.InstanceMu_AMD64)

				spuRepo.EXPECT().Query(gomock.Any(), &spurepo.QuerySpuFilter{
					Name:           "8C16G",
					ComponentTypes: []globaldedicatedv1.ComponentType{globaldedicatedv1.ComponentType_TIKV},
				}).Return([]*entity.Spu{spu}, nil).AnyTimes()
				skuRepo.EXPECT().Query(gomock.Any(), &spurepo.SkuFilter{
					Component: globaldedicatedv1.ComponentType_TIKV,
					SpuId:     lo.ToPtr(spu.Id),
				}).Return([]*entity.Sku{sku1, sku2}, nil).AnyTimes()
				skuRepo.EXPECT().QuerySpuMappings(gomock.Any(), []uint64{spu.Id}).Return(spuMappings, nil).AnyTimes()

				tierUnits, err := svc.SelectSkusByComponentType(ctx, globaldedicatedv1.ComponentType_TIKV, parameter)
				require.Error(t, err)
				assert.Contains(t, err.ErrMsg(), "unable to select sku")
				assert.Empty(t, tierUnits)
			})
		})
	})
}

func TestSelectSkuWithCpuGenerationFilterOnAWS(t *testing.T) {
	ctx := t.Context()
	ctrl := gomock.NewController(t, gomock.WithOverridableExpectations())
	defer ctrl.Finish()

	spuRepo := mock_spu.NewMockRepoSpu(ctrl)
	skuRepo := mock_spu.NewMockSkuRepo(ctrl)

	svc := NewSkuService(spuRepo, skuRepo, nil, nil)

	parameter := &SelectSkusRequest{
		TenantId:      0, // fill in cases
		ProjectId:     222,
		Provider:      globalcommonv1.CloudProvider_AWS,
		Region:        globalcommonv1.Region_AWS_US_EAST_1,
		CpuArch:       lo.ToPtr(globaldedicatedv1.InstanceMu_ARM64),
		CpuGeneration: nil, // fill in cases
		SpuName:       lo.ToPtr("8C32G"),
	}
	spu := &entity.Spu{
		Id:            1,
		Key:           "8C32G",
		ComponentType: globaldedicatedv1.ComponentType_TIKV,
		CpuCore:       8,
		MemorySizeGi:  32,
	}
	// sku1 is a sku with g4 cpu generation
	sku1 := &entity.Sku{
		Id:             1,
		Key:            "aws:tikv:m8g.2xlarge:gp3:4000:297:Standard",
		CloudProvider:  globalcommonv1.CloudProvider_AWS,
		ComponentType:  globaldedicatedv1.ComponentType_TIKV,
		SpuId:          1,
		SpuName:        "",
		SpuKey:         "8C32G",
		SpuDisplayName: "",
		Priority:       1,
		InstanceMuId:   1,
		CpuArch:        globaldedicatedv1.InstanceMu_ARM64,
		CpuGeneration:  globaldedicatedv1.InstanceMu_GRAVITON4,
		InstanceType:   "m8g.2xlarge",
		CpuCore:        8,
		MemoryGi:       32,
	}
	// sku2 is a sku with g2 cpu generation
	sku2 := &entity.Sku{
		Id:             2,
		Key:            "aws:tikv:m6g.2xlarge:gp3:4000:297:Standard",
		CloudProvider:  globalcommonv1.CloudProvider_AWS,
		ComponentType:  globaldedicatedv1.ComponentType_TIKV,
		SpuId:          1,
		SpuName:        "",
		SpuKey:         "8C32G",
		SpuDisplayName: "",
		Priority:       1,
		InstanceMuId:   2,
		CpuArch:        globaldedicatedv1.InstanceMu_ARM64,
		CpuGeneration:  globaldedicatedv1.InstanceMu_GRAVITON2,
		InstanceType:   "m6g.2xlarge",
		CpuCore:        8,
		MemoryGi:       32,
	}
	spuMappings := map[uint64][]*entity.SpuMapping{
		spu.Id: {
			{
				SpuId:    spu.Id,
				SkuKey:   sku1.Key,
				Provider: globalcommonv1.CloudProvider_AWS,
				Regions: map[globalcommonv1.Region_Name]bool{
					globalcommonv1.Region_AWS_US_WEST_2: true,
					globalcommonv1.Region_AWS_US_EAST_1: true,
				},
				TenantIDs: map[uint64]bool{
					111: true,
				},
			},
			{
				SpuId:    spu.Id,
				SkuKey:   sku2.Key,
				Provider: globalcommonv1.CloudProvider_AWS,
				Regions: map[globalcommonv1.Region_Name]bool{
					globalcommonv1.Region_AWS_US_WEST_2: true,
					globalcommonv1.Region_AWS_US_EAST_1: true,
				},
				TenantIDs: map[uint64]bool{
					111: true,
					222: true,
				},
			},
		},
	}

	t.Run("create cluster", func(t *testing.T) {
		t.Run("should use g4 for tenant 111", func(t *testing.T) {
			parameter.TenantId = 111

			spuRepo.EXPECT().Query(gomock.Any(), &spurepo.QuerySpuFilter{
				Name:           "8C32G",
				ComponentTypes: []globaldedicatedv1.ComponentType{globaldedicatedv1.ComponentType_TIKV},
			}).Return([]*entity.Spu{spu}, nil).AnyTimes()
			skuRepo.EXPECT().Query(gomock.Any(), &spurepo.SkuFilter{
				Component: globaldedicatedv1.ComponentType_TIKV,
				SpuId:     lo.ToPtr(spu.Id),
			}).Return([]*entity.Sku{sku1, sku2}, nil).AnyTimes()
			skuRepo.EXPECT().QuerySpuMappings(gomock.Any(), []uint64{spu.Id}).Return(spuMappings, nil).AnyTimes()

			skus, err := svc.SelectSkusByComponentType(ctx, globaldedicatedv1.ComponentType_TIKV, parameter)
			require.NoError(t, err)
			sku := skus[0]
			assert.Equal(t, sku1.Id, sku.Id)
			assert.Equal(t, sku1.Key, sku.Key)
			assert.Equal(t, globaldedicatedv1.InstanceMu_ARM64, sku.GetCpuArch())
			assert.Equal(t, globaldedicatedv1.ComponentType_TIKV, sku.GetComponentType())
			assert.Equal(t, globaldedicatedv1.InstanceMu_GRAVITON4, sku.GetCpuGeneration())
		})

		t.Run("should use g2 for tenant 222", func(t *testing.T) {
			parameter.TenantId = 222

			spuRepo.EXPECT().Query(gomock.Any(), &spurepo.QuerySpuFilter{
				Name:           "8C32G",
				ComponentTypes: []globaldedicatedv1.ComponentType{globaldedicatedv1.ComponentType_TIKV},
			}).Return([]*entity.Spu{spu}, nil).AnyTimes()
			skuRepo.EXPECT().Query(gomock.Any(), &spurepo.SkuFilter{
				Component: globaldedicatedv1.ComponentType_TIKV,
				SpuId:     lo.ToPtr(spu.Id),
			}).Return([]*entity.Sku{sku1, sku2}, nil).AnyTimes()
			skuRepo.EXPECT().QuerySpuMappings(gomock.Any(), []uint64{spu.Id}).Return(spuMappings, nil).AnyTimes()

			skus, err := svc.SelectSkusByComponentType(ctx, globaldedicatedv1.ComponentType_TIKV, parameter)
			require.NoError(t, err)
			sku := skus[0]
			assert.Equal(t, sku2.Id, sku.Id)
			assert.Equal(t, sku2.Key, sku.Key)
			assert.Equal(t, globaldedicatedv1.InstanceMu_ARM64, sku.GetCpuArch())
			assert.Equal(t, globaldedicatedv1.ComponentType_TIKV, sku.GetComponentType())
			assert.Equal(t, globaldedicatedv1.InstanceMu_GRAVITON2, sku.GetCpuGeneration())
		})
	})

	t.Run("scale cluster", func(t *testing.T) {
		t.Run("should use g4 for tenant 111 if cluster is g4", func(t *testing.T) {
			parameter.TenantId = 111
			parameter.CpuGeneration = lo.ToPtr(globaldedicatedv1.InstanceMu_GRAVITON4)

			spuRepo.EXPECT().Query(gomock.Any(), &spurepo.QuerySpuFilter{
				Name:           "8C32G",
				ComponentTypes: []globaldedicatedv1.ComponentType{globaldedicatedv1.ComponentType_TIKV},
			}).Return([]*entity.Spu{spu}, nil).AnyTimes()
			skuRepo.EXPECT().Query(gomock.Any(), &spurepo.SkuFilter{
				Component: globaldedicatedv1.ComponentType_TIKV,
				SpuId:     lo.ToPtr(spu.Id),
			}).Return([]*entity.Sku{sku1, sku2}, nil).AnyTimes()
			skuRepo.EXPECT().QuerySpuMappings(gomock.Any(), []uint64{spu.Id}).Return(spuMappings, nil).AnyTimes()

			skus, err := svc.SelectSkusByComponentType(ctx, globaldedicatedv1.ComponentType_TIKV, parameter)
			require.NoError(t, err)
			sku := skus[0]
			assert.Equal(t, sku1.Id, sku.Id)
			assert.Equal(t, sku1.Key, sku.Key)
			assert.Equal(t, globaldedicatedv1.InstanceMu_ARM64, sku.GetCpuArch())
			assert.Equal(t, globaldedicatedv1.ComponentType_TIKV, sku.GetComponentType())
			assert.Equal(t, globaldedicatedv1.InstanceMu_GRAVITON4, sku.GetCpuGeneration())
		})

		t.Run("should use g2 for tenant 111 if cluster is g2", func(t *testing.T) {
			parameter.TenantId = 111
			parameter.CpuGeneration = lo.ToPtr(globaldedicatedv1.InstanceMu_GRAVITON2)

			spuRepo.EXPECT().Query(gomock.Any(), &spurepo.QuerySpuFilter{
				Name:           "8C32G",
				ComponentTypes: []globaldedicatedv1.ComponentType{globaldedicatedv1.ComponentType_TIKV},
			}).Return([]*entity.Spu{spu}, nil).AnyTimes()
			skuRepo.EXPECT().Query(gomock.Any(), &spurepo.SkuFilter{
				Component: globaldedicatedv1.ComponentType_TIKV,
				SpuId:     lo.ToPtr(spu.Id),
			}).Return([]*entity.Sku{sku1, sku2}, nil).AnyTimes()
			skuRepo.EXPECT().QuerySpuMappings(gomock.Any(), []uint64{spu.Id}).Return(spuMappings, nil).AnyTimes()

			skus, err := svc.SelectSkusByComponentType(ctx, globaldedicatedv1.ComponentType_TIKV, parameter)
			require.NoError(t, err)
			sku := skus[0]
			assert.Equal(t, sku2.Id, sku.Id)
			assert.Equal(t, sku2.Key, sku.Key)
			assert.Equal(t, globaldedicatedv1.InstanceMu_ARM64, sku.GetCpuArch())
			assert.Equal(t, globaldedicatedv1.ComponentType_TIKV, sku.GetComponentType())
			assert.Equal(t, globaldedicatedv1.InstanceMu_GRAVITON2, sku.GetCpuGeneration())
		})

		t.Run("should use g2 for tenant 222 if cluster is g2", func(t *testing.T) {
			parameter.TenantId = 222
			parameter.CpuGeneration = lo.ToPtr(globaldedicatedv1.InstanceMu_GRAVITON2)

			spuRepo.EXPECT().Query(gomock.Any(), &spurepo.QuerySpuFilter{
				Name:           "8C32G",
				ComponentTypes: []globaldedicatedv1.ComponentType{globaldedicatedv1.ComponentType_TIKV},
			}).Return([]*entity.Spu{spu}, nil).AnyTimes()
			skuRepo.EXPECT().Query(gomock.Any(), &spurepo.SkuFilter{
				Component: globaldedicatedv1.ComponentType_TIKV,
				SpuId:     lo.ToPtr(spu.Id),
			}).Return([]*entity.Sku{sku1, sku2}, nil).AnyTimes()
			skuRepo.EXPECT().QuerySpuMappings(gomock.Any(), []uint64{spu.Id}).Return(spuMappings, nil).AnyTimes()

			skus, err := svc.SelectSkusByComponentType(ctx, globaldedicatedv1.ComponentType_TIKV, parameter)
			require.NoError(t, err)
			sku := skus[0]
			assert.Equal(t, sku2.Id, sku.Id)
			assert.Equal(t, sku2.Key, sku.Key)
			assert.Equal(t, globaldedicatedv1.InstanceMu_ARM64, sku.GetCpuArch())
			assert.Equal(t, globaldedicatedv1.ComponentType_TIKV, sku.GetComponentType())
			assert.Equal(t, globaldedicatedv1.InstanceMu_GRAVITON2, sku.GetCpuGeneration())
		})

		t.Run("should return error if tenant 222 requests g4", func(t *testing.T) {
			parameter.TenantId = 222
			parameter.CpuGeneration = lo.ToPtr(globaldedicatedv1.InstanceMu_GRAVITON4)

			spuRepo.EXPECT().Query(gomock.Any(), &spurepo.QuerySpuFilter{
				Name:           "8C32G",
				ComponentTypes: []globaldedicatedv1.ComponentType{globaldedicatedv1.ComponentType_TIKV},
			}).Return([]*entity.Spu{spu}, nil).AnyTimes()
			skuRepo.EXPECT().Query(gomock.Any(), &spurepo.SkuFilter{
				Component: globaldedicatedv1.ComponentType_TIKV,
				SpuId:     lo.ToPtr(spu.Id),
			}).Return([]*entity.Sku{sku1, sku2}, nil).AnyTimes()
			skuRepo.EXPECT().QuerySpuMappings(gomock.Any(), []uint64{spu.Id}).Return(spuMappings, nil).AnyTimes()

			skus, err := svc.SelectSkusByComponentType(ctx, globaldedicatedv1.ComponentType_TIKV, parameter)
			require.Error(t, err)
			assert.Contains(t, err.ErrMsg(), "unable to select sku")
			assert.Empty(t, skus)
		})
	})
}

func TestSelectSkuWithCpuGenerationFilterOnGCP(t *testing.T) {
	ctx := t.Context()
	ctrl := gomock.NewController(t, gomock.WithOverridableExpectations())
	defer ctrl.Finish()

	spuRepo := mock_spu.NewMockRepoSpu(ctrl)
	skuRepo := mock_spu.NewMockSkuRepo(ctrl)

	svc := NewSkuService(spuRepo, skuRepo, nil, nil)

	parameter := &SelectSkusRequest{
		TenantId:      111,
		ProjectId:     222,
		Provider:      globalcommonv1.CloudProvider_GCP,
		Region:        globalcommonv1.Region_GCP_US_WEST1,
		CpuArch:       nil, // fill in cases
		CpuGeneration: nil, // fill in cases
		SpuName:       lo.ToPtr("8C32G"),
	}
	spu := &entity.Spu{
		Id:            1,
		Key:           "8C32G",
		ComponentType: globaldedicatedv1.ComponentType_TIKV,
		CpuCore:       8,
		MemorySizeGi:  32,
	}
	n4Sku := &entity.Sku{
		Id:             1,
		Key:            "gcp:tikv:n4-standard-8:hyperdisk-balanced:4000:297",
		CloudProvider:  globalcommonv1.CloudProvider_GCP,
		ComponentType:  globaldedicatedv1.ComponentType_TIKV,
		SpuId:          1,
		SpuName:        "",
		SpuKey:         "8C32G",
		SpuDisplayName: "",
		Priority:       1,
		InstanceMuId:   1,
		CpuArch:        globaldedicatedv1.InstanceMu_AMD64,
		CpuGeneration:  globaldedicatedv1.InstanceMu_GENERATION4_INTEL_XEON,
		InstanceType:   "n4-standard-8",
		CpuCore:        8,
		MemoryGi:       32,
	}
	n2Sku := &entity.Sku{
		Id:             2,
		Key:            "gcp:tikv:n2-standard-8:pd-ssd:0:0",
		CloudProvider:  globalcommonv1.CloudProvider_GCP,
		ComponentType:  globaldedicatedv1.ComponentType_TIKV,
		SpuId:          1,
		SpuName:        "",
		SpuKey:         "8C32G",
		SpuDisplayName: "",
		Priority:       1,
		InstanceMuId:   2,
		CpuArch:        globaldedicatedv1.InstanceMu_AMD64,
		CpuGeneration:  globaldedicatedv1.InstanceMu_CPU_GENERATION_UNSPECIFIED,
		InstanceType:   "n2-standard-8",
		CpuCore:        8,
		MemoryGi:       32,
	}
	c4aSku := &entity.Sku{
		Id:             3,
		Key:            "gcp:tikv:c4a-standard-8:pd-ssd:0:0",
		CloudProvider:  globalcommonv1.CloudProvider_GCP,
		ComponentType:  globaldedicatedv1.ComponentType_TIKV,
		SpuId:          1,
		SpuName:        "",
		SpuKey:         "8C32G",
		SpuDisplayName: "",
		Priority:       1,
		InstanceMuId:   3,
		CpuArch:        globaldedicatedv1.InstanceMu_ARM64,
		CpuGeneration:  globaldedicatedv1.InstanceMu_GENERATION4_AXION,
		InstanceType:   "c4a-standard-8",
		CpuCore:        8,
		MemoryGi:       32,
	}
	spuMappings := map[uint64][]*entity.SpuMapping{
		spu.Id: {
			{
				SpuId:    spu.Id,
				SkuKey:   n4Sku.Key,
				Provider: globalcommonv1.CloudProvider_GCP,
				Regions: map[globalcommonv1.Region_Name]bool{
					globalcommonv1.Region_GCP_US_WEST2: true,
					globalcommonv1.Region_GCP_US_WEST1: true,
				},
				TenantIDs: map[uint64]bool{
					0: true,
				},
			},
			{
				SpuId:    spu.Id,
				SkuKey:   n2Sku.Key,
				Provider: globalcommonv1.CloudProvider_GCP,
				Regions: map[globalcommonv1.Region_Name]bool{
					globalcommonv1.Region_GCP_US_WEST2: true,
					globalcommonv1.Region_GCP_US_WEST1: true,
				},
				TenantIDs: map[uint64]bool{
					0: true,
				},
			},
			{
				SpuId:    spu.Id,
				SkuKey:   c4aSku.Key,
				Provider: globalcommonv1.CloudProvider_GCP,
				Regions: map[globalcommonv1.Region_Name]bool{
					globalcommonv1.Region_GCP_US_WEST2: true,
					globalcommonv1.Region_GCP_US_WEST1: true,
				},
				TenantIDs: map[uint64]bool{
					0: true,
				},
			},
		},
	}

	t.Run("require amd64 and gen4 intel xeon, return n4", func(t *testing.T) {
		parameter.CpuArch = lo.ToPtr(globaldedicatedv1.InstanceMu_AMD64)
		parameter.CpuGeneration = lo.ToPtr(globaldedicatedv1.InstanceMu_GENERATION4_INTEL_XEON)

		spuRepo.EXPECT().Query(gomock.Any(), &spurepo.QuerySpuFilter{
			Name:           "8C32G",
			ComponentTypes: []globaldedicatedv1.ComponentType{globaldedicatedv1.ComponentType_TIKV},
		}).Return([]*entity.Spu{spu}, nil).AnyTimes()
		skuRepo.EXPECT().Query(gomock.Any(), &spurepo.SkuFilter{
			Component: globaldedicatedv1.ComponentType_TIKV,
			SpuId:     lo.ToPtr(spu.Id),
		}).Return([]*entity.Sku{n4Sku, n2Sku, c4aSku}, nil).AnyTimes()
		skuRepo.EXPECT().QuerySpuMappings(gomock.Any(), []uint64{spu.Id}).Return(spuMappings, nil).AnyTimes()

		skus, err := svc.SelectSkusByComponentType(ctx, globaldedicatedv1.ComponentType_TIKV, parameter)
		require.NoError(t, err)
		sku := skus[0]
		assert.Equal(t, n4Sku.Id, sku.Id)
		assert.Equal(t, n4Sku.Key, sku.Key)
		assert.Equal(t, globaldedicatedv1.InstanceMu_AMD64, sku.GetCpuArch())
		assert.Equal(t, globaldedicatedv1.ComponentType_TIKV, sku.GetComponentType())
		assert.Equal(t, globaldedicatedv1.InstanceMu_GENERATION4_INTEL_XEON, sku.GetCpuGeneration())
	})

	t.Run("require amd64 and gen4 amd epyc, return error", func(t *testing.T) {
		parameter.CpuArch = lo.ToPtr(globaldedicatedv1.InstanceMu_AMD64)
		parameter.CpuGeneration = lo.ToPtr(globaldedicatedv1.InstanceMu_GENERATION4_AMD_EPYC)

		skus, err := svc.SelectSkusByComponentType(ctx, globaldedicatedv1.ComponentType_TIKV, parameter)
		require.Error(t, err)
		assert.Contains(t, err.ErrMsg(), "unable to select sku")
		assert.Empty(t, skus)
	})

	t.Run("require amd64 and gen4 axion, return error", func(t *testing.T) {
		parameter.CpuArch = lo.ToPtr(globaldedicatedv1.InstanceMu_AMD64)
		parameter.CpuGeneration = lo.ToPtr(globaldedicatedv1.InstanceMu_GENERATION4_AXION)

		skus, err := svc.SelectSkusByComponentType(ctx, globaldedicatedv1.ComponentType_TIKV, parameter)
		require.Error(t, err)
		assert.Contains(t, err.ErrMsg(), "unable to select sku")
		assert.Empty(t, skus)
	})

	t.Run("require arm64 and gen4 axion, return c4a", func(t *testing.T) {
		parameter.CpuArch = lo.ToPtr(globaldedicatedv1.InstanceMu_ARM64)
		parameter.CpuGeneration = lo.ToPtr(globaldedicatedv1.InstanceMu_GENERATION4_AXION)

		skus, err := svc.SelectSkusByComponentType(ctx, globaldedicatedv1.ComponentType_TIKV, parameter)
		require.NoError(t, err)
		sku := skus[0]
		assert.Equal(t, c4aSku.Id, sku.Id)
		assert.Equal(t, c4aSku.Key, sku.Key)
		assert.Equal(t, globaldedicatedv1.InstanceMu_ARM64, sku.GetCpuArch())
		assert.Equal(t, globaldedicatedv1.ComponentType_TIKV, sku.GetComponentType())
		assert.Equal(t, globaldedicatedv1.InstanceMu_GENERATION4_AXION, sku.GetCpuGeneration())
	})

	t.Run("require arm64 and gen4 intel xeon, return error", func(t *testing.T) {
		parameter.CpuArch = lo.ToPtr(globaldedicatedv1.InstanceMu_ARM64)
		parameter.CpuGeneration = lo.ToPtr(globaldedicatedv1.InstanceMu_GENERATION4_INTEL_XEON)

		skus, err := svc.SelectSkusByComponentType(ctx, globaldedicatedv1.ComponentType_TIKV, parameter)
		require.Error(t, err)
		assert.Contains(t, err.ErrMsg(), "unable to select sku")
		assert.Empty(t, skus)
	})

	t.Run("require amd64 and non-gen4, return n2", func(t *testing.T) {
		parameter.CpuArch = lo.ToPtr(globaldedicatedv1.InstanceMu_AMD64)
		parameter.CpuGeneration = lo.ToPtr(globaldedicatedv1.InstanceMu_CPU_GENERATION_UNSPECIFIED)

		skus, err := svc.SelectSkusByComponentType(ctx, globaldedicatedv1.ComponentType_TIKV, parameter)
		require.NoError(t, err)
		sku := skus[0]
		assert.Equal(t, n2Sku.Id, sku.Id)
		assert.Equal(t, n2Sku.Key, sku.Key)
		assert.Equal(t, n2Sku.GetCpuArch(), sku.GetCpuArch())
		assert.Equal(t, globaldedicatedv1.ComponentType_TIKV, sku.GetComponentType())
		assert.Equal(t, globaldedicatedv1.InstanceMu_CPU_GENERATION_UNSPECIFIED, sku.GetCpuGeneration())
	})

	t.Run("only require amd64, return n4", func(t *testing.T) {
		parameter.CpuArch = lo.ToPtr(globaldedicatedv1.InstanceMu_AMD64)
		parameter.CpuGeneration = nil

		skus, err := svc.SelectSkusByComponentType(ctx, globaldedicatedv1.ComponentType_TIKV, parameter)
		require.NoError(t, err)
		sku := skus[0]
		assert.Equal(t, n4Sku.Id, sku.Id)
		assert.Equal(t, n4Sku.Key, sku.Key)
		assert.Equal(t, globaldedicatedv1.InstanceMu_AMD64, sku.GetCpuArch())
		assert.Equal(t, globaldedicatedv1.ComponentType_TIKV, sku.GetComponentType())
	})

	t.Run("only require arm64, return c4a", func(t *testing.T) {
		parameter.CpuArch = lo.ToPtr(globaldedicatedv1.InstanceMu_ARM64)
		parameter.CpuGeneration = nil

		skus, err := svc.SelectSkusByComponentType(ctx, globaldedicatedv1.ComponentType_TIKV, parameter)
		require.NoError(t, err)
		sku := skus[0]
		assert.Equal(t, c4aSku.Id, sku.Id)
		assert.Equal(t, c4aSku.Key, sku.Key)
		assert.Equal(t, globaldedicatedv1.InstanceMu_ARM64, sku.GetCpuArch())
		assert.Equal(t, globaldedicatedv1.ComponentType_TIKV, sku.GetComponentType())
	})
}