package cluster

import (
	"context"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	clusterngv1beta "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	dataplanemanagerv1 "github.com/tidbcloud/cluster-service-ng/api/tidbcloud/dataplane-manager/v1"
	"go.uber.org/mock/gomock"
	"google.golang.org/protobuf/types/known/timestamppb"

	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/mock_regionalserver"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/entity"
	clusterrepo "github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/repository/cluster"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/repository/cluster/mock_cluster"
	regionrepo "github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/repository/region"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/repository/region/mock_region"
	resourcepoolrepo "github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/repository/resourcepool"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/repository/resourcepool/mock_resourcepool"
	tidbgrouprepo "github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/repository/tidbgroup"
	encryptionkeysvc "github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/service/encryptionkey"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
)

func TestListPassesResourcePoolIDsToRepository(t *testing.T) {
	ctrl := gomock.NewController(t)
	clusterRepo := mock_cluster.NewMockRepo(ctrl)
	svc := &serviceImpl{clusterRepo: clusterRepo}

	clusterRepo.EXPECT().
		Query(gomock.Any(), gomock.Any()).
		DoAndReturn(func(
			_ context.Context,
			filter *clusterrepo.Filter,
		) ([]*entity.Cluster, bizerr.BizError) {
			require.Equal(t, []uint64{42, 43}, filter.SharedPoolIds)
			require.Equal(t, []uint64{200}, filter.OrganizationIds)
			require.Equal(t, []entity.ServicePlan{globalcommonv1.ServicePlan_BYOC}, filter.ServicePlans)
			return nil, nil
		})

	clusters, err := svc.List(t.Context(), &globalngv1.ListClustersRequest{
		OrganizationIds: []uint64{200},
		ServicePlans:    []globalcommonv1.ServicePlan{globalcommonv1.ServicePlan_BYOC},
		ResourcePoolIds: []uint64{42, 43},
	})

	require.Nil(t, err)
	require.Nil(t, clusters)
}

func TestFullResourcePoolDisplayNamesUsesGlobalResourcePoolRepository(t *testing.T) {
	ctrl := gomock.NewController(t)
	resourcePoolRepo := mock_resourcepool.NewMockRepo(ctrl)
	svc := &serviceImpl{resourcePoolRepo: resourcePoolRepo}
	clusters := []*entity.Cluster{
		{Cluster: &globalngv1.Cluster{
			ClusterId:      1001,
			OrganizationId: 200,
			ServicePlan:    globalcommonv1.ServicePlan_BYOC,
			SharedPoolId:   "42",
		}},
		{Cluster: &globalngv1.Cluster{
			ClusterId:      1002,
			OrganizationId: 200,
			ServicePlan:    globalcommonv1.ServicePlan_BYOC,
			SharedPoolId:   "43",
		}},
		{Cluster: &globalngv1.Cluster{
			ClusterId:    1003,
			ServicePlan:  globalcommonv1.ServicePlan_Premium,
			SharedPoolId: "99",
		}},
		{Cluster: &globalngv1.Cluster{
			ClusterId:    1004,
			ServicePlan:  globalcommonv1.ServicePlan_BYOC,
			SharedPoolId: "not-an-id",
		}},
	}

	resourcePoolRepo.EXPECT().
		Query(gomock.Any(), gomock.Any()).
		DoAndReturn(func(
			_ context.Context,
			filter *resourcepoolrepo.Filter,
		) ([]*entity.ResourcePool, int64, bizerr.BizError) {
			assert.ElementsMatch(t, []uint64{42, 43}, filter.Ids)
			assert.Equal(t, []entity.ServicePlan{globalcommonv1.ServicePlan_BYOC}, filter.ServicePlans)
			assert.Equal(t, int(globalngv1.ResourcePool_ORGANIZATION), filter.OwnerType)
			assert.Equal(t, []uint64{200}, filter.OwnerIds)
			assert.True(t, filter.ShowDeleted)
			assert.Equal(t, int64(2), filter.PageSize)
			return []*entity.ResourcePool{
				{ResourcePool: &globalngv1.ResourcePool{Id: 42, Name: "orders"}},
				{ResourcePool: &globalngv1.ResourcePool{Id: 43, Name: "analytics"}},
			}, 2, nil
		})

	err := svc.fullResourcePoolDisplayNames(t.Context(), clusters, true)

	require.Nil(t, err)
	assert.Equal(t, "orders", clusters[0].GetResourcePoolDisplayName())
	assert.Equal(t, "analytics", clusters[1].GetResourcePoolDisplayName())
	assert.Empty(t, clusters[2].GetResourcePoolDisplayName())
	assert.Empty(t, clusters[3].GetResourcePoolDisplayName())
}

