package cluster

import (
	"cmp"
	"context"
	"slices"
	"testing"
	"time"

	. "github.com/onsi/gomega"
	"github.com/samber/lo"
	accountsdk "github.com/tidbcloud/account/idl/pbgen/proto/account"
	"go.uber.org/mock/gomock"
	grpcmetadata "google.golang.org/grpc/metadata"

	openapicommonv1beta1 "github.com/tidbcloud/tidb-management-service/api/spec/tidb_cloud_open_api/common/v1beta1"
	openapidedicatedv1beta1 "github.com/tidbcloud/tidb-management-service/api/spec/tidb_cloud_open_api/dedicated/v1beta1"
	"github.com/tidbcloud/tidb-management-service/tidb-mgmt-service/domain/adapter/account/mock_account"
	"github.com/tidbcloud/tidb-management-service/tidb-mgmt-service/pkg/common"
	"github.com/tidbcloud/tidb-management-service/utilities/grpc/metadata"
)

func TestScaleDedicatedCluster(t *testing.T) {
	g := NewWithT(t)

	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mockAccountCli := mock_account.NewMockClient(ctrl)
	mockAccountCli.EXPECT().
		GetProject(gomock.Any(), gomock.Any()).
		DoAndReturn(func(ctx context.Context, projectId uint64) (*accountsdk.Project, error) {
			return &accountsdk.Project{
				Id:       projectId,
				TenantId: 1, // org id
			}, nil
		}).
		AnyTimes()
	mockAccountCli.EXPECT().
		GetTenantThirdPartyAccountProviders(gomock.Any(), gomock.Any()).
		DoAndReturn(func(ctx context.Context, id uint64) ([]openapicommonv1beta1.Region_CloudProvider, error) {
			return []openapicommonv1beta1.Region_CloudProvider{}, nil
		}).
		AnyTimes()
	mockAccountCli.EXPECT().
		UpsertClusterRelations(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
		Return(nil).
		AnyTimes()

	service, err := initServiceWithFakeClients(mockAccountCli)
	g.Expect(err).ShouldNot(HaveOccurred())

	// Create a cluster first
	ctx := grpcmetadata.NewIncomingContext(t.Context(), grpcmetadata.New(map[string]string{
		metadata.UserIDHeaderKey:           "1",
		metadata.OrgIDHeaderKey:            "1",
		metadata.DefaultProjectIDHeaderKey: "1",
		metadata.ProjectListHeaderKey:      "1,2,3",
		metadata.IsCheckPermissionKey:      "true",
	}))

	req := &openapidedicatedv1beta1.CreateClusterRequest{
		Cluster: &openapidedicatedv1beta1.Cluster{
			DisplayName: "scale-test-1",
			RegionId:    "aws-us-west-2",
			TidbNodeSetting: &openapidedicatedv1beta1.Cluster_TidbNodeSetting{
				NodeSpecKey: "8C16G",
				TidbNodeGroups: []*openapidedicatedv1beta1.TidbNodeGroup{
					{
						NodeCount: 1,
					},
				},
			},
			TikvNodeSetting: &openapidedicatedv1beta1.Cluster_StorageNodeSetting{
				NodeSpecKey:   "8C16G",
				NodeCount:     3,
				StorageSizeGi: 50,
				StorageType:   openapidedicatedv1beta1.Cluster_StorageNodeSetting_Standard,
			},
			Labels: map[string]string{
				common.ProjectIDLabelKey: "1",
			},
		},
	}
	c, err := service.CreateCluster(ctx, req)
	g.Expect(err).ShouldNot(HaveOccurred())
	g.Expect(c.TidbNodeSetting).ShouldNot(BeNil())
	g.Expect(c.TidbNodeSetting.TidbNodeGroups).ShouldNot(BeNil())
	g.Expect(c.TidbNodeSetting.TidbNodeGroups).Should(HaveLen(1))
	g.Expect(c.TidbNodeSetting.GetNodeSpecKey()).Should(Equal("8C16G"))
	g.Expect(c.TidbNodeSetting.TidbNodeGroups[0].GetNodeCount()).Should(Equal(int32(1)))
	g.Expect(c.TikvNodeSetting).ShouldNot(BeNil())
	g.Expect(c.TikvNodeSetting.GetNodeSpecKey()).Should(Equal("8C16G"))
	g.Expect(c.TikvNodeSetting.GetNodeCount()).Should(Equal(int32(3)))

	time.Sleep(3 * time.Second) // wait for the workflow dag to finish
	clusterID := c.ClusterId

	t.Run("Scale should success", func(t *testing.T) {
		cluster, err := service.UpdateCluster(ctx, &openapidedicatedv1beta1.UpdateClusterRequest{
			Cluster: &openapidedicatedv1beta1.UpdateClusterRequest_Cluster{
				ClusterId: clusterID,

				TikvNodeSetting: &openapidedicatedv1beta1.UpdateClusterRequest_StorageNodeSetting{
					NodeSpecKey:   "16C32G",
					NodeCount:     lo.ToPtr(int32(3)),
					StorageSizeGi: int32(50),
				},
			},
		})
		g.Expect(err).ShouldNot(HaveOccurred())
		g.Expect(cluster).ShouldNot(BeNil())
		g.Expect(cluster.ClusterId).Should(Equal(clusterID))
		g.Expect(cluster.DisplayName).Should(Equal("scale-test-1"))
		g.Expect(cluster.State).Should(Equal(openapicommonv1beta1.ClusterState_MODIFYING))
		g.Expect(cluster.GetRegionDisplayName()).Should(Equal(fakeRegionDisplayName(cluster.GetRegionId())))

		// Wait for the cluster to be scale done
		time.Sleep(3 * time.Second)

		// Get the cluster again
		cluster, err = service.GetCluster(ctx, &openapidedicatedv1beta1.GetClusterRequest{ClusterId: clusterID})
		g.Expect(err).ShouldNot(HaveOccurred())
		g.Expect(cluster).ShouldNot(BeNil())
		g.Expect(cluster.ClusterId).Should(Equal(clusterID))
		g.Expect(cluster.DisplayName).Should(Equal("scale-test-1"))
		g.Expect(cluster.State).Should(Equal(openapicommonv1beta1.ClusterState_ACTIVE))
		g.Expect(cluster.TidbNodeSetting.TidbNodeGroups).ShouldNot(BeNil())
		g.Expect(cluster.TidbNodeSetting.TidbNodeGroups).Should(HaveLen(1))
		g.Expect(cluster.TidbNodeSetting.GetNodeSpecKey()).Should(Equal("8C16G"))
		g.Expect(cluster.TidbNodeSetting.TidbNodeGroups[0].GetNodeCount()).Should(Equal(int32(1)))
		g.Expect(cluster.TikvNodeSetting).ShouldNot(BeNil())
		g.Expect(cluster.TikvNodeSetting.GetNodeSpecKey()).Should(Equal("16C32G"))
		g.Expect(cluster.TikvNodeSetting.GetNodeCount()).Should(Equal(int32(3)))
	})

	t.Run("Scale 2 should success", func(t *testing.T) {
		cluster, err := service.UpdateCluster(ctx, &openapidedicatedv1beta1.UpdateClusterRequest{
			Cluster: &openapidedicatedv1beta1.UpdateClusterRequest_Cluster{
				ClusterId: clusterID,

				TidbNodeSetting: &openapidedicatedv1beta1.UpdateClusterRequest_TidbNodeSetting{
					NodeSpecKey: "8C16G",
					TidbNodeGroups: []*openapidedicatedv1beta1.UpdateClusterRequest_TidbNodeSetting_TidbNodeGroup{
						{
							NodeCount: lo.ToPtr(int32(3)),
						},
					},
				},
			},
		})
		g.Expect(err).ShouldNot(HaveOccurred())
		g.Expect(cluster).ShouldNot(BeNil())
		g.Expect(cluster.ClusterId).Should(Equal(clusterID))
		g.Expect(cluster.DisplayName).Should(Equal("scale-test-1"))
		g.Expect(cluster.State).Should(Equal(openapicommonv1beta1.ClusterState_MODIFYING))
		g.Expect(cluster.GetRegionDisplayName()).Should(Equal(fakeRegionDisplayName(cluster.GetRegionId())))

		// Wait for the cluster to be scale done
		time.Sleep(3 * time.Second)

		// Get the cluster again
		cluster, err = service.GetCluster(ctx, &openapidedicatedv1beta1.GetClusterRequest{ClusterId: clusterID})
		g.Expect(err).ShouldNot(HaveOccurred())
		g.Expect(cluster).ShouldNot(BeNil())
		g.Expect(cluster.ClusterId).Should(Equal(clusterID))
		g.Expect(cluster.DisplayName).Should(Equal("scale-test-1"))
		g.Expect(cluster.State).Should(Equal(openapicommonv1beta1.ClusterState_ACTIVE))
		g.Expect(cluster.TidbNodeSetting.TidbNodeGroups).ShouldNot(BeNil())
		g.Expect(cluster.TidbNodeSetting.TidbNodeGroups).Should(HaveLen(1))
		g.Expect(cluster.TidbNodeSetting.GetNodeSpecKey()).Should(Equal("8C16G"))
		g.Expect(cluster.TidbNodeSetting.TidbNodeGroups[0].GetNodeCount()).Should(Equal(int32(3)))
		g.Expect(cluster.TikvNodeSetting).ShouldNot(BeNil())
		g.Expect(cluster.TikvNodeSetting.GetNodeSpecKey()).Should(Equal("16C32G"))
		g.Expect(cluster.TikvNodeSetting.GetNodeCount()).Should(Equal(int32(3)))
	})

	t.Run("Scale tiproxy should success", func(t *testing.T) {
		cluster, err := service.UpdateCluster(ctx, &openapidedicatedv1beta1.UpdateClusterRequest{
			Cluster: &openapidedicatedv1beta1.UpdateClusterRequest_Cluster{
				ClusterId: clusterID,
				TidbNodeSetting: &openapidedicatedv1beta1.UpdateClusterRequest_TidbNodeSetting{
					NodeSpecKey: "8C16G",
					TidbNodeGroups: []*openapidedicatedv1beta1.UpdateClusterRequest_TidbNodeSetting_TidbNodeGroup{
						{
							NodeCount: lo.ToPtr(int32(3)),
							TiproxySetting: &openapidedicatedv1beta1.TidbNodeGroup_TiProxySetting{
								NodeSpecKey: "2C4G",
								NodeCount:   lo.ToPtr(int32(2)),
							},
						},
					},
				},
			},
		})
		g.Expect(err).ShouldNot(HaveOccurred())
		g.Expect(cluster).ShouldNot(BeNil())

		// Wait for the cluster to be scale done
		time.Sleep(3 * time.Second)

		// Get the cluster again
		cluster, err = service.GetCluster(ctx, &openapidedicatedv1beta1.GetClusterRequest{ClusterId: clusterID})
		g.Expect(err).ShouldNot(HaveOccurred())
		g.Expect(cluster).ShouldNot(BeNil())
		g.Expect(cluster.State).Should(Equal(openapicommonv1beta1.ClusterState_ACTIVE))
		g.Expect(cluster.TidbNodeSetting.TidbNodeGroups[0].TiproxySetting).ShouldNot(BeNil())
		g.Expect(cluster.TidbNodeSetting.TidbNodeGroups[0].TiproxySetting.GetNodeSpecKey()).
			Should(Equal("2C4G"))
		g.Expect(cluster.TidbNodeSetting.TidbNodeGroups[0].TiproxySetting.GetNodeCount()).
			Should(Equal(int32(2)))
	})

	t.Run("TiProxy type is not allowed to be updated", func(t *testing.T) {
		_, err := service.UpdateCluster(ctx, &openapidedicatedv1beta1.UpdateClusterRequest{
			Cluster: &openapidedicatedv1beta1.UpdateClusterRequest_Cluster{
				ClusterId: clusterID,
				TidbNodeSetting: &openapidedicatedv1beta1.UpdateClusterRequest_TidbNodeSetting{
					NodeSpecKey: "8C16G",
					TidbNodeGroups: []*openapidedicatedv1beta1.UpdateClusterRequest_TidbNodeSetting_TidbNodeGroup{
						{
							NodeCount: lo.ToPtr(int32(3)),
							TiproxySetting: &openapidedicatedv1beta1.TidbNodeGroup_TiProxySetting{
								NodeSpecKey: "8C16G",
								NodeCount:   lo.ToPtr(int32(2)),
							},
						},
					},
				},
			},
		})
		g.Expect(err).Should(HaveOccurred())
		g.Expect(err.Error()).Should(ContainSubstring("tiproxy node type cannot be updated"))
	})

	t.Run("Scale tiproxy to 0 should success", func(t *testing.T) {
		cluster, err := service.UpdateCluster(ctx, &openapidedicatedv1beta1.UpdateClusterRequest{
			Cluster: &openapidedicatedv1beta1.UpdateClusterRequest_Cluster{
				ClusterId: clusterID,
				TidbNodeSetting: &openapidedicatedv1beta1.UpdateClusterRequest_TidbNodeSetting{
					NodeSpecKey: "8C16G",
					TidbNodeGroups: []*openapidedicatedv1beta1.UpdateClusterRequest_TidbNodeSetting_TidbNodeGroup{
						{
							NodeCount: lo.ToPtr(int32(3)),
							TiproxySetting: &openapidedicatedv1beta1.TidbNodeGroup_TiProxySetting{
								NodeCount: lo.ToPtr(int32(0)),
							},
						},
					},
				},
			},
		})
		g.Expect(err).ShouldNot(HaveOccurred())
		g.Expect(cluster).ShouldNot(BeNil())

		// Wait for the cluster to be scale done
		time.Sleep(3 * time.Second)

		// Get the cluster again
		cluster, err = service.GetCluster(ctx, &openapidedicatedv1beta1.GetClusterRequest{ClusterId: clusterID})
		g.Expect(err).ShouldNot(HaveOccurred())
		g.Expect(cluster).ShouldNot(BeNil())
		g.Expect(cluster.State).Should(Equal(openapicommonv1beta1.ClusterState_ACTIVE))
		g.Expect(cluster.TidbNodeSetting.TidbNodeGroups[0].GetTiproxySetting().GetNodeCount()).
			Should(Equal(int32(0)))
	})

	t.Run("TiProxy type should be required", func(t *testing.T) {
		_, err := service.UpdateCluster(ctx, &openapidedicatedv1beta1.UpdateClusterRequest{
			Cluster: &openapidedicatedv1beta1.UpdateClusterRequest_Cluster{
				ClusterId: clusterID,
				TidbNodeSetting: &openapidedicatedv1beta1.UpdateClusterRequest_TidbNodeSetting{
					NodeSpecKey: "8C16G",
					TidbNodeGroups: []*openapidedicatedv1beta1.UpdateClusterRequest_TidbNodeSetting_TidbNodeGroup{
						{
							NodeCount: lo.ToPtr(int32(3)),
							TiproxySetting: &openapidedicatedv1beta1.TidbNodeGroup_TiProxySetting{
								NodeCount: lo.ToPtr(int32(2)),
							},
						},
					},
				},
			},
		})
		g.Expect(err).Should(HaveOccurred())
		g.Expect(err.Error()).Should(ContainSubstring("tiproxy node type cannot be empty"))
	})

	t.Run("Multiple tiproxy should have same type", func(t *testing.T) {
		_, err := service.UpdateCluster(ctx, &openapidedicatedv1beta1.UpdateClusterRequest{
			Cluster: &openapidedicatedv1beta1.UpdateClusterRequest_Cluster{
				ClusterId: clusterID,
				TidbNodeSetting: &openapidedicatedv1beta1.UpdateClusterRequest_TidbNodeSetting{
					NodeSpecKey: "8C16G",
					TidbNodeGroups: []*openapidedicatedv1beta1.UpdateClusterRequest_TidbNodeSetting_TidbNodeGroup{
						{
							NodeCount: lo.ToPtr(int32(3)),
							TiproxySetting: &openapidedicatedv1beta1.TidbNodeGroup_TiProxySetting{
								NodeSpecKey: "2C4G",
								NodeCount:   lo.ToPtr(int32(2)),
							},
						},
						{
							NodeCount:       lo.ToPtr(int32(2)),
							TidbNodeGroupId: "2",
							TiproxySetting: &openapidedicatedv1beta1.TidbNodeGroup_TiProxySetting{
								NodeSpecKey: "8C16G",
								NodeCount:   lo.ToPtr(int32(1)),
							},
						},
					},
				},
			},
		})
		g.Expect(err).Should(HaveOccurred())
		g.Expect(err.Error()).
			Should(ContainSubstring("multi tiproxy group is not allowed set different node spec"))
	})

	t.Run("multiple tiproxy groups should success", func(t *testing.T) {
		cluster, err := service.UpdateCluster(ctx, &openapidedicatedv1beta1.UpdateClusterRequest{
			Cluster: &openapidedicatedv1beta1.UpdateClusterRequest_Cluster{
				ClusterId: clusterID,
				TidbNodeSetting: &openapidedicatedv1beta1.UpdateClusterRequest_TidbNodeSetting{
					NodeSpecKey: "8C16G",
					TidbNodeGroups: []*openapidedicatedv1beta1.UpdateClusterRequest_TidbNodeSetting_TidbNodeGroup{
						{
							NodeCount: lo.ToPtr(int32(3)),
							TiproxySetting: &openapidedicatedv1beta1.TidbNodeGroup_TiProxySetting{
								NodeSpecKey: "2C4G",
								NodeCount:   lo.ToPtr(int32(2)),
							},
						},
						{
							NodeCount:       lo.ToPtr(int32(2)),
							TidbNodeGroupId: "2",
							TiproxySetting: &openapidedicatedv1beta1.TidbNodeGroup_TiProxySetting{
								NodeSpecKey: "2C4G",
								NodeCount:   lo.ToPtr(int32(1)),
							},
						},
					},
				},
			},
		})
		g.Expect(err).ShouldNot(HaveOccurred())
		g.Expect(cluster).ShouldNot(BeNil())

		// Wait for the cluster to be scale done
		time.Sleep(3 * time.Second)

		// Get the cluster again
		cluster, err = service.GetCluster(ctx, &openapidedicatedv1beta1.GetClusterRequest{ClusterId: clusterID})
		g.Expect(err).ShouldNot(HaveOccurred())
		g.Expect(cluster).ShouldNot(BeNil())
		g.Expect(cluster.State).Should(Equal(openapicommonv1beta1.ClusterState_ACTIVE))
		g.Expect(cluster.TidbNodeSetting.TidbNodeGroups).ShouldNot(BeNil())
		g.Expect(cluster.TidbNodeSetting.TidbNodeGroups).Should(HaveLen(2))
		g.Expect(cluster.TidbNodeSetting.GetNodeSpecKey()).Should(Equal("8C16G"))

		tidbGroups := cluster.TidbNodeSetting.TidbNodeGroups
		slices.SortFunc(tidbGroups, func(a, b *openapidedicatedv1beta1.TidbNodeGroup) int {
			return cmp.Compare(a.GetTidbNodeGroupId(), b.GetTidbNodeGroupId())
		})

		g.Expect(tidbGroups[0].GetNodeCount()).Should(Equal(int32(3)))
		g.Expect(tidbGroups[0].GetTiproxySetting().GetNodeSpecKey()).
			Should(Equal("2C4G"))
		g.Expect(tidbGroups[0].GetTiproxySetting().GetNodeCount()).Should(Equal(int32(2)))
		g.Expect(tidbGroups[1].GetNodeCount()).Should(Equal(int32(2)))
		g.Expect(tidbGroups[1].GetTiproxySetting().GetNodeSpecKey()).
			Should(Equal("2C4G"))
		g.Expect(tidbGroups[1].GetTiproxySetting().GetNodeCount()).Should(Equal(int32(1)))
	})

	t.Run("TiProxy type cannot be updated", func(t *testing.T) {
		_, err := service.UpdateCluster(ctx, &openapidedicatedv1beta1.UpdateClusterRequest{
			Cluster: &openapidedicatedv1beta1.UpdateClusterRequest_Cluster{
				ClusterId: clusterID,
				TidbNodeSetting: &openapidedicatedv1beta1.UpdateClusterRequest_TidbNodeSetting{
					NodeSpecKey: "8C16G",
					TidbNodeGroups: []*openapidedicatedv1beta1.UpdateClusterRequest_TidbNodeSetting_TidbNodeGroup{
						{
							NodeCount: lo.ToPtr(int32(3)),
							TiproxySetting: &openapidedicatedv1beta1.TidbNodeGroup_TiProxySetting{
								NodeSpecKey: "8C16G",
								NodeCount:   lo.ToPtr(int32(2)),
							},
						},
					},
				},
			},
		})
		g.Expect(err).Should(HaveOccurred())
		g.Expect(err.Error()).Should(ContainSubstring("tiproxy node type cannot be updated"))
	})
}