package cluster

import (
	"context"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/mock/gomock"

	globalcommonv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/common/v1"
	globalngv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/nextgen/v1"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/adapter/regionalserver"
	mockregionalserver "github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/adapter/regionalserver/mock_regionalserver"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/entity"
	mockencryptionkeyrepo "github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/repository/encryptionkey/mock_encryptionkey"
	mockregionrepo "github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/repository/region/mock_region"
	encryptionkeysvc "github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/service/encryptionkey"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
)

func TestHandleDualLayerEncryptionKeyFromSourceClusterValidatesCmekAccess(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	ctx := t.Context()
	encryptionKeyRepo := mockencryptionkeyrepo.NewMockRepo(ctrl)
	regionRepo := mockregionrepo.NewMockRepo(ctrl)
	regionalCli := mockregionalserver.NewMockClient(ctrl)
	svc := &serviceImpl{
		encryptionKeyRepo: encryptionKeyRepo,
		regionRepo:        regionRepo,
		regionalCli:       regionalCli,
	}

	const (
		keyID              uint64 = 77
		organizationID     uint64 = 42
		sourceClusterID    uint64 = 123
		targetSharedPoolID        = "1001"
	)
	keyArn := "arn:aws:kms:us-east-1:123456789012:key/source-key"
	sourceCluster := &entity.Cluster{
		Cluster: &globalngv1.Cluster{
			ClusterId:      sourceClusterID,
			OrganizationId: organizationID,
		},
		EncryptionKeyId: keyID,
	}
	req := &globalngv1.CreateClusterRequest{
		Cluster: &globalngv1.Cluster{
			OrganizationId: organizationID,
			SharedPoolId:   targetSharedPoolID,
		},
	}
	key := &entity.EncryptionKey{
		Id:              keyID,
		TenantId:        organizationID,
		ProviderName:    globalcommonv1.CloudProvider_AWS,
		RegionName:      globalcommonv1.Region_AWS_US_EAST_1,
		ServicePlan:     globalcommonv1.ServicePlan_Premium,
		CustomerManaged: true,
		MasterKey:       keyArn,
		RegionId:        101,
	}
	region := &entity.Region{
		Region: &globalcommonv1.Region{
			Id:              globalcommonv1.Region_AWS_US_EAST_1,
			CloudProviderId: globalcommonv1.CloudProvider_AWS,
			ServicePlan:     globalcommonv1.ServicePlan_Premium,
		},
		RegionID: 101,
	}

	encryptionKeyRepo.EXPECT().Load(gomock.Any(), keyID, false).Return(key, nil)
	regionRepo.EXPECT().Get(gomock.Any(), uint64(101)).Return(region, nil)
	regionalCli.EXPECT().
		CheckKms(gomock.Any(), region, gomock.AssignableToTypeOf(&regionalserver.CheckKmsRequest{})).
		DoAndReturn(func(_ any, _ *entity.Region, checkReq *regionalserver.CheckKmsRequest) bizerr.BizError {
			assert.Equal(t, targetSharedPoolID, checkReq.SharedPoolId)
			assert.Equal(t, "aws", checkReq.Provider)
			assert.Equal(t, keyArn, checkReq.KeyId)
			assert.Equal(t, "us-east-1", checkReq.Region)
			assert.Equal(t, map[string]string{"ExternalId": "42"}, checkReq.EncryptionContext)
			return nil
		})

	gotKeyID, err := svc.handleDualLayerEncryptionKeyFromSourceCluster(ctx, req, 0, "", sourceCluster)

	require.Nil(t, err)
	assert.Equal(t, keyID, gotKeyID)
}

func TestHandleDualLayerEncryptionKeyFromSourceClusterFailsWhenCmekAccessInvalid(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	ctx := t.Context()
	encryptionKeyRepo := mockencryptionkeyrepo.NewMockRepo(ctrl)
	regionRepo := mockregionrepo.NewMockRepo(ctrl)
	regionalCli := mockregionalserver.NewMockClient(ctrl)
	svc := &serviceImpl{
		encryptionKeyRepo: encryptionKeyRepo,
		regionRepo:        regionRepo,
		regionalCli:       regionalCli,
	}

	const keyID uint64 = 77
	key := &entity.EncryptionKey{
		Id:              keyID,
		TenantId:        42,
		ProviderName:    globalcommonv1.CloudProvider_AWS,
		RegionName:      globalcommonv1.Region_AWS_US_EAST_1,
		ServicePlan:     globalcommonv1.ServicePlan_Premium,
		CustomerManaged: true,
		MasterKey:       "arn:aws:kms:us-east-1:123456789012:key/source-key",
		RegionId:        101,
	}
	region := &entity.Region{
		Region: &globalcommonv1.Region{
			Id:              globalcommonv1.Region_AWS_US_EAST_1,
			CloudProviderId: globalcommonv1.CloudProvider_AWS,
			ServicePlan:     globalcommonv1.ServicePlan_Premium,
		},
		RegionID: 101,
	}

	encryptionKeyRepo.EXPECT().Load(gomock.Any(), keyID, false).Return(key, nil)
	regionRepo.EXPECT().Get(gomock.Any(), uint64(101)).Return(region, nil)
	regionalCli.EXPECT().
		CheckKms(gomock.Any(), region, gomock.AssignableToTypeOf(&regionalserver.CheckKmsRequest{})).
		Return(bizerr.New(bizerr.ErrCodeParams, "kms key disabled"))

	gotKeyID, err := svc.handleDualLayerEncryptionKeyFromSourceCluster(ctx,
		&globalngv1.CreateClusterRequest{Cluster: &globalngv1.Cluster{OrganizationId: 42, SharedPoolId: "1001"}},
		0,
		"",
		&entity.Cluster{
			Cluster: &globalngv1.Cluster{
				ClusterId:      123,
				OrganizationId: 42,
			},
			EncryptionKeyId: keyID,
		})

	require.NotNil(t, err)
	assert.Zero(t, gotKeyID)
	assert.Equal(t, restoreKmsUnavailableErrorMsg, err.ErrMsg())
}

