package encryptionkey

import (
	"context"
	"strconv"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	clusterngv1beta "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	"go.uber.org/mock/gomock"
	"gorm.io/gorm/clause"

	globalcommonv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/common/v1"
	globalngv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/nextgen/v1"
	regionalserveradapter "github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/adapter/regionalserver"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/entity"
	encryptionkeyrepo "github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/repository/encryptionkey"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/infra/repoimpl/schema"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
	"github.com/tidbcloud/tidb-management-service/utilities/idgenerator"
)

func TestGetCustomerManagedEncryptionKeyByClusterId(t *testing.T) {
	ctx := t.Context()
	sfid := idgenerator.NewSonyflakeID(1, "")

	t.Run("HasCMEK", func(t *testing.T) {
		ts := newTestService(t)
		key := createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid)
		cluster := createTestCluster(ctx, t, ts.clusterRepo, sfid, func(c *entity.Cluster) {
			c.EncryptionKeyId = key.Id
		})

		result, bizErr := ts.svc.GetCustomerManagedEncryptionKeyByClusterId(ctx,
			&globalngv1.GetCustomerManagedEncryptionKeyByClusterIdRequest{
				ClusterId: cluster.GetClusterId(),
			})
		require.Nil(t, bizErr)
		assert.Equal(t, key.Id, result.GetCustomerManagedEncryptionKeyId())
		assert.Equal(t, globalcommonv1.CloudProvider_AWS, result.GetCloudProviderId())
		assert.Equal(t, "arn:aws:kms:us-west-2:123456789012:key/test-key-id", result.GetAwsKms().GetKmsKeyArn())
	})

	t.Run("HasDefaultKey_NotFound", func(t *testing.T) {
		ts := newTestService(t)
		defaultKey := createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid, func(k *entity.EncryptionKey) {
			k.CustomerManaged = false
			k.MasterKey = "default_key"
		})
		cluster := createTestCluster(ctx, t, ts.clusterRepo, sfid, func(c *entity.Cluster) {
			c.EncryptionKeyId = defaultKey.Id
		})

		_, bizErr := ts.svc.GetCustomerManagedEncryptionKeyByClusterId(ctx,
			&globalngv1.GetCustomerManagedEncryptionKeyByClusterIdRequest{
				ClusterId: cluster.GetClusterId(),
			})
		require.NotNil(t, bizErr)
		assert.Contains(t, bizErr.Error(), "does not use a customer-managed")
	})

	t.Run("NoEncryption_NotFound", func(t *testing.T) {
		ts := newTestService(t)
		cluster := createTestCluster(ctx, t, ts.clusterRepo, sfid)

		_, bizErr := ts.svc.GetCustomerManagedEncryptionKeyByClusterId(ctx,
			&globalngv1.GetCustomerManagedEncryptionKeyByClusterIdRequest{
				ClusterId: cluster.GetClusterId(),
			})
		require.NotNil(t, bizErr)
		assert.Contains(t, bizErr.Error(), "no encryption key")
	})
}

