package cluster

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

func TestScaleClusterN2N4Coexist(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, scale out 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)

		// Mimic the cluster is in available state
		resCluster.State = globaldedicatedv1.Cluster_AVAILABLE
		err = svc.clusterRepo.Save(t.Context(), resCluster)
		require.NoError(t, err)

		scaleReq := &globaldedicatedv1.ScaleClusterWithMultiComponentRequest{
			ClusterId: resCluster.GetClusterId(),
			ClusterComponents: []*globaldedicatedv1.Cluster_ClusterComponent{
				{
					ComponentType: globaldedicatedv1.ComponentType_TIKV,
					Count:         lo.ToPtr(uint32(6)),
				},
			},
		}
		scaledCluster, err := svc.Scale(t.Context(), scaleReq)
		require.NoError(t, err)

		pdComponent, _ := lo.Find(scaledCluster.ClusterComponents, func(c *globaldedicatedv1.Cluster_ClusterComponent) bool {
			return c.ComponentType == 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, scale out 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)

		// Mimic the cluster is in available state
		resCluster.State = globaldedicatedv1.Cluster_AVAILABLE
		err = svc.clusterRepo.Save(t.Context(), resCluster)
		require.NoError(t, err)

		scaleReq := &globaldedicatedv1.ScaleClusterWithMultiComponentRequest{
			ClusterId: resCluster.GetClusterId(),
			ClusterComponents: []*globaldedicatedv1.Cluster_ClusterComponent{
				{
					ComponentType: globaldedicatedv1.ComponentType_TIKV,
					Count:         lo.ToPtr(uint32(6)),
				},
			},
		}
		scaledCluster, err := svc.Scale(t.Context(), scaleReq)
		require.NoError(t, err)

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