func TestHandleDualLayerEncryptionKeyFromSourceClusterSkipsPreflightForDefaultKey(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	ctx := t.Context()
	encryptionKeyRepo := mockencryptionkeyrepo.NewMockRepo(ctrl)
	encryptionKeySvc := &restoreEncryptionKeyService{
		defaultKey: &entity.EncryptionKey{
			Id:              88,
			TenantId:        42,
			ProviderName:    globalcommonv1.CloudProvider_AWS,
			RegionName:      globalcommonv1.Region_AWS_US_EAST_1,
			ServicePlan:     globalcommonv1.ServicePlan_Premium,
			CustomerManaged: false,
			MasterKey:       "",
			RegionId:        101,
		},
	}
	svc := &serviceImpl{
		encryptionKeyRepo: encryptionKeyRepo,
		encryptionKeySvc:  encryptionKeySvc,
	}

	const sourceKeyID uint64 = 77
	encryptionKeyRepo.EXPECT().Load(gomock.Any(), sourceKeyID, false).Return(&entity.EncryptionKey{
		Id:              sourceKeyID,
		TenantId:        42,
		ProviderName:    globalcommonv1.CloudProvider_AWS,
		RegionName:      globalcommonv1.Region_AWS_US_EAST_1,
		ServicePlan:     globalcommonv1.ServicePlan_Premium,
		CustomerManaged: false,
		MasterKey:       "existing-default-key",
		RegionId:        101,
	}, nil)

	gotKeyID, err := svc.handleDualLayerEncryptionKeyFromSourceCluster(ctx,
		&globalngv1.CreateClusterRequest{Cluster: &globalngv1.Cluster{OrganizationId: 42, SharedPoolId: "1001"}},
		12,
		"target-dataplane",
		&entity.Cluster{
			Cluster: &globalngv1.Cluster{
				ClusterId:      123,
				OrganizationId: 42,
			},
			EncryptionKeyId: sourceKeyID,
		})

	require.Nil(t, err)
	assert.Equal(t, uint64(88), gotKeyID)
	assert.Equal(t, 1, encryptionKeySvc.defaultKeyCalls)
}

func TestHandleDualLayerEncryptionKeyFromSourceClusterPreservesRegionalCheckKmsErrorCode(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	ctx := t.Context()
	encryptionKeyRepo := mockencryptionkeyrepo.NewMockRepo(ctrl)
	regionRepo := mockregionrepo.NewMockRepo(ctrl)
	regionalCli := mockregionalserver.NewMockClient(ctrl)
	svc := &serviceImpl{
		encryptionKeyRepo: encryptionKeyRepo,
		regionRepo:        regionRepo,
		regionalCli:       regionalCli,
	}

	const keyID uint64 = 77
	key := &entity.EncryptionKey{
		Id:              keyID,
		TenantId:        42,
		ProviderName:    globalcommonv1.CloudProvider_AWS,
		RegionName:      globalcommonv1.Region_AWS_US_EAST_1,
		ServicePlan:     globalcommonv1.ServicePlan_Premium,
		CustomerManaged: true,
		MasterKey:       "arn:aws:kms:us-east-1:123456789012:key/source-key",
		RegionId:        101,
	}
	region := &entity.Region{
		Region: &globalcommonv1.Region{
			Id:              globalcommonv1.Region_AWS_US_EAST_1,
			CloudProviderId: globalcommonv1.CloudProvider_AWS,
			ServicePlan:     globalcommonv1.ServicePlan_Premium,
		},
		RegionID: 101,
	}

	encryptionKeyRepo.EXPECT().Load(gomock.Any(), keyID, false).Return(key, nil)
	regionRepo.EXPECT().Get(gomock.Any(), uint64(101)).Return(region, nil)
	regionalCli.EXPECT().
		CheckKms(gomock.Any(), region, gomock.AssignableToTypeOf(&regionalserver.CheckKmsRequest{})).
		Return(bizerr.New(bizerr.ErrCodeRegionalAPI, "regional server unavailable"))

	gotKeyID, err := svc.handleDualLayerEncryptionKeyFromSourceCluster(ctx,
		&globalngv1.CreateClusterRequest{Cluster: &globalngv1.Cluster{OrganizationId: 42, SharedPoolId: "1001"}},
		0,
		"",
		&entity.Cluster{
			Cluster: &globalngv1.Cluster{
				ClusterId:      123,
				OrganizationId: 42,
			},
			EncryptionKeyId: keyID,
		})

	require.NotNil(t, err)
	assert.Zero(t, gotKeyID)
	assert.Equal(t, bizerr.ErrCodeRegionalAPI, err.ErrCode())
	assert.Equal(t, "regional server unavailable", err.ErrMsg())
}

type restoreEncryptionKeyService struct {
	encryptionkeysvc.Service

	defaultKey      *entity.EncryptionKey
	defaultKeyCalls int
}

func (s *restoreEncryptionKeyService) CreateDefaultKeyOrGrantPermission(
	context.Context,
	uint64,
	uint64,
	string,
) (*entity.EncryptionKey, bizerr.BizError) {
	s.defaultKeyCalls++
	return s.defaultKey, nil
}