package cluster

import (
	"context"
	"encoding/base64"
	"encoding/binary"
	"encoding/hex"
	"strconv"
	"testing"

	. "github.com/onsi/gomega"
	"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"
	openapiserverlessv1beta1 "github.com/tidbcloud/tidb-management-service/api/spec/tidb_cloud_open_api/serverless/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/adapter/billing"
	"github.com/tidbcloud/tidb-management-service/tidb-mgmt-service/domain/adapter/billing/mock_billing"
	"github.com/tidbcloud/tidb-management-service/tidb-mgmt-service/domain/entity"
	"github.com/tidbcloud/tidb-management-service/tidb-mgmt-service/infra/billingcli"
	"github.com/tidbcloud/tidb-management-service/tidb-mgmt-service/infra/repository"
	"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/bizerr"
	"github.com/tidbcloud/tidb-management-service/utilities/grpc/metadata"
)

func newAlibabaMarketplaceAccountMock(t *testing.T) *mock_account.MockClient {
	t.Helper()

	ctrl := gomock.NewController(t)
	accountMock := mock_account.NewMockClient(ctrl)
	accountMock.EXPECT().
		GetTenantThirdPartyAccountProviders(gomock.Any(), uint64(1)).
		Return([]openapicommonv1beta1.Region_CloudProvider{openapicommonv1beta1.Region_alicloud}, nil).
		AnyTimes()

	return accountMock
}

