package cluster

import (
	"context"
	"math/rand/v2"
	"strconv"
	"testing"

	"github.com/samber/lo"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	regionaldedicatedv1 "github.com/tidbcloud/infra-provider/api/tidbcloud/regional-server/v1"
	"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/adapter/account"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/adapter/account/mock_account"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/adapter/billing"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/adapter/billing/mock_billing"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/adapter/clusterconfig/mock_clusterconfig"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/adapter/clustermetric/mock_clustermetric"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/adapter/regionserver/mock_regionserver"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/adapter/remoteencryptor/mock_remoteencryptor"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/entity"
	spusvc "github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/service/spu"
	clusterrepoimpl "github.com/tidbcloud/tidb-management-service/dedicated-global-service/infra/repoimpl/cluster"
	regionimpl "github.com/tidbcloud/tidb-management-service/dedicated-global-service/infra/repoimpl/region"
	spurepoimpl "github.com/tidbcloud/tidb-management-service/dedicated-global-service/infra/repoimpl/spu"
)

func TestCreateClusterN2N4Coexist(t *testing.T) {
	ctrl := gomock.NewController(t, gomock.WithOverridableExpectations())

	encryptor := mock_remoteencryptor.NewMockClient(ctrl)
	encryptor.EXPECT().EncryptBase64(gomock.Any(), gomock.Any()).Return("encrypted", nil).AnyTimes()
	encryptor.EXPECT().DecryptBase64(gomock.Any(), gomock.Any()).Return("password", nil).AnyTimes()

	accountCli := mock_account.NewMockClient(ctrl)
	accountCli.EXPECT().GetProject(gomock.Any(), uint64(ProjectID)).Return(&account.ProjectInfo{
		ID:             ProjectID,
		TenantID:       TenantID,
		Name:           "test",
		AwsCmekEnabled: false,
	}, nil).AnyTimes()
	accountCli.EXPECT().GetTenant(gomock.Any(), uint64(TenantID)).Return(&account.TenantInfo{
		ID:   TenantID,
		Name: "test",
		Plan: account.TenantPlanPOC,
	}, nil).AnyTimes()

	regionalCli := mock_regionserver.NewMockClient(ctrl)
	regionalCli.EXPECT().ListVpcs(gomock.Any(), gomock.Any(), &regionaldedicatedv1.ListVpcsRequest{
		TenantId:  strconv.FormatUint(TenantID, 10),
		ProjectId: strconv.FormatUint(ProjectID, 10),
		Statuses: []regionaldedicatedv1.Vpc_Status{
			regionaldedicatedv1.Vpc_ACTIVE,
			regionaldedicatedv1.Vpc_CREATING,
			regionaldedicatedv1.Vpc_FAILED,
		},
	}).Return(&regionaldedicatedv1.ListVpcsResponse{
		Vpcs: []*regionaldedicatedv1.Vpc{
			{
				VpcId:      "vpc-xxxxxxxxxxxxxxxxx",
				TenantIds:  []string{strconv.FormatUint(TenantID, 10)},
				ProjectIds: []string{strconv.FormatUint(ProjectID, 10)},
				CidrBlock:  "0.0.0.0/0",
			},
		},
	}, nil).AnyTimes()

	metricCli := mock_clustermetric.NewMockClient(ctrl)

	billingCli := mock_billing.NewMockClient(ctrl)
	billingCli.EXPECT().
		BillingActionCheck(gomock.Any(), billing.Action_Create, uint64(TenantID), gomock.Any()).
		Return(&billing.ActionCheckResult{Pass: true}, nil).
		AnyTimes()

	clusterConfigCli := mock_clusterconfig.NewMockClient(ctrl)
	clusterConfigCli.EXPECT().GetClusterVersion(gomock.Any(), gomock.Any()).Return(ClusterVersion, nil).AnyTimes()
	clusterConfigCli.EXPECT().
		ListComponentConfigurations(gomock.Any(), gomock.Any()).
		Return([]*regionaldedicatedv1.ComponentConfiguration{}, nil).
		AnyTimes()

	svc := newTestService(encryptor, accountCli, regionalCli, metricCli, billingCli, clusterConfigCli, false)

	t.Run("n2 and n4 tikv mapped to the same pd spu, create n2 tikv should succeed", func(t *testing.T) {
		req := &globaldedicatedv1.CreateClusterRequest{
			Cluster: &globaldedicatedv1.Cluster{
				ClusterId:       uint64(rand.Int64()),
				OrganizationId:  TenantID,
				ProjectId:       ProjectID,
				CloudProviderId: globalcommonv1.CloudProvider_GCP,
				RegionId:        globalcommonv1.Region_GCP_US_WEST1,
				Port:            4000,
				ClusterComponents: []*globaldedicatedv1.Cluster_ClusterComponent{
					{
						ComponentType: globaldedicatedv1.ComponentType_TIDB,
						SpuName:       "TidbGcpX86N2",
						Count:         lo.ToPtr(uint32(1)),
					},
					{
						ComponentType: globaldedicatedv1.ComponentType_TIKV,
						SpuName:       "TikvGcpX86N2",
						Count:         lo.ToPtr(uint32(3)),
						Storages: map[string]*globaldedicatedv1.Cluster_ClusterComponent_Storage{
							globaldedicatedv1.Sku_DATA.String(): {
								VolumeType: lo.ToPtr(globaldedicatedv1.StorageMu_PD_SSD),
								SizeGi:     555,
							},
							globaldedicatedv1.Sku_RAFT_LOG.String(): {
								VolumeType: lo.ToPtr(globaldedicatedv1.StorageMu_VOLUME_TYPE_UNSPECIFIED),
							},
						},
					},
				},
			},
		}
		resCluster, err := svc.Create(t.Context(), req)
		require.NoError(t, err)

		pdComponent, _ := lo.Find(resCluster.GetClusterComponents(), func(c *globaldedicatedv1.Cluster_ClusterComponent) bool {
			return c.GetComponentType() == globaldedicatedv1.ComponentType_PD
		})
		assert.Equal(t, "PdGcp", pdComponent.GetSpuName())
		assert.Equal(t, "PdGcpX86N2", pdComponent.GetSkuKey())
	})

	t.Run("n2 and n4 tikv mapped to the same pd spu, create n4 tikv should succeed", func(t *testing.T) {
		req := &globaldedicatedv1.CreateClusterRequest{
			Cluster: &globaldedicatedv1.Cluster{
				ClusterId:       uint64(rand.Int64()),
				OrganizationId:  TenantID,
				ProjectId:       ProjectID,
				CloudProviderId: globalcommonv1.CloudProvider_GCP,
				RegionId:        globalcommonv1.Region_GCP_US_WEST1,
				Port:            4000,
				ClusterComponents: []*globaldedicatedv1.Cluster_ClusterComponent{
					{
						ComponentType: globaldedicatedv1.ComponentType_TIDB,
						SpuName:       "TidbGcpX86N4",
						Count:         lo.ToPtr(uint32(1)),
					},
					{
						ComponentType: globaldedicatedv1.ComponentType_TIKV,
						SpuName:       "TikvGcpX86N4",
						Count:         lo.ToPtr(uint32(3)),
						Storages: map[string]*globaldedicatedv1.Cluster_ClusterComponent_Storage{
							globaldedicatedv1.Sku_DATA.String(): {
								VolumeType: lo.ToPtr(globaldedicatedv1.StorageMu_PD_SSD),
								SizeGi:     555,
							},
							globaldedicatedv1.Sku_RAFT_LOG.String(): {
								VolumeType: lo.ToPtr(globaldedicatedv1.StorageMu_VOLUME_TYPE_UNSPECIFIED),
							},
						},
					},
				},
			},
		}
		resCluster, err := svc.Create(t.Context(), req)
		require.NoError(t, err)

		pdComponent, _ := lo.Find(resCluster.GetClusterComponents(), func(c *globaldedicatedv1.Cluster_ClusterComponent) bool {
			return c.GetComponentType() == globaldedicatedv1.ComponentType_PD
		})
		assert.Equal(t, "PdGcp", pdComponent.GetSpuName())
		assert.Equal(t, "PdGcpX86N4", pdComponent.GetSkuKey())
	})
}