func TestFullResourcePoolLifecycleProjectsPoolStateToAvailableCluster(t *testing.T) {
	ctrl := gomock.NewController(t)
	resourcePoolRepo := mock_resourcepool.NewMockRepo(ctrl)
	svc := &serviceImpl{resourcePoolRepo: resourcePoolRepo}
	resumeAt := timestamppb.Now()
	clusters := []*entity.Cluster{
		{Cluster: &globalngv1.Cluster{ClusterId: 1001, State: globalngv1.Cluster_AVAILABLE, SharedPoolId: "42"}},
		{Cluster: &globalngv1.Cluster{ClusterId: 1002, State: globalngv1.Cluster_AVAILABLE, SharedPoolId: "43"}},
		{Cluster: &globalngv1.Cluster{ClusterId: 1003, State: globalngv1.Cluster_AVAILABLE, SharedPoolId: "44"}},
		{Cluster: &globalngv1.Cluster{ClusterId: 1004, State: globalngv1.Cluster_MODIFYING, SharedPoolId: "45"}},
	}
	resourcePoolRepo.EXPECT().Query(gomock.Any(), gomock.Any()).DoAndReturn(func(
		_ context.Context,
		filter *resourcepoolrepo.Filter,
	) ([]*entity.ResourcePool, int64, bizerr.BizError) {
		assert.ElementsMatch(t, []uint64{42, 43, 44, 45}, filter.Ids)
		assert.False(t, filter.ShowDeleted)
		return []*entity.ResourcePool{
			{ResourcePool: &globalngv1.ResourcePool{Id: 42, LifecycleState: globalngv1.ResourcePool_PAUSING}},
			{ResourcePool: &globalngv1.ResourcePool{Id: 43, LifecycleState: globalngv1.ResourcePool_PAUSED, AutomaticResumeAt: resumeAt}},
			{ResourcePool: &globalngv1.ResourcePool{Id: 44, LifecycleState: globalngv1.ResourcePool_RESUMING}},
			{ResourcePool: &globalngv1.ResourcePool{Id: 45, LifecycleState: globalngv1.ResourcePool_PAUSED, AutomaticResumeAt: resumeAt}},
		}, 4, nil
	})

	err := svc.fullResourcePoolLifecycle(t.Context(), clusters, false)

	require.Nil(t, err)
	assert.Equal(t, globalngv1.Cluster_PAUSING, clusters[0].GetState())
	assert.Equal(t, globalngv1.Cluster_PAUSED, clusters[1].GetState())
	assert.Same(t, resumeAt, clusters[1].GetAutomaticResumeAt())
	assert.Equal(t, globalngv1.Cluster_RESUMING, clusters[2].GetState())
	assert.Equal(t, globalngv1.Cluster_MODIFYING, clusters[3].GetState())
}

func TestResourcePoolLifecycleStateFiltersQueryAvailableAndFilterProjectedState(t *testing.T) {
	queryStates := resourcePoolLifecycleQueryStates([]globalngv1.Cluster_State{globalngv1.Cluster_PAUSED})
	require.ElementsMatch(t, []globalngv1.Cluster_State{globalngv1.Cluster_AVAILABLE}, queryStates)

	clusters := []*entity.Cluster{
		{Cluster: &globalngv1.Cluster{State: globalngv1.Cluster_PAUSED}},
		{Cluster: &globalngv1.Cluster{State: globalngv1.Cluster_AVAILABLE}},
	}
	filtered := filterResourcePoolLifecycleClusters(
		clusters,
		[]globalngv1.Cluster_State{globalngv1.Cluster_PAUSED},
		nil,
	)

	require.Len(t, filtered, 1)
	assert.Equal(t, globalngv1.Cluster_PAUSED, filtered[0].GetState())

	filtered = filterResourcePoolLifecycleClusters(
		clusters,
		[]globalngv1.Cluster_State{globalngv1.Cluster_AVAILABLE},
		nil,
	)
	require.Len(t, filtered, 1)
	assert.Equal(t, globalngv1.Cluster_AVAILABLE, filtered[0].GetState())
}