func newDenyBillingActionCheckMock(t *testing.T, reason string, times int) *mock_billing.MockClient {
	t.Helper()

	ctrl := gomock.NewController(t)
	billingMock := mock_billing.NewMockClient(ctrl)
	billingMock.EXPECT().
		BillingActionCheck(
			gomock.Any(),
			billing.ActionCreate,
			uint64(1),
			gomock.AssignableToTypeOf(&billing.Resource{}),
		).
		DoAndReturn(func(
			ctx context.Context,
			action billing.Action,
			tenantID uint64,
			resource *billing.Resource,
		) (*billing.ActionCheckResult, bizerr.BizError) {
			g := NewWithT(t)
			g.Expect(resource.Type).Should(Equal(billing.ResourceTypeServerlessCluster))
			g.Expect(resource.Provider).ShouldNot(BeNil())
			g.Expect(*resource.Provider).Should(Equal(billing.ProviderAlicloud))
			return &billing.ActionCheckResult{Pass: false, Reason: &reason}, nil
		}).
		Times(times)
	billingMock.EXPECT().GetSpendLimitSettings(gomock.Any(), gomock.Any()).AnyTimes()
	billingMock.EXPECT().SendClusterEventToUNC(gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes()

	return billingMock
}

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

	service, err := initService()
	g.Expect(err).ShouldNot(HaveOccurred())
	dbProvider := db.NewProvider(nil)

	type testCase struct {
		name      string
		inputsFn  func() (context.Context, *openapiserverlessv1beta1.CreateClusterRequest)
		outputsFn func(*openapiserverlessv1beta1.Cluster, error)
	}

	cases := []testCase{
		{
			name: "Create should success",
			inputsFn: func() (context.Context, *openapiserverlessv1beta1.CreateClusterRequest) {
				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 := &openapiserverlessv1beta1.CreateClusterRequest{
					Cluster: &openapiserverlessv1beta1.Cluster{
						DisplayName: "create-test-1",
						Region:      &openapicommonv1beta1.Region{Name: "regions/aws-us-east-1"},
						Labels: map[string]string{
							common.ProjectIDLabelKey: "1",
						},
					},
				}
				return ctx, req
			},
			outputsFn: func(cluster *openapiserverlessv1beta1.Cluster, err error) {
				g.Expect(err).ShouldNot(HaveOccurred())
				// 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.DeployTypeServerless))
				g.Expect(record.CreatedBy).Should(Equal("foo@bar.com"))
				g.Expect(record.Labels).Should(HaveKeyWithValue(metadata.UserIDHeaderKey, "1"))
				g.Expect(record.EngineType).Should(Equal(openapicommonv1beta1.EngineType_MYSQL_COMPATIBLE.String()))
				g.Expect(cluster.EngineType).Should(Equal(openapicommonv1beta1.EngineType_MYSQL_COMPATIBLE))
			},
		},
		{
			name: "Create pg starter cluster should success",
			inputsFn: func() (context.Context, *openapiserverlessv1beta1.CreateClusterRequest) {
				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 := &openapiserverlessv1beta1.CreateClusterRequest{
					Cluster: &openapiserverlessv1beta1.Cluster{
						DisplayName: "create-test-pg-1",
						Region:      &openapicommonv1beta1.Region{Name: "regions/aws-us-east-1"},
						EngineType:  openapicommonv1beta1.EngineType_POSTGRESQL_COMPATIBLE,
						Labels: map[string]string{
							common.ProjectIDLabelKey: "1",
						},
					},
				}
				return ctx, req
			},
			outputsFn: func(cluster *openapiserverlessv1beta1.Cluster, err error) {
				g.Expect(err).ShouldNot(HaveOccurred())
				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.EngineType).Should(Equal(openapicommonv1beta1.EngineType_POSTGRESQL_COMPATIBLE.String()))
				g.Expect(cluster.EngineType).Should(Equal(openapicommonv1beta1.EngineType_POSTGRESQL_COMPATIBLE))
			},
		},
		{
			name: "Create cluster with no pubilc endpoint should success",
			inputsFn: func() (context.Context, *openapiserverlessv1beta1.CreateClusterRequest) {
				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 := &openapiserverlessv1beta1.CreateClusterRequest{
					Cluster: &openapiserverlessv1beta1.Cluster{
						DisplayName: "create-test-2",
						Region:      &openapicommonv1beta1.Region{Name: "regions/aws-us-east-1"},
						Endpoints: &openapiserverlessv1beta1.Cluster_Endpoints{
							Public: &openapiserverlessv1beta1.Cluster_Endpoints_Public{Disabled: true},
						},
						Labels: map[string]string{
							common.ProjectIDLabelKey: "1",
						},
					},
				}
				return ctx, req
			},
			outputsFn: func(cluster *openapiserverlessv1beta1.Cluster, err error) {
				g.Expect(err).ShouldNot(HaveOccurred())
				// 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.DeployTypeServerless))
				g.Expect(record.CreatedBy).Should(Equal("foo@bar.com"))
				g.Expect(record.Labels).Should(HaveKeyWithValue(metadata.UserIDHeaderKey, "1"))
				g.Expect(cluster.Endpoints.Public.Disabled).Should(BeTrue())
			},
		},
		{
			name: "Create cluster with encryption config should success",
			inputsFn: func() (context.Context, *openapiserverlessv1beta1.CreateClusterRequest) {
				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 := &openapiserverlessv1beta1.CreateClusterRequest{
					Cluster: &openapiserverlessv1beta1.Cluster{
						DisplayName: "create-test-3",
						Region:      &openapicommonv1beta1.Region{Name: "regions/aws-us-east-1"},
						Endpoints: &openapiserverlessv1beta1.Cluster_Endpoints{
							Public: &openapiserverlessv1beta1.Cluster_Endpoints_Public{Disabled: true},
						},
						AutoScaling: &openapiserverlessv1beta1.Cluster_AutoScaling{
							MinRcu: 2001,
							MaxRcu: 3001,
						},
						EncryptionConfig: &openapiserverlessv1beta1.Cluster_EncryptionConfig{
							EnhancedEncryptionEnabled: true,
						},
						Labels: map[string]string{
							common.ProjectIDLabelKey: "1",
						},
					},
				}
				return ctx, req
			},
			outputsFn: func(cluster *openapiserverlessv1beta1.Cluster, err error) {
				g.Expect(err).ShouldNot(HaveOccurred())
				// 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-3"))
				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.DeployTypeServerless))
				g.Expect(record.CreatedBy).Should(Equal("foo@bar.com"))
				g.Expect(record.Labels).Should(HaveKeyWithValue(metadata.UserIDHeaderKey, "1"))
				g.Expect(cluster.Endpoints.Public.Disabled).Should(BeTrue())
				g.Expect(cluster.EncryptionConfig.EnhancedEncryptionEnabled).Should(BeTrue())
			},
		},
		{
			name: "Create free clusters failed when free quota is exceeded",
			inputsFn: func() (context.Context, *openapiserverlessv1beta1.CreateClusterRequest) {
				var record schema.Cluster
				// except create-test-1, others should be non-free
				err = dbProvider.ReadDB(t.Context()).
					Model(&schema.Cluster{}).
					Where("display_name = ?", "create-test-1").
					Take(&record).
					Error
				g.Expect(err).ShouldNot(HaveOccurred())
				billingcli.TestSpendLimitSetting.RangeID = record.ClusterID
				freeClusterQuotaConfig.ServerlessFreeQuotaMaxCount = 1
				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 := &openapiserverlessv1beta1.CreateClusterRequest{
					Cluster: &openapiserverlessv1beta1.Cluster{
						DisplayName: "create-test-4",
						Region:      &openapicommonv1beta1.Region{Name: "regions/aws-us-east-1"},
						Labels: map[string]string{
							common.ProjectIDLabelKey: "1",
						},
					},
				}
				return ctx, req
			},
			outputsFn: func(cluster *openapiserverlessv1beta1.Cluster, err error) {
				g.Expect(err).Should(HaveOccurred())
				g.Expect(err.Error()).Should(ContainSubstring("reached the maximum number of free clusters"))
				billingcli.TestSpendLimitSetting.RangeID = 0
				freeClusterQuotaConfig.ServerlessFreeQuotaMaxCount = 100
			},
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			ctx, req := tc.inputsFn()
			cluster, err := service.CreateCluster(ctx, req)
			tc.outputsFn(cluster, err)
		})
	}
}

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

	const denyReason = "billing denied alicloud marketplace cluster creation"
	service, err := initServiceWithAccountAndBillingCli(
		newAlibabaMarketplaceAccountMock(t),
		newDenyBillingActionCheckMock(t, denyReason, 1),
	)
	g.Expect(err).ShouldNot(HaveOccurred())

	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",
	}))

	cluster, err := service.CreateCluster(ctx, &openapiserverlessv1beta1.CreateClusterRequest{
		Cluster: &openapiserverlessv1beta1.Cluster{
			DisplayName: "create-test-alibaba-marketplace-starter",
			Region:      &openapicommonv1beta1.Region{Name: "regions/alicloud-cn-hangzhou"},
			Labels: map[string]string{
				common.ProjectIDLabelKey: "1",
			},
		},
	})

	g.Expect(cluster).Should(BeNil())
	g.Expect(err).Should(HaveOccurred())
	g.Expect(err.Error()).Should(ContainSubstring(denyReason))
}

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

	const denyReason = "billing denied batch alicloud marketplace cluster creation"
	service, err := initServiceWithAccountAndBillingCli(
		newAlibabaMarketplaceAccountMock(t),
		newDenyBillingActionCheckMock(t, denyReason, 1),
	)
	g.Expect(err).ShouldNot(HaveOccurred())

	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",
	}))

	resp, err := service.BatchCreateClusters(ctx, &openapiserverlessv1beta1.BatchCreateClustersRequest{
		Requests: []*openapiserverlessv1beta1.CreateClusterRequest{
			{
				Cluster: &openapiserverlessv1beta1.Cluster{
					DisplayName: "batch-create-test-alibaba-marketplace-starter",
					Region:      &openapicommonv1beta1.Region{Name: "regions/alicloud-cn-hangzhou"},
					SpendingLimit: &openapiserverlessv1beta1.Cluster_SpendingLimit{
						Monthly: 1,
					},
					Labels: map[string]string{
						common.ProjectIDLabelKey: "1",
					},
				},
			},
		},
	})

	g.Expect(resp).Should(BeNil())
	g.Expect(err).Should(HaveOccurred())
	g.Expect(err.Error()).Should(ContainSubstring(denyReason))
}

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

	const denyReason = "billing denied claim alicloud marketplace cluster creation"
	service, err := initServiceWithAccountAndBillingCli(
		newAlibabaMarketplaceAccountMock(t),
		newDenyBillingActionCheckMock(t, denyReason, 1),
	)
	g.Expect(err).ShouldNot(HaveOccurred())

	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",
	}))

	dbProvider := db.NewProvider(nil)
	clusterRepo := repository.NewClusterRepository(dbProvider)
	zeroRentRecordRepo := repository.NewZeroRentRecordRepository(dbProvider)

	zeroCluster := &entity.Cluster{
		ClusterId:      entity.GenClusterID(),
		ProjectId:      1,
		OrganizationId: 1,
		DisplayName:    "rented-shadow-cluster",
		DeployType:     entity.DeployTypeServerless,
		CloudProvider:  openapicommonv1beta1.Region_aws,
		Region:         "us-east-1",
		UserId:         "1",
		CreatedBy:      "foo@bar.com",
		Labels: map[string]string{
			common.ProjectIDLabelKey: "1",
			metadata.UserIDHeaderKey: "1",
		},
		ServicePlan: openapicommonv1beta1.ServicePlan_Starter,
		ShadowState: entity.ShadowStateRented,
	}
	_, err = clusterRepo.InsertNewCluster(ctx, zeroCluster)
	g.Expect(err).ShouldNot(HaveOccurred())

	secretBytes := [16]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}
	err = zeroRentRecordRepo.Upsert(ctx, zeroCluster.ClusterId, hex.EncodeToString(secretBytes[:]), "127.0.0.1", "test")
	g.Expect(err).ShouldNot(HaveOccurred())

	var zeroIDPayload [24]byte
	binary.BigEndian.PutUint64(zeroIDPayload[:8], zeroCluster.ClusterId)
	copy(zeroIDPayload[8:], secretBytes[:])
	zeroID := base64.RawURLEncoding.EncodeToString(zeroIDPayload[:])

	claimedCluster, err := service.ClaimClusterFromZero(ctx, &openapiserverlessv1beta1.ClaimClusterFromZeroRequest{
		ZeroId: zeroID,
		Cluster: &openapiserverlessv1beta1.Cluster{
			DisplayName: "claim-test-alibaba-marketplace-starter",
			Region:      &openapicommonv1beta1.Region{Name: "regions/alicloud-cn-hangzhou"},
			Labels: map[string]string{
				common.ProjectIDLabelKey: strconv.FormatUint(zeroCluster.ProjectId, 10),
			},
		},
	})

	g.Expect(claimedCluster).Should(BeNil())
	g.Expect(err).Should(HaveOccurred())
	g.Expect(err.Error()).Should(ContainSubstring(denyReason))

	var count int64
	err = dbProvider.ReadDB(t.Context()).
		Model(&schema.Cluster{}).
		Where("display_name = ?", "claim-test-alibaba-marketplace-starter").
		Count(&count).
		Error
	g.Expect(err).ShouldNot(HaveOccurred())
	g.Expect(count).Should(Equal(int64(0)))
}