package cluster

import (
	"context"
	"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/domain/entity"
	"github.com/tidbcloud/tidb-management-service/tidb-mgmt-service/infra/repository/db"
	"github.com/tidbcloud/tidb-management-service/tidb-mgmt-service/infra/repository/schema"
	"github.com/tidbcloud/tidb-management-service/tidb-mgmt-service/pkg/common"
	"github.com/tidbcloud/tidb-management-service/utilities/grpc/metadata"
)

func TestCreateDedicatedCluster(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())
	dbProvider := db.NewProvider(nil)

	t.Run("Create should success", func(t *testing.T) {
		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.UserEmailHeaderKey:        "foo@bar.com",
			metadata.IsCheckPermissionKey:      "true",
		}))
		req := &openapidedicatedv1beta1.CreateClusterRequest{
			Cluster: &openapidedicatedv1beta1.Cluster{
				DisplayName: "create-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",
				},
			},
		}

		cluster, err := service.CreateCluster(ctx, req)
		g.Expect(err).ShouldNot(HaveOccurred())

		time.Sleep(3 * time.Second) // wait for the workflow dag to finish

		// Check the cluster in database
		var record schema.Cluster
		err = dbProvider.ReadDB(t.Context()).
			Model(&schema.Cluster{}).
			Where("cluster_id = ?", cluster.ClusterId).
			Take(&record).
			Error
		g.Expect(err).ShouldNot(HaveOccurred())
		g.Expect(record).ShouldNot(BeNil())
		g.Expect(record.DisplayName).Should(Equal("create-test-1"))
		g.Expect(record.OrganizationID).Should(Equal(uint64(1)))
		g.Expect(record.ProjectID).Should(Equal(uint64(1)))
		g.Expect(record.ClusterID).ShouldNot(Equal(0))
		g.Expect(record.DeployType).Should(Equal(entity.DeployTypeDedicated))
		g.Expect(record.CreatedBy).Should(Equal("foo@bar.com"))
		g.Expect(record.Labels).Should(HaveKeyWithValue(metadata.UserIDHeaderKey, "1"))
		g.Expect(cluster.GetTikvNodeSetting().GetStorageType()).
			Should(Equal(openapidedicatedv1beta1.Cluster_StorageNodeSetting_Standard))
	})

	t.Run("Create BASIC v2 tikv should success", func(t *testing.T) {
		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.UserEmailHeaderKey:        "foo@bar.com",
			metadata.IsCheckPermissionKey:      "true",
		}))
		req := &openapidedicatedv1beta1.CreateClusterRequest{
			Cluster: &openapidedicatedv1beta1.Cluster{
				DisplayName: "create-test-2",
				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",
				},
			},
		}

		cluster, err := service.CreateCluster(ctx, req)
		g.Expect(err).ShouldNot(HaveOccurred())

		time.Sleep(3 * time.Second) // wait for the workflow dag to finish

		// Check the cluster in database
		var record schema.Cluster
		err = dbProvider.ReadDB(t.Context()).
			Model(&schema.Cluster{}).
			Where("cluster_id = ?", cluster.ClusterId).
			Take(&record).
			Error
		g.Expect(err).ShouldNot(HaveOccurred())
		g.Expect(record).ShouldNot(BeNil())
		g.Expect(record.DisplayName).Should(Equal("create-test-2"))
		g.Expect(record.OrganizationID).Should(Equal(uint64(1)))
		g.Expect(record.ProjectID).Should(Equal(uint64(1)))
		g.Expect(record.ClusterID).ShouldNot(Equal(0))
		g.Expect(record.DeployType).Should(Equal(entity.DeployTypeDedicated))
		g.Expect(record.CreatedBy).Should(Equal("foo@bar.com"))
		g.Expect(record.Labels).Should(HaveKeyWithValue(metadata.UserIDHeaderKey, "1"))
		g.Expect(cluster.GetTikvNodeSetting().GetStorageType()).
			Should(Equal(openapidedicatedv1beta1.Cluster_StorageNodeSetting_Standard))
	})

	t.Run("Create cluster with tiproxy should success", func(t *testing.T) {
		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.UserEmailHeaderKey:        "foo@bar.com",
			metadata.IsCheckPermissionKey:      "true",
		}))
		req := &openapidedicatedv1beta1.CreateClusterRequest{
			Cluster: &openapidedicatedv1beta1.Cluster{
				DisplayName: "create-test-tiproxy",
				RegionId:    "aws-us-west-2",
				TidbNodeSetting: &openapidedicatedv1beta1.Cluster_TidbNodeSetting{
					NodeSpecKey: "8C16G",
					TidbNodeGroups: []*openapidedicatedv1beta1.TidbNodeGroup{
						{
							NodeCount: 2,
							TiproxySetting: &openapidedicatedv1beta1.TidbNodeGroup_TiProxySetting{
								NodeSpecKey: "2C4G",
								NodeCount:   lo.ToPtr(int32(2)),
							},
						},
					},
				},
				TikvNodeSetting: &openapidedicatedv1beta1.Cluster_StorageNodeSetting{
					NodeSpecKey:   "8C16G",
					NodeCount:     3,
					StorageSizeGi: 50,
					StorageType:   openapidedicatedv1beta1.Cluster_StorageNodeSetting_Standard,
				},
				Labels: map[string]string{
					common.ProjectIDLabelKey: "1",
				},
			},
		}

		cluster, err := service.CreateCluster(ctx, req)
		g.Expect(err).ShouldNot(HaveOccurred())

		time.Sleep(3 * time.Second) // wait for the workflow dag to finish

		// Check the cluster in database
		var record schema.Cluster
		err = dbProvider.ReadDB(t.Context()).
			Model(&schema.Cluster{}).
			Where("cluster_id = ?", cluster.ClusterId).
			Take(&record).
			Error
		g.Expect(err).ShouldNot(HaveOccurred())
		g.Expect(record).ShouldNot(BeNil())
		g.Expect(record.DisplayName).Should(Equal("create-test-tiproxy"))
		g.Expect(len(cluster.GetTidbNodeSetting().GetTidbNodeGroups())).Should(Equal(1))
		g.Expect(cluster.GetTidbNodeSetting().GetTidbNodeGroups()[0].GetTiproxySetting().GetNodeCount()).
			Should(Equal(int32(2)))
		g.Expect(cluster.GetTidbNodeSetting().GetTidbNodeGroups()[0].GetTiproxySetting().GetNodeSpecKey()).
			Should(Equal("2C4G"))
	})
}