func TestResourcePoolLifecycleExcludeStateQueryKeepsAvailableCandidates(t *testing.T) {
	require.Empty(t, resourcePoolLifecycleQueryExcludeStates([]globalngv1.Cluster_State{
		globalngv1.Cluster_PAUSED,
	}))
	require.ElementsMatch(t,
		[]globalngv1.Cluster_State{globalngv1.Cluster_CREATING},
		resourcePoolLifecycleQueryExcludeStates([]globalngv1.Cluster_State{
			globalngv1.Cluster_PAUSED,
			globalngv1.Cluster_CREATING,
		}),
	)
	require.Empty(t, resourcePoolLifecycleQueryExcludeStates([]globalngv1.Cluster_State{
		globalngv1.Cluster_AVAILABLE,
	}))

	filtered := filterResourcePoolLifecycleClusters(
		[]*entity.Cluster{
			{Cluster: &globalngv1.Cluster{State: globalngv1.Cluster_PAUSED}},
			{Cluster: &globalngv1.Cluster{State: globalngv1.Cluster_AVAILABLE}},
		},
		nil,
		[]globalngv1.Cluster_State{globalngv1.Cluster_PAUSED},
	)
	require.Len(t, filtered, 1)
	require.Equal(t, globalngv1.Cluster_AVAILABLE, filtered[0].GetState())

	filtered = filterResourcePoolLifecycleClusters(
		[]*entity.Cluster{
			{Cluster: &globalngv1.Cluster{State: globalngv1.Cluster_PAUSED}},
			{Cluster: &globalngv1.Cluster{State: globalngv1.Cluster_AVAILABLE}},
		},
		nil,
		[]globalngv1.Cluster_State{globalngv1.Cluster_AVAILABLE},
	)
	require.Len(t, filtered, 1)
	require.Equal(t, globalngv1.Cluster_PAUSED, filtered[0].GetState())
}

func TestFullClusterInfoSplitsPublicAndVpcPeeringEndpoints(t *testing.T) {
	ctx := t.Context()
	ctrl := gomock.NewController(t)
	regionRepo := mock_region.NewMockRepo(ctrl)
	regionalCli := mock_regionalserver.NewMockClient(ctrl)
	svc := &serviceImpl{
		regionRepo:    regionRepo,
		tidbgroupRepo: &stubTidbGroupRepo{},
		regionalCli:   regionalCli,
	}

	cluster := &entity.Cluster{
		Cluster:  &globalngv1.Cluster{ClusterId: 1001},
		RegionId: 2001,
	}
	region := &entity.Region{
		Region:   &globalcommonv1.Region{Id: globalcommonv1.Region_AWS_US_WEST_2},
		RegionID: 2001,
	}

	regionRepo.EXPECT().ListRegions(gomock.Any(), gomock.Any()).Return([]*entity.Region{region}, nil)
	regionalCli.EXPECT().GetTidbGroup(gomock.Any(), region, gomock.Any()).Return(&clusterngv1beta.TiDBGroup{
		Id:        "3001",
		ClusterId: "1001",
		SqlEndpoint: &clusterngv1beta.TiDBGroup_SqlEndpoint{
			PublicLb: &clusterngv1beta.TiDBGroup_SqlEndpoint_LoadBalancer{
				Enabled: true,
				Host:    "public.example.com",
				Port:    4000,
				IpAllowList: []*clusterngv1beta.TiDBGroup_SqlEndpoint_LoadBalancer_IpAddress{
					{CidrBlock: "10.0.0.0/24"},
				},
			},
			PrivateLb: &clusterngv1beta.TiDBGroup_SqlEndpoint_LoadBalancer{
				Enabled: true,
				Host:    "private.example.com",
				Port:    5000,
			},
		},
	}, nil)

	err := svc.fullClusterInfo(ctx, []*entity.Cluster{cluster}, false)

	require.Nil(t, err)
	require.Len(t, cluster.Endpoints, 2)
	assert.NotNil(t, cluster.Endpoints[0].PublicEndpoint)
	assert.Nil(t, cluster.Endpoints[0].VpcPeeringEndpoint)
	assert.Equal(t, "public.example.com", cluster.Endpoints[0].PublicEndpoint.Host)
	assert.NotNil(t, cluster.Endpoints[1].VpcPeeringEndpoint)
	assert.Nil(t, cluster.Endpoints[1].PublicEndpoint)
	assert.Equal(t, "private.example.com", cluster.Endpoints[1].VpcPeeringEndpoint.Host)
}