func TestListCustomerManagedEncryptionKeys(t *testing.T) {
	ctx := t.Context()
	sfid := idgenerator.NewSonyflakeID(1, "")

	t.Run("MatchingFilter", func(t *testing.T) {
		ts := newTestService(t)

		// Create 2 CMEKs and 1 default key
		cmek1 := createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid)
		cmek2 := createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid, func(k *entity.EncryptionKey) {
			k.MasterKey = "arn:aws:kms:us-west-2:123456789012:key/second-key"
		})
		createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid, func(k *entity.EncryptionKey) {
			k.CustomerManaged = false
			k.MasterKey = "default_key"
		})

		resp, bizErr := ts.svc.ListCustomerManagedEncryptionKeys(ctx,
			&globalngv1.ListCustomerManagedEncryptionKeysRequest{
				OrganizationId: testOrgID,
			})
		require.Nil(t, bizErr)
		require.GreaterOrEqual(t, len(resp.GetCustomerManagedEncryptionKeys()), 2)

		// Verify both CMEKs are in the result
		ids := make(map[uint64]bool)
		for _, c := range resp.GetCustomerManagedEncryptionKeys() {
			ids[c.GetCustomerManagedEncryptionKeyId()] = true
		}
		assert.True(t, ids[cmek1.Id], "cmek1 should be listed")
		assert.True(t, ids[cmek2.Id], "cmek2 should be listed")
	})

	t.Run("FilterByRegion", func(t *testing.T) {
		ts := newTestService(t)

		createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid)
		createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid, func(k *entity.EncryptionKey) {
			k.RegionName = globalcommonv1.Region_AWS_US_EAST_1
			k.RegionId = 4
		})

		resp, bizErr := ts.svc.ListCustomerManagedEncryptionKeys(ctx,
			&globalngv1.ListCustomerManagedEncryptionKeysRequest{
				OrganizationId: testOrgID,
				RegionIds:      []globalcommonv1.Region_Name{globalcommonv1.Region_AWS_US_EAST_1},
			})
		require.Nil(t, bizErr)
		for _, c := range resp.GetCustomerManagedEncryptionKeys() {
			assert.Equal(t, globalcommonv1.Region_AWS_US_EAST_1, c.GetRegionId())
		}
	})

	t.Run("OnlyDefaultKeys_EmptyResult", func(t *testing.T) {
		ts := newTestService(t)

		// Create only a default key under a unique org
		uniqueOrg := uint64(99999)
		createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid, func(k *entity.EncryptionKey) {
			k.TenantId = uniqueOrg
			k.CustomerManaged = false
			k.MasterKey = "default_key"
		})

		resp, bizErr := ts.svc.ListCustomerManagedEncryptionKeys(ctx,
			&globalngv1.ListCustomerManagedEncryptionKeysRequest{
				OrganizationId: uniqueOrg,
			})
		require.Nil(t, bizErr)
		assert.Empty(t, resp.GetCustomerManagedEncryptionKeys())
	})

	t.Run("ServicePlan_RoundTrip", func(t *testing.T) {
		ts := newTestService(t)

		uniqueOrg := uint64(88801)
		cmek := createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid, func(k *entity.EncryptionKey) {
			k.TenantId = uniqueOrg
			k.ServicePlan = globalcommonv1.ServicePlan_Premium
		})

		resp, bizErr := ts.svc.ListCustomerManagedEncryptionKeys(ctx,
			&globalngv1.ListCustomerManagedEncryptionKeysRequest{
				OrganizationId: uniqueOrg,
			})
		require.Nil(t, bizErr)
		require.Len(t, resp.GetCustomerManagedEncryptionKeys(), 1)

		got := resp.GetCustomerManagedEncryptionKeys()[0]
		assert.Equal(t, cmek.Id, got.GetCustomerManagedEncryptionKeyId())
		assert.Equal(t, globalcommonv1.ServicePlan_Premium, got.GetServicePlan(),
			"ServicePlan should round-trip through DB correctly")
	})

	t.Run("ServicePlan_FallbackToRegion", func(t *testing.T) {
		ts := newTestService(t)

		// Insert a record directly via GORM with empty ServicePlan to simulate
		// legacy rows that predate the service_plan column.
		uniqueOrg := uint64(88802)
		wdb := (*dbProvider.Load()).WriteDB(ctx)
		legacyKey := &schema.EncryptionKey{
			ID:                 sfid.MustNew(),
			TenantId:           uniqueOrg,
			ProviderId:         1,
			RegionId:           1,
			ProviderIdentityId: 1,
			CustomerManaged:    true,
			ServicePlan:        "", // empty — simulates old record
			MasterKey:          "arn:aws:kms:us-west-2:123456789012:key/legacy-key",
			InitialDataplaneId: "dataplane-1",
		}
		require.NoError(t, wdb.Create(legacyKey).Error)

		resp, bizErr := ts.svc.ListCustomerManagedEncryptionKeys(ctx,
			&globalngv1.ListCustomerManagedEncryptionKeysRequest{
				OrganizationId: uniqueOrg,
			})
		require.Nil(t, bizErr)
		require.Len(t, resp.GetCustomerManagedEncryptionKeys(), 1)

		got := resp.GetCustomerManagedEncryptionKeys()[0]
		assert.Equal(t, legacyKey.ID, got.GetCustomerManagedEncryptionKeyId())
		assert.Equal(t, globalcommonv1.ServicePlan_Premium, got.GetServicePlan(),
			"empty service_plan should fall back to region's ServicePlan (Premium)")
	})
}