func initN2N4SkuRecords(ctx context.Context) error {
	spuRepo := spurepoimpl.NewSpuRepo(*dbProvider.Load())
	skuRepo := spurepoimpl.NewSkuRepo(*dbProvider.Load())
	instanceMuRepo := spurepoimpl.NewInstanceMuRepo(*dbProvider.Load())
	storageMuRepo := spurepoimpl.NewStorageMuRepo(*dbProvider.Load())
	regionRepo := regionimpl.NewRegionRepo(*dbProvider.Load())
	clusterRepo := clusterrepoimpl.NewClusterRepo(*dbProvider.Load())
	skuSvc := spusvc.NewSkuService(spuRepo, skuRepo, instanceMuRepo, storageMuRepo)
	spuSvc := spusvc.NewSpuService(spuRepo, skuRepo, regionRepo, clusterRepo, skuSvc, nil, *dbProvider.Load())

	// instance MUs.
	// GCP x86 N2 instance - older generation
	instanceIdGcpX86N2, err := instanceMuRepo.Save(ctx, &entity.InstanceMu{
		CloudProvider: globalcommonv1.CloudProvider_GCP,
		InstanceType:  "n2-test",
		CpuArch:       globaldedicatedv1.InstanceMu_AMD64,
		CpuCore:       8,
		MemoryGi:      16,
	})
	if err != nil {
		return err
	}

	// GCP x86 N4 instance - newer generation
	instanceIdGcpX86N4, err := instanceMuRepo.Save(ctx, &entity.InstanceMu{
		CloudProvider: globalcommonv1.CloudProvider_GCP,
		InstanceType:  "n4-test",
		CpuArch:       globaldedicatedv1.InstanceMu_AMD64,
		CpuCore:       8,
		MemoryGi:      16,
	})
	if err != nil {
		return err
	}

	// storage MUs.
	pdssdId, err := storageMuRepo.Save(ctx, &entity.StorageMu{
		CloudProvider: globalcommonv1.CloudProvider_GCP,
		VolumeType:    globaldedicatedv1.StorageMu_PD_SSD,
	})
	if err != nil {
		return err
	}

	// tidb spu.
	tidbSpuGcpX86N2, err := spuSvc.Create(ctx, &entity.Spu{
		DisplayName:   "TidbGcpX86N2",
		Name:          "TidbGcpX86N2",
		ComponentType: globaldedicatedv1.ComponentType_TIDB,
		CpuCore:       8,
		MemorySizeGi:  16,
		Mappings: []*globaldedicatedv1.Spu_Mapping{
			{
				CloudProvider:   globalcommonv1.CloudProvider_GCP,
				PhysicalRegions: []globalcommonv1.Region_Name{globalcommonv1.Region_NAME_UNSPECIFIED},
				OrganizationIds: []uint64{0},
				SkuKey:          "TidbGcpX86N2",
			},
		},
	})
	if err != nil {
		return err
	}

	tidbSpuGcpX86N4, err := spuSvc.Create(ctx, &entity.Spu{
		DisplayName:   "TidbGcpX86N4",
		Name:          "TidbGcpX86N4",
		ComponentType: globaldedicatedv1.ComponentType_TIDB,
		CpuCore:       8,
		MemorySizeGi:  16,
		Mappings: []*globaldedicatedv1.Spu_Mapping{
			{
				CloudProvider:   globalcommonv1.CloudProvider_GCP,
				PhysicalRegions: []globalcommonv1.Region_Name{globalcommonv1.Region_NAME_UNSPECIFIED},
				OrganizationIds: []uint64{0},
				SkuKey:          "TidbGcpX86N4",
			},
		},
	})
	if err != nil {
		return err
	}

	// tidb sku.
	_, err = skuRepo.Save(ctx, &entity.Sku{
		Key:           "TidbGcpX86N2",
		SpuId:         tidbSpuGcpX86N2.GetId(),
		CloudProvider: globalcommonv1.CloudProvider_GCP,
		ComponentType: globaldedicatedv1.ComponentType_TIDB,
		InstanceMuId:  uint64(instanceIdGcpX86N2),
	})
	if err != nil {
		return err
	}

	_, err = skuRepo.Save(ctx, &entity.Sku{
		Key:           "TidbGcpX86N4",
		SpuId:         tidbSpuGcpX86N4.GetId(),
		CloudProvider: globalcommonv1.CloudProvider_GCP,
		ComponentType: globaldedicatedv1.ComponentType_TIDB,
		InstanceMuId:  uint64(instanceIdGcpX86N4),
	})
	if err != nil {
		return err
	}

	// pd spu.
	pdSpuGcp, err := spuSvc.Create(ctx, &entity.Spu{
		DisplayName:   "PdGcp",
		Name:          "PdGcp",
		ComponentType: globaldedicatedv1.ComponentType_PD,
		CpuCore:       8,
		MemorySizeGi:  16,
		Mappings: []*globaldedicatedv1.Spu_Mapping{
			{
				CloudProvider:   globalcommonv1.CloudProvider_GCP,
				PhysicalRegions: []globalcommonv1.Region_Name{globalcommonv1.Region_NAME_UNSPECIFIED},
				OrganizationIds: []uint64{0},
				SkuKey:          "PdGcpX86N4",
			},
			{
				CloudProvider:   globalcommonv1.CloudProvider_GCP,
				PhysicalRegions: []globalcommonv1.Region_Name{globalcommonv1.Region_NAME_UNSPECIFIED},
				OrganizationIds: []uint64{0},
				SkuKey:          "PdGcpX86N2",
			},
		},
	})
	if err != nil {
		return err
	}

	// pd sku.
	_, err = skuRepo.Save(ctx, &entity.Sku{
		Key:              "PdGcpX86N2",
		SpuId:            pdSpuGcp.GetId(),
		CloudProvider:    globalcommonv1.CloudProvider_GCP,
		ComponentType:    globaldedicatedv1.ComponentType_PD,
		InstanceMuId:     uint64(instanceIdGcpX86N2),
		DefaultNodeCount: 3,
		DataStorage: &globaldedicatedv1.Sku_Storage{
			StorageMuId:   uint64(pdssdId),
			SizeGiDefault: 10,
		},
		Conditions: []*globaldedicatedv1.Sku_Condition{
			{
				Key:      globaldedicatedv1.Sku_Condition_ATTACH_TIERUNIT,
				Operator: globaldedicatedv1.Sku_Condition_IN,
				Values: []string{
					"TikvGcpX86N2",
				},
			},
		},
	})
	if err != nil {
		return err
	}

	_, err = skuRepo.Save(ctx, &entity.Sku{
		Key:              "PdGcpX86N4",
		SpuId:            pdSpuGcp.GetId(),
		CloudProvider:    globalcommonv1.CloudProvider_GCP,
		ComponentType:    globaldedicatedv1.ComponentType_PD,
		InstanceMuId:     uint64(instanceIdGcpX86N4),
		DefaultNodeCount: 3,
		DataStorage: &globaldedicatedv1.Sku_Storage{
			StorageMuId:   uint64(pdssdId),
			SizeGiDefault: 10,
		},
		Conditions: []*globaldedicatedv1.Sku_Condition{
			{
				Key:      globaldedicatedv1.Sku_Condition_ATTACH_TIERUNIT,
				Operator: globaldedicatedv1.Sku_Condition_IN,
				Values: []string{
					"TikvGcpX86N4",
				},
			},
		},
	})
	if err != nil {
		return err
	}

	// tikv spu.
	tikvSpuGcpX86N2, err := spuSvc.Create(ctx, &entity.Spu{
		DisplayName:   "TikvGcpX86N2",
		Name:          "TikvGcpX86N2",
		ComponentType: globaldedicatedv1.ComponentType_TIKV,
		CpuCore:       8,
		MemorySizeGi:  16,
		Mappings: []*globaldedicatedv1.Spu_Mapping{
			{
				CloudProvider:   globalcommonv1.CloudProvider_GCP,
				PhysicalRegions: []globalcommonv1.Region_Name{globalcommonv1.Region_NAME_UNSPECIFIED},
				OrganizationIds: []uint64{0},
				SkuKey:          "TikvGcpX86N2",
			},
		},
	})
	if err != nil {
		return err
	}

	tikvSpuGcpX86N4, err := spuSvc.Create(ctx, &entity.Spu{
		DisplayName:   "TikvGcpX86N4",
		Name:          "TikvGcpX86N4",
		ComponentType: globaldedicatedv1.ComponentType_TIKV,
		CpuCore:       8,
		MemorySizeGi:  16,
		Mappings: []*globaldedicatedv1.Spu_Mapping{
			{
				CloudProvider:   globalcommonv1.CloudProvider_GCP,
				PhysicalRegions: []globalcommonv1.Region_Name{globalcommonv1.Region_NAME_UNSPECIFIED},
				OrganizationIds: []uint64{0},
				SkuKey:          "TikvGcpX86N4",
			},
		},
	})
	if err != nil {
		return err
	}

	// tikv skus.
	_, err = skuRepo.Save(ctx, &entity.Sku{
		Key:           "TikvGcpX86N2",
		SpuId:         tikvSpuGcpX86N2.GetId(),
		CloudProvider: globalcommonv1.CloudProvider_GCP,
		ComponentType: globaldedicatedv1.ComponentType_TIKV,
		InstanceMuId:  uint64(instanceIdGcpX86N2),
		DataStorage: &globaldedicatedv1.Sku_Storage{
			StorageMuId:   uint64(pdssdId),
			SizeGiDefault: 10,
		},
	})
	if err != nil {
		return err
	}

	_, err = skuRepo.Save(ctx, &entity.Sku{
		Key:           "TikvGcpX86N4",
		SpuId:         tikvSpuGcpX86N4.GetId(),
		CloudProvider: globalcommonv1.CloudProvider_GCP,
		ComponentType: globaldedicatedv1.ComponentType_TIKV,
		InstanceMuId:  uint64(instanceIdGcpX86N4),
		DataStorage: &globaldedicatedv1.Sku_Storage{
			StorageMuId:   uint64(pdssdId),
			SizeGiDefault: 10,
		},
	})
	if err != nil {
		return err
	}

	return nil
}