func TestListSetsPublicEndpointReachabilityFromDNSProbe(t *testing.T) {
	prevCheckDNSPropagation := checkDNSPropagation
	t.Cleanup(func() { checkDNSPropagation = prevCheckDNSPropagation })
	var probedHosts []string
	checkDNSPropagation = func(_ context.Context, host string) (bool, string) {
		probedHosts = append(probedHosts, host)
		return true, ""
	}

	ctrl := gomock.NewController(t)
	clusterRepo := mock_cluster.NewMockRepo(ctrl)
	regionRepo := mock_region.NewMockRepo(ctrl)
	regionalCli := mock_regionalserver.NewMockClient(ctrl)
	svc := &serviceImpl{
		clusterRepo:   clusterRepo,
		regionRepo:    regionRepo,
		tidbgroupRepo: &stubTidbGroupRepo{},
		regionalCli:   regionalCli,
	}

	cluster := &entity.Cluster{
		Cluster: &globalngv1.Cluster{
			ClusterId: 1001,
		},
		RegionId: 2001,
	}
	region := &entity.Region{
		Region:   &globalcommonv1.Region{Id: globalcommonv1.Region_AWS_US_WEST_2},
		RegionID: 2001,
	}

	clusterRepo.EXPECT().Query(gomock.Any(), gomock.Any()).Return([]*entity.Cluster{cluster}, nil)
	regionRepo.EXPECT().ListRegions(gomock.Any(), gomock.Any()).Return([]*entity.Region{region}, nil)
	regionalCli.EXPECT().GetTidbGroup(gomock.Any(), region, gomock.Any()).Return(&clusterngv1beta.TiDBGroup{
		Id:            "3001",
		ClusterId:     "1001",
		DataplaneName: "dp-1",
		SqlEndpoint: &clusterngv1beta.TiDBGroup_SqlEndpoint{
			PublicLb: &clusterngv1beta.TiDBGroup_SqlEndpoint_LoadBalancer{
				Enabled: true,
				Host:    "tidb.example.com",
				Port:    4000,
			},
		},
	}, nil)
	regionalCli.EXPECT().GetDataPlane(gomock.Any(), region, &dataplanemanagerv1.GetDataPlaneRequest{
		DataPlaneId: "dp-1",
	}).Return(&dataplanemanagerv1.DataPlane{
		Vpc: &dataplanemanagerv1.DataPlane_Vpc{
			NatIps: []string{"1.1.1.1", "2.2.2.2"},
		},
	}, nil)

	clusters, err := svc.List(t.Context(), &globalngv1.ListClustersRequest{ClusterIds: []uint64{1001}})

	require.Nil(t, err)
	require.Len(t, clusters, 1)
	require.Equal(t, []string{"tidb.example.com"}, probedHosts)
	require.Len(t, clusters[0].Endpoints, 1)
	assert.Equal(t, []string{"1.1.1.1", "2.2.2.2"}, clusters[0].Endpoints[0].GetPublicEndpoint().GetEgressIps())
	reachability := clusters[0].Endpoints[0].GetPublicEndpoint().GetReachability()
	require.NotNil(t, reachability)
	assert.True(t, reachability.GetReachable())
	require.NotNil(t, reachability.GetDetail())
	assert.True(t, reachability.GetDetail().GetServiceActive())
	assert.True(t, reachability.GetDetail().GetEndpointActive())
	assert.True(t, reachability.GetDetail().GetDnsReachable())
	assert.Empty(t, reachability.GetDetail().GetMessage())
}

func TestFullPublicEndpointEgressIPsSetsOnlyPublicEndpoints(t *testing.T) {
	ctrl := gomock.NewController(t)
	regionalCli := mock_regionalserver.NewMockClient(ctrl)
	svc := &serviceImpl{regionalCli: regionalCli}
	region := &entity.Region{RegionID: 2001}
	cluster := &entity.Cluster{
		Cluster: &globalngv1.Cluster{
			DataplaneId: "dp-1",
			Endpoints: []*globalngv1.Endpoint{
				{PublicEndpoint: &globalngv1.Endpoint_PublicEndpoint{Host: "public.example.com", Port: 4000}},
				{VpcPeeringEndpoint: &globalngv1.Endpoint_VpcPeeringEndpoint{Host: "private.example.com", Port: 4000}},
			},
		},
		RegionId: 2001,
	}

	regionalCli.EXPECT().GetDataPlane(gomock.Any(), region, &dataplanemanagerv1.GetDataPlaneRequest{
		DataPlaneId: "dp-1",
	}).Return(&dataplanemanagerv1.DataPlane{
		Vpc: &dataplanemanagerv1.DataPlane_Vpc{NatIps: []string{"3.3.3.3"}},
	}, nil)

	svc.fullPublicEndpointEgressIPs(t.Context(), []*entity.Region{region}, []*entity.Cluster{cluster})

	assert.Equal(t, []string{"3.3.3.3"}, cluster.Endpoints[0].GetPublicEndpoint().GetEgressIps())
	assert.Nil(t, cluster.Endpoints[1].GetPublicEndpoint())
}