func TestCvtEncryptionKeyToCmek(t *testing.T) {
	ctx := t.Context()
	ts := newTestService(t)

	t.Run("AWS", func(t *testing.T) {
		key := &entity.EncryptionKey{
			Id:              100,
			TenantId:        testOrgID,
			ProviderName:    globalcommonv1.CloudProvider_AWS,
			RegionName:      globalcommonv1.Region_AWS_US_WEST_2,
			ServicePlan:     globalcommonv1.ServicePlan_Premium,
			CustomerManaged: true,
			MasterKey:       "arn:aws:kms:us-west-2:123456789012:key/test-key",
		}
		identity := &entity.RegionProviderIdentity{
			RegionProviderIdentity: &globalngv1.RegionProviderIdentity{
				Id:       1,
				Identity: "123456789012",
			},
			RegionId: 1,
		}

		cmek, bizErr := ts.svc.CvtEncryptionKeyToCmek(ctx, key, identity)
		require.Nil(t, bizErr)
		assert.Equal(t, uint64(100), cmek.GetCustomerManagedEncryptionKeyId())
		assert.Equal(t, globalcommonv1.CloudProvider_AWS, cmek.GetCloudProviderId())
		assert.Equal(t, "arn:aws:kms:us-west-2:123456789012:key/test-key", cmek.GetAwsKms().GetKmsKeyArn())
		assert.Equal(t, "123456789012", cmek.GetAwsPrincipal().GetAccountId())
		assert.Equal(t, strconv.FormatUint(testOrgID, 10), cmek.GetAwsPrincipal().GetExternalId())
	})

	t.Run("Aliyun", func(t *testing.T) {
		key := &entity.EncryptionKey{
			Id:              200,
			TenantId:        testOrgID,
			ProviderName:    globalcommonv1.CloudProvider_ALICLOUD,
			RegionName:      globalcommonv1.Region_NAME_UNSPECIFIED,
			ServicePlan:     globalcommonv1.ServicePlan_Premium,
			CustomerManaged: true,
			MasterKey:       "acs:kms:ap-southeast-1:12345:key/key-abc-123",
		}
		identity := &entity.RegionProviderIdentity{
			RegionProviderIdentity: &globalngv1.RegionProviderIdentity{
				Id:       2,
				Identity: "acs:ram::12345:role/aliyun-role",
			},
			RegionId: 1,
		}

		cmek, bizErr := ts.svc.CvtEncryptionKeyToCmek(ctx, key, identity)
		require.Nil(t, bizErr)
		assert.Equal(t, uint64(200), cmek.GetCustomerManagedEncryptionKeyId())
		assert.Equal(t, globalcommonv1.CloudProvider_ALICLOUD, cmek.GetCloudProviderId())
		assert.Equal(t, "acs:kms:ap-southeast-1:12345:key/key-abc-123", cmek.GetAliyunKms().GetKmsKeyArn())
		assert.Equal(t, "12345", cmek.GetAliyunPrincipal().GetAccountId())
	})

	t.Run("AliyunAccountIDIdentity", func(t *testing.T) {
		key := &entity.EncryptionKey{
			Id:              201,
			TenantId:        testOrgID,
			ProviderName:    globalcommonv1.CloudProvider_ALICLOUD,
			RegionName:      globalcommonv1.Region_NAME_UNSPECIFIED,
			ServicePlan:     globalcommonv1.ServicePlan_Premium,
			CustomerManaged: true,
			MasterKey:       "acs:kms:ap-southeast-1:12345:key/key-abc-123",
		}
		identity := &entity.RegionProviderIdentity{
			RegionProviderIdentity: &globalngv1.RegionProviderIdentity{
				Id:       3,
				Identity: "12345",
			},
			RegionId: 1,
		}

		cmek, bizErr := ts.svc.CvtEncryptionKeyToCmek(ctx, key, identity)
		require.Nil(t, bizErr)
		assert.Equal(t, "12345", cmek.GetAliyunPrincipal().GetAccountId())
	})

	t.Run("NilInputs", func(t *testing.T) {
		_, bizErr := ts.svc.CvtEncryptionKeyToCmek(ctx, nil, nil)
		require.NotNil(t, bizErr)
	})
}

func TestVerifyCmekAccessIamPrincipal(t *testing.T) {
	ctx := t.Context()
	ts := newTestService(t)

	t.Run("AWS_ValidArn", func(t *testing.T) {
		ts.mockRegionalCli.EXPECT().
			CheckKms(gomock.Any(), gomock.Any(), gomock.Any()).
			DoAndReturn(func(
				_ context.Context,
				_ *entity.Region,
				req *regionalserveradapter.CheckKmsRequest,
			) bizerr.BizError {
				assert.Equal(t, "1", req.SharedPoolId)
				assert.Equal(t, map[string]string{
					"ExternalId": strconv.FormatUint(testOrgID, 10),
				}, req.EncryptionContext)
				return nil
			})

		resp, bizErr := ts.svc.VerifyCmekAccessIamPrincipal(ctx,
			&globalngv1.VerifyCmekAccessIamPrincipalRequest{
				Key: newTestAwsCmekFull(),
			})
		require.Nil(t, bizErr)
		assert.True(t, resp.GetValid())
	})

	t.Run("ExistingClusterUsesClusterSharedPool", func(t *testing.T) {
		cluster := createTestCluster(ctx, t, ts.clusterRepo, idgenerator.NewSonyflakeID(2, ""), func(c *entity.Cluster) {
			c.SharedPoolId = "sp-existing"
		})
		ts.mockRegionalCli.EXPECT().
			CheckKms(gomock.Any(), gomock.Any(), gomock.Any()).
			DoAndReturn(func(
				_ context.Context,
				_ *entity.Region,
				req *regionalserveradapter.CheckKmsRequest,
			) bizerr.BizError {
				assert.Equal(t, "sp-existing", req.SharedPoolId)
				return nil
			})

		resp, bizErr := ts.svc.VerifyCmekAccessIamPrincipal(ctx,
			&globalngv1.VerifyCmekAccessIamPrincipalRequest{
				Key:       newTestAwsCmekFull(),
				ClusterId: cluster.GetClusterId(),
			})
		require.Nil(t, bizErr)
		assert.True(t, resp.GetValid())
	})

	t.Run("ValidateCustomerManagedEncryptionKeyUsesCreateSharedPoolTarget", func(t *testing.T) {
		for _, tt := range []struct {
			name                  string
			sharedPoolID          string
			expectedCheckPoolID   string
			expectedResolverCalls int32
		}{
			{
				name:                  "explicit shared pool",
				sharedPoolID:          "sp-validate",
				expectedCheckPoolID:   "sp-validate",
				expectedResolverCalls: 0,
			},
			{
				name:                  "default shared pool",
				sharedPoolID:          "",
				expectedCheckPoolID:   "1",
				expectedResolverCalls: 1,
			},
		} {
			t.Run(tt.name, func(t *testing.T) {
				ts := newTestService(t)
				ts.mockRegionalCli.EXPECT().
					CheckKms(gomock.Any(), gomock.Any(), gomock.Any()).
					DoAndReturn(func(
						_ context.Context,
						_ *entity.Region,
						req *regionalserveradapter.CheckKmsRequest,
					) bizerr.BizError {
						assert.Equal(t, tt.expectedCheckPoolID, req.SharedPoolId)
						return nil
					})

				resolverCallsBefore := ts.resourcePoolResolver.calls.Load()
				bizErr := ts.svc.ValidateCustomerManagedEncryptionKey(ctx, newTestAwsCmekFull(), tt.sharedPoolID)

				require.Nil(t, bizErr)
				assert.Equal(t, resolverCallsBefore+tt.expectedResolverCalls, ts.resourcePoolResolver.calls.Load())
			})
		}
	})

	t.Run("Aliyun_AccountIDIdentityUsesAccountID", func(t *testing.T) {
		cleanTestData(ctx, t)
		ts := newTestService(t)
		seedAliyunProviderIdentity(ctx, t, "12345")
		expectedAccountID := "12345"

		principal, bizErr := ts.svc.GetCmekAccessIamPrincipal(ctx, &globalngv1.GetCmekAccessIamPrincipalRequest{
			OrganizationId:  testOrgID,
			CloudProviderId: globalcommonv1.CloudProvider_ALICLOUD,
			RegionId:        globalcommonv1.Region_ALICLOUD_AP_SOUTHEAST_1,
			ServicePlan:     globalcommonv1.ServicePlan_Premium,
		})
		require.Nil(t, bizErr)
		assert.Equal(t, expectedAccountID, principal.GetAliyun().GetAccountId())

		ts.mockRegionalCli.EXPECT().
			CheckKms(gomock.Any(), gomock.Any(), gomock.Any()).
			DoAndReturn(func(
				_ context.Context,
				_ *entity.Region,
				req *regionalserveradapter.CheckKmsRequest,
			) bizerr.BizError {
				assert.Empty(t, req.EncryptionContext)
				return nil
			})

		resp, bizErr := ts.svc.VerifyCmekAccessIamPrincipal(ctx, &globalngv1.VerifyCmekAccessIamPrincipalRequest{
			Key: newTestAliyunCmekFull(expectedAccountID),
		})
		require.Nil(t, bizErr)
		assert.True(t, resp.GetValid())
	})

	t.Run("AWS_CheckKmsFails", func(t *testing.T) {
		ts.mockRegionalCli.EXPECT().CheckKms(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(bizerr.New(bizerr.ErrCodeParams, "kms access denied"))

		resp, bizErr := ts.svc.VerifyCmekAccessIamPrincipal(ctx,
			&globalngv1.VerifyCmekAccessIamPrincipalRequest{
				Key: newTestAwsCmekFull(),
			})
		require.Nil(t, resp)
		require.NotNil(t, bizErr)
		assert.Equal(t, bizerr.ErrCodeParams, bizErr.ErrCode())
		assert.Equal(t, "kms access denied", bizErr.ErrMsg())
		assert.Contains(t, bizErr.Error(), "data-plane kms check failed")
		assert.NotEqual(t, cloudResourcesNotReadyMessage, bizErr.ErrMsg())
	})

	t.Run("AWS_CheckKmsRegionalUnavailable", func(t *testing.T) {
		ts.mockRegionalCli.EXPECT().CheckKms(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(bizerr.New(bizerr.ErrCodeRegionalAPI, "regional server is not ready"))
		ts.mockRegionalCli.EXPECT().GetSharedPool(gomock.Any(), gomock.Any(), gomock.Any()).
			DoAndReturn(func(
				_ context.Context,
				_ *entity.Region,
				req *clusterngv1beta.GetSharedPoolRequest,
			) (*clusterngv1beta.SharedPool, bizerr.BizError) {
				assert.Equal(t, "1", req.GetSharedPoolId())
				return &clusterngv1beta.SharedPool{Phase: clusterngv1beta.SharedPool_AVAILABLE}, nil
			})

		resp, bizErr := ts.svc.VerifyCmekAccessIamPrincipal(ctx,
			&globalngv1.VerifyCmekAccessIamPrincipalRequest{
				Key: newTestAwsCmekFull(),
			})
		require.Nil(t, resp)
		require.NotNil(t, bizErr)
		assert.Equal(t, bizerr.ErrCodeRegionalAPI, bizErr.ErrCode())
		assert.Equal(t, "regional server is not ready", bizErr.ErrMsg())
		assert.Contains(t, bizErr.Error(), "data-plane kms check failed")
		assert.Contains(t, bizErr.Error(), "regional server is not ready")
	})

	t.Run("AWS_CheckKmsPhysicalClusterNotReadyWhenSharedPoolCreating", func(t *testing.T) {
		ts := newTestService(t)
		ts.mockRegionalCli.EXPECT().CheckKms(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(bizerr.New(bizerr.ErrCodeRegionalAPI, "regional server is not ready"))
		ts.mockRegionalCli.EXPECT().GetSharedPool(gomock.Any(), gomock.Any(), gomock.Any()).
			DoAndReturn(func(
				_ context.Context,
				_ *entity.Region,
				req *clusterngv1beta.GetSharedPoolRequest,
			) (*clusterngv1beta.SharedPool, bizerr.BizError) {
				assert.Equal(t, "1", req.GetSharedPoolId())
				return &clusterngv1beta.SharedPool{Phase: clusterngv1beta.SharedPool_CREATING}, nil
			})

		resp, bizErr := ts.svc.VerifyCmekAccessIamPrincipal(ctx,
			&globalngv1.VerifyCmekAccessIamPrincipalRequest{
				Key: newTestAwsCmekFull(),
			})
		require.Nil(t, resp)
		require.NotNil(t, bizErr)
		assert.Equal(t, bizerr.ErrCodePhysicalClusterNotReady, bizErr.ErrCode())
		assert.Equal(t, cloudResourcesNotReadyMessage, bizErr.ErrMsg())
		assert.Contains(t, bizErr.Error(), "physical cluster is not ready")
	})

	t.Run("AWS_CheckKmsPhysicalClusterNotReadyWhenSharedPoolMissing", func(t *testing.T) {
		ts := newTestService(t)
		ts.mockRegionalCli.EXPECT().CheckKms(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(bizerr.New(bizerr.ErrCodeRecordNotFound, "shared pool not found"))
		ts.mockRegionalCli.EXPECT().GetSharedPool(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(nil, bizerr.New(bizerr.ErrCodeRecordNotFound, "shared pool not found"))

		resp, bizErr := ts.svc.VerifyCmekAccessIamPrincipal(ctx,
			&globalngv1.VerifyCmekAccessIamPrincipalRequest{
				Key: newTestAwsCmekFull(),
			})
		require.Nil(t, resp)
		require.NotNil(t, bizErr)
		assert.Equal(t, bizerr.ErrCodePhysicalClusterNotReady, bizErr.ErrCode())
		assert.Equal(t, cloudResourcesNotReadyMessage, bizErr.ErrMsg())
		assert.Contains(t, bizErr.Error(), "physical cluster is not ready")
	})

	t.Run("AWS_CheckKmsPhysicalClusterNotReadyWhenSharedPoolTargetMissing", func(t *testing.T) {
		svc := &service{}
		for _, sharedPoolID := range []string{"", "0"} {
			bizErr := svc.kmsVerifyFailureError(ctx, nil, kmsCheckTarget{sharedPoolID: sharedPoolID},
				bizerr.New(bizerr.ErrCodeRegionalAPI, "regional server is not ready"),
				"data-plane kms check failed")
			require.NotNil(t, bizErr)
			assert.Equal(t, bizerr.ErrCodePhysicalClusterNotReady, bizErr.ErrCode())
			assert.Equal(t, cloudResourcesNotReadyMessage, bizErr.ErrMsg())
			assert.Contains(t, bizErr.Error(), "physical cluster is not ready")
		}
	})

	t.Run("AWS_CheckKmsReadinessProbeFailureLiftsOriginalRegional", func(t *testing.T) {
		ts := newTestService(t)
		ts.mockRegionalCli.EXPECT().CheckKms(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(bizerr.New(bizerr.ErrCodeRegionalAPI, "regional server is not ready"))
		ts.mockRegionalCli.EXPECT().GetSharedPool(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(nil, bizerr.New(bizerr.ErrCodeRegionalAPI, "get shared pool failed"))

		resp, bizErr := ts.svc.VerifyCmekAccessIamPrincipal(ctx,
			&globalngv1.VerifyCmekAccessIamPrincipalRequest{
				Key: newTestAwsCmekFull(),
			})
		require.Nil(t, resp)
		require.NotNil(t, bizErr)
		assert.Equal(t, bizerr.ErrCodeRegionalAPI, bizErr.ErrCode())
		assert.Equal(t, "regional server is not ready", bizErr.ErrMsg())
		assert.Contains(t, bizErr.Error(), "data-plane kms check failed")
		assert.NotContains(t, bizErr.Error(), "get shared pool failed")
	})

	t.Run("AWS_InvalidArn", func(t *testing.T) {
		resolverCalls := ts.resourcePoolResolver.calls.Load()
		cmek := newTestAwsCmekFull()
		cmek.Key = &globalngv1.CustomerManagedEncryptionKey_AwsKms_{
			AwsKms: &globalngv1.CustomerManagedEncryptionKey_AwsKms{
				KmsKeyArn: "not-a-valid-arn",
			},
		}

		resp, bizErr := ts.svc.VerifyCmekAccessIamPrincipal(ctx,
			&globalngv1.VerifyCmekAccessIamPrincipalRequest{
				Key: cmek,
			})
		require.Nil(t, bizErr)
		assert.False(t, resp.GetValid())
		assert.Contains(t, resp.GetFailReason(), "arn")
		assert.Equal(t, resolverCalls, ts.resourcePoolResolver.calls.Load(),
			"invalid KMS input must not create or resolve a default shared pool")
	})

	t.Run("AWS_WrongAccountId", func(t *testing.T) {
		cmek := newTestAwsCmekFull()
		cmek.Principal = &globalngv1.CustomerManagedEncryptionKey_AwsPrincipal{
			AwsPrincipal: &globalngv1.CmekAccessIamPrincipal_AwsCmekPrincipal{
				AccountId: "wrong-account",
			},
		}

		_, bizErr := ts.svc.VerifyCmekAccessIamPrincipal(ctx,
			&globalngv1.VerifyCmekAccessIamPrincipalRequest{
				Key: cmek,
			})
		require.NotNil(t, bizErr)
		assert.Contains(t, bizErr.Error(), "account not matched")
	})

	t.Run("UnsupportedServicePlan", func(t *testing.T) {
		resolverCalls := ts.resourcePoolResolver.calls.Load()
		cmek := newTestAwsCmekFull()
		cmek.ServicePlan = globalcommonv1.ServicePlan_SERVICE_PLAN_UNSPECIFIED

		_, bizErr := ts.svc.VerifyCmekAccessIamPrincipal(ctx,
			&globalngv1.VerifyCmekAccessIamPrincipalRequest{
				Key: cmek,
			})
		require.NotNil(t, bizErr)
		assert.Contains(t, bizErr.Error(), "only support Premium")
		assert.Equal(t, resolverCalls, ts.resourcePoolResolver.calls.Load(),
			"invalid service plan must not create or resolve a default shared pool")
	})
}

func TestCreateCustomerManagedEncryptionKey(t *testing.T) {
	ctx := t.Context()

	t.Run("VerifyInvalidReturnsBizError", func(t *testing.T) {
		cleanTestData(ctx, t)
		ts := newTestService(t)
		ts.mockRegionalCli.EXPECT().CheckKms(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(bizerr.New(bizerr.ErrCodeParams, "kms access denied"))

		cmek, bizErr := ts.svc.CreateCustomerManagedEncryptionKey(
			ctx,
			newTestAwsCmekFull(),
			1,
			"shared-pool-1",
			"dataplane-1",
		)

		require.Nil(t, cmek)
		require.NotNil(t, bizErr)
		assert.Equal(t, bizerr.ErrCodeParams, bizErr.ErrCode())
		assert.Equal(t, "kms access denied", bizErr.ErrMsg())
		assert.Contains(t, bizErr.Error(), "data-plane kms check failed")
		assert.NotEqual(t, cloudResourcesNotReadyMessage, bizErr.ErrMsg())
	})

	t.Run("AliyunVerifyInvalidReturnsBizError", func(t *testing.T) {
		cleanTestData(ctx, t)
		ts := newTestService(t)
		seedAliyunProviderIdentity(ctx, t, "12345")
		ts.mockRegionalCli.EXPECT().CheckKms(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(bizerr.New(bizerr.ErrCodeParams, "Forbidden.KeyNotFound"))

		cmek, bizErr := ts.svc.CreateCustomerManagedEncryptionKey(
			ctx,
			newTestAliyunCmekFull(""),
			1,
			"2048593891413721088",
			"dataplane-1",
		)

		require.Nil(t, cmek)
		require.NotNil(t, bizErr)
		assert.Equal(t, bizerr.ErrCodeParams, bizErr.ErrCode())
		assert.Equal(t, "Forbidden.KeyNotFound", bizErr.ErrMsg())
		assert.Contains(t, bizErr.Error(), "data-plane kms check failed")
		assert.NotEqual(t, cloudResourcesNotReadyMessage, bizErr.ErrMsg())
	})

	t.Run("RepoCreateFailureReturnsBizError", func(t *testing.T) {
		cleanTestData(ctx, t)
		ts := newTestService(t)
		svc := NewService(
			ts.regionRepo,
			ts.providerIdentityRepo,
			ts.regionService,
			failingCreateEncryptionKeyRepo{Repo: ts.encryptionKeyRepo},
			ts.clusterRepo,
			ts.bindingRepo,
			*dbTransactioner.Load(),
			ts.mockRegionalCli,
			ts.mockWorkflow,
			ts.resourcePoolResolver,
			idgenerator.NewSonyflakeID(9, ""),
		)
		ts.mockRegionalCli.EXPECT().CheckKms(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil)

		cmek, bizErr := svc.CreateCustomerManagedEncryptionKey(
			ctx,
			newTestAwsCmekFull(),
			1,
			"shared-pool-1",
			"dataplane-1",
		)

		require.Nil(t, cmek)
		require.NotNil(t, bizErr)
		assert.Contains(t, bizErr.Error(), "createDefaultKey customer managed encryption key failed")
		assert.Contains(t, bizErr.Error(), "injected encryption key create failure")
	})
}

func TestCreateDefaultKeyOrGrantPermissionRegionalFailureReturnsCloudResourcesNotReady(t *testing.T) {
	ctx := t.Context()
	cleanTestData(ctx, t)
	ts := newTestService(t)
	sfid := idgenerator.NewSonyflakeID(10, "")
	createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid, func(k *entity.EncryptionKey) {
		k.CustomerManaged = false
		k.MasterKey = ""
		k.ProviderIdentityId = 1
		k.InitialDataplaneId = "source-dataplane"
	})

	ts.mockRegionalCli.EXPECT().
		AddEncryptionKeyTrustedDataplane(gomock.Any(), gomock.Any(), gomock.Any()).
		Return(bizerr.New(bizerr.ErrCodeRegionalAPI, "regional server is not ready"))

	key, bizErr := ts.svc.CreateDefaultKeyOrGrantPermission(ctx, testOrgID, 1, "target-dataplane")

	require.Nil(t, key)
	require.NotNil(t, bizErr)
	assert.Equal(t, cloudResourcesNotReadyMessage, bizErr.ErrMsg())
	assert.Contains(t, bizErr.Error(), "grant default encryption key to dataplane failed")
	assert.Contains(t, bizErr.Error(), "regional server is not ready")
}

func TestGetCmekAccessIamPrincipal(t *testing.T) {
	ctx := t.Context()
	ts := newTestService(t)

	t.Run("AWS_Success", func(t *testing.T) {
		result, bizErr := ts.svc.GetCmekAccessIamPrincipal(ctx,
			&globalngv1.GetCmekAccessIamPrincipalRequest{
				OrganizationId:  testOrgID,
				CloudProviderId: globalcommonv1.CloudProvider_AWS,
				RegionId:        globalcommonv1.Region_AWS_US_WEST_2,
				ServicePlan:     globalcommonv1.ServicePlan_Premium,
			})
		require.Nil(t, bizErr)
		assert.Equal(t, "123456789012", result.GetAws().GetAccountId())
	})

	t.Run("UnsupportedProvider", func(t *testing.T) {
		_, bizErr := ts.svc.GetCmekAccessIamPrincipal(ctx,
			&globalngv1.GetCmekAccessIamPrincipalRequest{
				OrganizationId:  testOrgID,
				CloudProviderId: globalcommonv1.CloudProvider_GCP,
				RegionId:        globalcommonv1.Region_GCP_US_EAST1,
				ServicePlan:     globalcommonv1.ServicePlan_Premium,
			})
		require.NotNil(t, bizErr)
		assert.Contains(t, bizErr.Error(), "only support Premium on AWS or Aliyun")
	})
}

func TestAliyunAccountID(t *testing.T) {
	assert.Equal(t, "12345", aliyunAccountID("12345"))
	assert.Equal(t, "12345", aliyunAccountID("acs:ram::12345:role/admin"))
	assert.Equal(t, "12345", aliyunAccountID("acs:ram::12345:*"))
	assert.Equal(t, "not-an-account", aliyunAccountID("not-an-account"))
	assert.Empty(t, aliyunAccountID(""))
}

type failingCreateEncryptionKeyRepo struct {
	encryptionkeyrepo.Repo
}

func (failingCreateEncryptionKeyRepo) Create(_ context.Context, _ *entity.EncryptionKey) bizerr.BizError {
	return bizerr.New(bizerr.ErrCodeDB, "injected encryption key create failure")
}

func newTestAliyunCmekFull(accountID string) *globalngv1.CustomerManagedEncryptionKey {
	return &globalngv1.CustomerManagedEncryptionKey{
		CloudProviderId: globalcommonv1.CloudProvider_ALICLOUD,
		RegionId:        globalcommonv1.Region_ALICLOUD_AP_SOUTHEAST_1,
		ServicePlan:     globalcommonv1.ServicePlan_Premium,
		OrganizationId:  testOrgID,
		Key: &globalngv1.CustomerManagedEncryptionKey_AliyunKms_{
			AliyunKms: &globalngv1.CustomerManagedEncryptionKey_AliyunKms{
				KmsKeyArn: "acs:kms:ap-southeast-1:12345:key/key-abc-123",
			},
		},
		Principal: &globalngv1.CustomerManagedEncryptionKey_AliyunPrincipal{
			AliyunPrincipal: &globalngv1.CmekAccessIamPrincipal_AliyunCmekPrincipal{
				AccountId: accountID,
			},
		},
	}
}

func seedAliyunProviderIdentity(ctx context.Context, t *testing.T, identity string) {
	t.Helper()
	wdb := (*dbProvider.Load()).WriteDB(ctx)
	require.NoError(t, wdb.Exec("UPDATE providers SET status = ? WHERE id = ?", "active", 2).Error)

	isPublic := true
	regionName := "alicloud-ap-southeast-1-cmek-test-" +
		strings.NewReplacer("/", "-", " ", "-").Replace(t.Name())
	cloudRegion := &schema.Region{
		Name:              regionName,
		RegionName:        "ap-southeast-1",
		ProviderId:        2,
		DisplayName:       "Singapore (ap-southeast-1)",
		CountryCode:       "SG",
		Area:              "APAC",
		Priority:          100,
		DefaultBackupTime: "10:00",
		ServiceAddress:    "https://alicloud-ap-southeast-1.tidbcloud.com",
		ServicePlan:       "Premium",
		IsPublic:          &isPublic,
		Status:            "active",
	}
	require.NoError(t, wdb.Create(cloudRegion).Error)
	require.NoError(t, wdb.Create(&schema.RegionServicePlan{
		RegionId:    uint64(cloudRegion.ID),
		ServicePlan: "Premium",
		IsPublic:    &isPublic,
	}).Error)
	require.NoError(t, wdb.Create(&schema.ProviderIdentity{
		ID:       uint64(cloudRegion.ID),
		RegionId: uint64(cloudRegion.ID),
		Identity: identity,
	}).Error)
	require.NoError(t, wdb.Clauses(clause.OnConflict{
		Columns: []clause.Column{
			{Name: "provider_id"},
			{Name: "service_plan"},
			{Name: "region_id"},
			{Name: "tenant_id"},
		},
		DoUpdates: clause.AssignmentColumns([]string{"identity"}),
	}).Create(&schema.DefaultProviderIdentity{
		ProviderId:  2,
		ServicePlan: "Premium",
		RegionId:    0,
		TenantId:    0,
		Identity:    identity,
	}).Error)
}