func TestFullConnectionReachabilitySetsPublicEndpointFromDNSProbe(t *testing.T) {
	prevCheckDNSPropagation := checkDNSPropagation
	t.Cleanup(func() { checkDNSPropagation = prevCheckDNSPropagation })
	var probedHosts []string
	checkDNSPropagation = func(_ context.Context, host string) (bool, string) {
		probedHosts = append(probedHosts, host)
		return false, "DNS_NOT_REACHABLE"
	}

	cluster := &entity.Cluster{
		Cluster: &globalngv1.Cluster{
			Endpoints: []*globalngv1.Endpoint{
				{
					PublicEndpoint: &globalngv1.Endpoint_PublicEndpoint{
						Host: "tidb.example.com",
						Port: 4000,
					},
				},
			},
		},
	}

	(&serviceImpl{}).fullConnectionReachability(t.Context(), cluster)

	require.Equal(t, []string{"tidb.example.com"}, probedHosts)
	reachability := cluster.Endpoints[0].GetPublicEndpoint().GetReachability()
	require.NotNil(t, reachability)
	assert.False(t, reachability.GetReachable())
	require.NotNil(t, reachability.GetDetail())
	assert.True(t, reachability.GetDetail().GetServiceActive())
	assert.True(t, reachability.GetDetail().GetEndpointActive())
	assert.False(t, reachability.GetDetail().GetDnsReachable())
	assert.Equal(t, "DNS_NOT_REACHABLE", reachability.GetDetail().GetMessage())
}

func TestFullConnectionReachabilityMarksPrivateEndpointsReachable(t *testing.T) {
	cluster := &entity.Cluster{
		Cluster: &globalngv1.Cluster{
			Endpoints: []*globalngv1.Endpoint{
				{
					PrivateLinkEndpoints: []*globalngv1.Endpoint_PrivateLinkEndpoint{
						{Host: "privatelink.example.com", Port: 4000},
					},
				},
				{
					VpcPeeringEndpoint: &globalngv1.Endpoint_VpcPeeringEndpoint{
						Host: "vpc.example.com",
						Port: 4000,
					},
				},
			},
		},
	}

	(&serviceImpl{}).fullConnectionReachability(t.Context(), cluster)

	privateLinkReachability := cluster.Endpoints[0].GetPrivateLinkEndpoints()[0].GetReachability()
	require.NotNil(t, privateLinkReachability)
	assert.True(t, privateLinkReachability.GetReachable())
	assert.True(t, privateLinkReachability.GetDetail().GetServiceActive())
	assert.True(t, privateLinkReachability.GetDetail().GetEndpointActive())
	assert.True(t, privateLinkReachability.GetDetail().GetDnsReachable())
	assert.Empty(t, privateLinkReachability.GetDetail().GetMessage())

	vpcPeeringReachability := cluster.Endpoints[1].GetVpcPeeringEndpoint().GetReachability()
	require.NotNil(t, vpcPeeringReachability)
	assert.True(t, vpcPeeringReachability.GetReachable())
	assert.True(t, vpcPeeringReachability.GetDetail().GetServiceActive())
	assert.True(t, vpcPeeringReachability.GetDetail().GetEndpointActive())
	assert.True(t, vpcPeeringReachability.GetDetail().GetDnsReachable())
	assert.Empty(t, vpcPeeringReachability.GetDetail().GetMessage())
}

func TestFullClusterInfoDerivesEssentialHighAvailabilityTypeFromSharedPool(t *testing.T) {
	tests := []struct {
		name                string
		servicePlan         globalcommonv1.ServicePlan
		initialType         globalngv1.Cluster_HighAvailabilityType
		sharedPoolNodeZones []string
		sharedPoolErr       bizerr.BizError
		want                globalngv1.Cluster_HighAvailabilityType
		expectSharedPools   bool
	}{
		{
			name:                "essential v2 with single node zone is zonal",
			servicePlan:         globalcommonv1.ServicePlan_Essential_V2,
			sharedPoolNodeZones: []string{"us-west-2a"},
			want:                globalngv1.Cluster_ZONAL,
			expectSharedPools:   true,
		},
		{
			name:                "essential v2 with multiple node zones is regional",
			servicePlan:         globalcommonv1.ServicePlan_Essential_V2,
			sharedPoolNodeZones: []string{"us-west-2a", "us-west-2b", "us-west-2c"},
			want:                globalngv1.Cluster_REGIONAL,
			expectSharedPools:   true,
		},
		{
			name:              "essential v2 without node zones is regional",
			servicePlan:       globalcommonv1.ServicePlan_Essential_V2,
			want:              globalngv1.Cluster_REGIONAL,
			expectSharedPools: true,
		},
		{
			name:                "premium keeps stored regional even with node zones",
			servicePlan:         globalcommonv1.ServicePlan_Premium,
			initialType:         globalngv1.Cluster_REGIONAL,
			sharedPoolNodeZones: []string{"us-west-2a"},
			want:                globalngv1.Cluster_REGIONAL,
			expectSharedPools:   false,
		},
		{
			name:              "shared pool lookup failure preserves stored value",
			servicePlan:       globalcommonv1.ServicePlan_Essential_V2,
			initialType:       globalngv1.Cluster_ZONAL,
			sharedPoolErr:     bizerr.New(bizerr.ErrCodeRegionalAPI, "shared pool unavailable"),
			want:              globalngv1.Cluster_ZONAL,
			expectSharedPools: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ctx := t.Context()
			ctrl := gomock.NewController(t)
			regionRepo := mock_region.NewMockRepo(ctrl)
			regionalCli := mock_regionalserver.NewMockClient(ctrl)
			svc := &serviceImpl{
				regionRepo:    regionRepo,
				tidbgroupRepo: &stubTidbGroupRepo{},
				regionalCli:   regionalCli,
			}

			cluster := &entity.Cluster{
				Cluster: &globalngv1.Cluster{
					ClusterId:            1001,
					ServicePlan:          tt.servicePlan,
					HighAvailabilityType: tt.initialType,
				},
				RegionId: 2001,
			}
			region := &entity.Region{
				Region:   &globalcommonv1.Region{Id: globalcommonv1.Region_AWS_US_WEST_2},
				RegionID: 2001,
			}

			regionRepo.EXPECT().ListRegions(gomock.Any(), gomock.Any()).Return([]*entity.Region{region}, nil)
			regionalCli.EXPECT().GetTidbGroup(gomock.Any(), region, gomock.Any()).Return(&clusterngv1beta.TiDBGroup{
				Id:           "3001",
				ClusterId:    "1001",
				SharedPoolId: "5001",
			}, nil)
			if tt.expectSharedPools {
				regionalCli.EXPECT().ListSharedPools(gomock.Any(), region, gomock.Any()).Return(&clusterngv1beta.ListSharedPoolsResponse{
					SharedPools: []*clusterngv1beta.SharedPool{
						{
							Id: "5001",
							Provider: &clusterngv1beta.Provider{
								NodeZones: tt.sharedPoolNodeZones,
							},
						},
					},
				}, tt.sharedPoolErr)
			}

			err := svc.fullClusterInfo(ctx, []*entity.Cluster{cluster}, false)

			require.Nil(t, err)
			assert.Equal(t, tt.want, cluster.GetHighAvailabilityType())
		})
	}
}

func TestFullClusterHighAvailabilityTypesListsSharedPoolsByRegion(t *testing.T) {
	ctx := t.Context()
	ctrl := gomock.NewController(t)
	regionalCli := mock_regionalserver.NewMockClient(ctrl)
	svc := &serviceImpl{
		regionalCli: regionalCli,
	}

	region := &entity.Region{
		Region:   &globalcommonv1.Region{Id: globalcommonv1.Region_AWS_US_WEST_2},
		RegionID: 2001,
	}
	clusters := []*entity.Cluster{
		{
			Cluster: &globalngv1.Cluster{
				ClusterId:    1001,
				ServicePlan:  globalcommonv1.ServicePlan_Essential_V2,
				SharedPoolId: "5001",
			},
			RegionId: 2001,
		},
		{
			Cluster: &globalngv1.Cluster{
				ClusterId:    1002,
				ServicePlan:  globalcommonv1.ServicePlan_Essential_V2,
				SharedPoolId: "5001",
			},
			RegionId: 2001,
		},
		{
			Cluster: &globalngv1.Cluster{
				ClusterId:    1003,
				ServicePlan:  globalcommonv1.ServicePlan_Essential_V2,
				SharedPoolId: "5002",
			},
			RegionId: 2001,
		},
	}

	regionalCli.EXPECT().
		ListSharedPools(gomock.Any(), region, gomock.Any()).
		DoAndReturn(func(
			_ context.Context,
			_ *entity.Region,
			req *clusterngv1beta.ListSharedPoolsRequest,
		) (*clusterngv1beta.ListSharedPoolsResponse, bizerr.BizError) {
			assert.ElementsMatch(t, []string{"5001", "5002"}, req.GetIds())
			return &clusterngv1beta.ListSharedPoolsResponse{
				SharedPools: []*clusterngv1beta.SharedPool{
					{
						Id: "5001",
						Provider: &clusterngv1beta.Provider{
							NodeZones: []string{"us-west-2a"},
						},
					},
					{
						Id: "5002",
						Provider: &clusterngv1beta.Provider{
							NodeZones: []string{"us-west-2a", "us-west-2b", "us-west-2c"},
						},
					},
				},
			}, nil
		}).
		Times(1)

	svc.fullClusterHighAvailabilityTypes(ctx, clusters, map[uint64]*entity.Region{
		2001: region,
	})

	assert.Equal(t, globalngv1.Cluster_ZONAL, clusters[0].GetHighAvailabilityType())
	assert.Equal(t, globalngv1.Cluster_ZONAL, clusters[1].GetHighAvailabilityType())
	assert.Equal(t, globalngv1.Cluster_REGIONAL, clusters[2].GetHighAvailabilityType())
}

func TestFullClusterDualLayerEncryptionUsesProviderIdentityId(t *testing.T) {
	ctx := t.Context()
	ctrl := gomock.NewController(t)
	providerIdentityRepo := mock_region.NewMockProviderIdentityRepo(ctrl)
	key := awsCmekEncryptionKey(10, 111)
	key.ProviderId = 999
	identity := &entity.RegionProviderIdentity{
		RegionProviderIdentity: &globalngv1.RegionProviderIdentity{
			Id:       111,
			Identity: "123456789012",
		},
	}
	encryptionKeySvc := &fakeEncryptionKeyService{keys: []*entity.EncryptionKey{key}}
	svc := &serviceImpl{
		encryptionKeySvc:           encryptionKeySvc,
		regionProviderIdentityRepo: providerIdentityRepo,
	}
	providerIdentityRepo.EXPECT().Load(gomock.Any(), uint64(111)).Return(identity, nil)

	cluster := &entity.Cluster{Cluster: &globalngv1.Cluster{ClusterId: 1}, EncryptionKeyId: 10}
	err := svc.fullClusterDualLayerEncryption(ctx, cluster)

	require.Nil(t, err)
	assert.Equal(t, 1, encryptionKeySvc.listCalls)
	assert.Equal(t, "60205", cluster.DualLayerDataEncryption.GetCmek().GetAwsPrincipal().GetExternalId())
}

func TestFullClustersDualLayerEncryptionBatchesEncryptionKeysAndCachesProviderIdentities(t *testing.T) {
	ctx := t.Context()
	ctrl := gomock.NewController(t)
	providerIdentityRepo := mock_region.NewMockProviderIdentityRepo(ctrl)
	identity := &entity.RegionProviderIdentity{
		RegionProviderIdentity: &globalngv1.RegionProviderIdentity{
			Id:       111,
			Identity: "123456789012",
		},
	}
	encryptionKeySvc := &fakeEncryptionKeyService{
		keys: []*entity.EncryptionKey{
			awsCmekEncryptionKey(10, 111),
			{
				Id:              20,
				CustomerManaged: false,
			},
		},
	}
	svc := &serviceImpl{
		encryptionKeySvc:           encryptionKeySvc,
		regionProviderIdentityRepo: providerIdentityRepo,
	}
	providerIdentityRepo.EXPECT().Load(gomock.Any(), uint64(111)).Return(identity, nil).Times(1)
	clusters := []*entity.Cluster{
		{Cluster: &globalngv1.Cluster{ClusterId: 1}},
		{Cluster: &globalngv1.Cluster{ClusterId: 2}, EncryptionKeyId: 10},
		{Cluster: &globalngv1.Cluster{ClusterId: 3}, EncryptionKeyId: 10},
		{Cluster: &globalngv1.Cluster{ClusterId: 4}, EncryptionKeyId: 20},
	}

	err := svc.fullClustersDualLayerEncryption(ctx, clusters)

	require.Nil(t, err)
	assert.Equal(t, 1, encryptionKeySvc.listCalls)
	assert.ElementsMatch(t, []uint64{10, 20}, encryptionKeySvc.listIds)
	assert.Equal(t, "60205", clusters[1].DualLayerDataEncryption.GetCmek().GetAwsPrincipal().GetExternalId())
	assert.Equal(t, "60205", clusters[2].DualLayerDataEncryption.GetCmek().GetAwsPrincipal().GetExternalId())
	assert.IsType(t, &globalngv1.Cluster_DualLayerDataEncryption_NoEncryption{}, clusters[0].DualLayerDataEncryption.GetKey())
	assert.IsType(t, &globalngv1.Cluster_DualLayerDataEncryption_DefaultKey{}, clusters[3].DualLayerDataEncryption.GetKey())
}

func awsCmekEncryptionKey(id, providerIdentityId uint64) *entity.EncryptionKey {
	return &entity.EncryptionKey{
		Id:                 id,
		TenantId:           60205,
		ProviderName:       globalcommonv1.CloudProvider_AWS,
		RegionName:         globalcommonv1.Region_AWS_US_WEST_2,
		ServicePlan:        globalcommonv1.ServicePlan_Premium,
		ProviderIdentityId: providerIdentityId,
		CustomerManaged:    true,
		MasterKey:          "arn:aws:kms:us-west-2:123456789012:key/test-key",
	}
}

type fakeEncryptionKeyService struct {
	encryptionkeysvc.Service

	keys      []*entity.EncryptionKey
	listCalls int
	listIds   []uint64
}

func (f *fakeEncryptionKeyService) ListEncryptionByIds(
	_ context.Context,
	encryptionIds []uint64,
) ([]*entity.EncryptionKey, bizerr.BizError) {
	f.listCalls++
	f.listIds = append([]uint64(nil), encryptionIds...)
	return f.keys, nil
}

func (f *fakeEncryptionKeyService) CvtEncryptionKeyToCmek(
	_ context.Context,
	key *entity.EncryptionKey,
	providerIdentity *entity.RegionProviderIdentity,
) (*globalngv1.CustomerManagedEncryptionKey, bizerr.BizError) {
	return &globalngv1.CustomerManagedEncryptionKey{
		CustomerManagedEncryptionKeyId: key.Id,
		OrganizationId:                 key.TenantId,
		CloudProviderId:                key.ProviderName,
		RegionId:                       key.RegionName,
		ServicePlan:                    key.ServicePlan,
		Key: &globalngv1.CustomerManagedEncryptionKey_AwsKms_{
			AwsKms: &globalngv1.CustomerManagedEncryptionKey_AwsKms{
				KmsKeyArn: key.MasterKey,
			},
		},
		Principal: &globalngv1.CustomerManagedEncryptionKey_AwsPrincipal{
			AwsPrincipal: &globalngv1.CmekAccessIamPrincipal_AwsCmekPrincipal{
				AccountId:  providerIdentity.GetIdentity(),
				ExternalId: "60205",
			},
		},
	}, nil
}

var _ encryptionkeysvc.Service = (*fakeEncryptionKeyService)(nil)
var _ regionrepo.ProviderIdentityRepo = (*mock_region.MockProviderIdentityRepo)(nil)

type stubTidbGroupRepo struct{}

func (s *stubTidbGroupRepo) Create(context.Context, *entity.TidbGroup) bizerr.BizError {
	panic("unexpected call")
}

func (s *stubTidbGroupRepo) Save(context.Context, *entity.TidbGroup) bizerr.BizError {
	panic("unexpected call")
}

func (s *stubTidbGroupRepo) Load(context.Context, uint64, uint64) (*entity.TidbGroup, bizerr.BizError) {
	panic("unexpected call")
}

func (s *stubTidbGroupRepo) Delete(context.Context, *entity.TidbGroup) bizerr.BizError {
	panic("unexpected call")
}

func (s *stubTidbGroupRepo) DeleteByClusterId(context.Context, uint64) bizerr.BizError {
	panic("unexpected call")
}

func (s *stubTidbGroupRepo) Query(context.Context, *tidbgrouprepo.Filter) ([]*entity.TidbGroup, bizerr.BizError) {
	return []*entity.TidbGroup{
		{Id: 3001, ClusterId: 1001},
	}, nil
}

var _ tidbgrouprepo.Repo = (*stubTidbGroupRepo)(nil)