package app

import (
	"context"
	"errors"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	"github.com/tidbcloud/cluster-service-ng/domain/server"
	"github.com/tidbcloud/cluster-service-ng/domain/server/br"
	"github.com/tidbcloud/cluster-service-ng/domain/server/mocks"
	"github.com/tidbcloud/cluster-service-ng/pkg/clustergateway"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/emptypb"
)

func TestCreateSharedPoolPassesVcpusLimit(t *testing.T) {
	ctx := context.Background()
	mockSvc := mocks.NewService(t)
	app := &appImpl{ngService: mockSvc}

	req := &apiv1.CreateSharedPoolRequest{
		SharedPool:           &apiv1.SharedPool{Id: "100"},
		InitSharedPoolConfig: &apiv1.CreateSharedPoolRequest_InitSharedPoolConfig{},
		VcpusLimit:           106,
	}

	mockSvc.On("CreateSharedPool", ctx, mock.MatchedBy(func(sharedPool *entity.SharedPool) bool {
		return sharedPool.GetId() == "100" && sharedPool.GetVcpusLimit() == 106
	}), (*entity.RequestUnit)(nil), req.GetInitSharedPoolConfig()).Return(&entity.SharedPool{
		SharedPool: apiv1.SharedPool{
			Id:              "100",
			VcpusLimit:      106,
			ObjectStorage:   &apiv1.SharedPool_ObjectStorage{},
			InternalStorage: &apiv1.SharedPool_ObjectStorage{},
		},
	}, nil).Once()

	resp, err := app.CreateSharedPool(ctx, req)
	require.NoError(t, err)
	assert.EqualValues(t, 106, resp.GetVcpusLimit())
	mockSvc.AssertExpectations(t)
}

func TestCreateSharedPoolPassesEnableCapacityPlanner(t *testing.T) {
	ctx := context.Background()
	mockSvc := mocks.NewService(t)
	app := &appImpl{ngService: mockSvc}
	req := &apiv1.CreateSharedPoolRequest{
		SharedPool:            &apiv1.SharedPool{Id: "100"},
		InitSharedPoolConfig:  &apiv1.CreateSharedPoolRequest_InitSharedPoolConfig{},
		CapacityMode:          apiv1.CapacityMode_RCU_DRIVEN,
		EnableCapacityPlanner: true,
	}

	mockSvc.On("CreateSharedPool", ctx, mock.MatchedBy(func(sharedPool *entity.SharedPool) bool {
		return sharedPool.GetId() == "100" && sharedPool.EnableCapacityPlanner
	}), (*entity.RequestUnit)(nil), req.GetInitSharedPoolConfig()).Return(&entity.SharedPool{
		SharedPool: apiv1.SharedPool{
			Id:              "100",
			ObjectStorage:   &apiv1.SharedPool_ObjectStorage{},
			InternalStorage: &apiv1.SharedPool_ObjectStorage{},
		},
	}, nil).Once()

	_, err := app.CreateSharedPool(ctx, req)
	require.NoError(t, err)
	mockSvc.AssertExpectations(t)
}

func TestCreateSharedPoolRejectsCapacityPlannerForTopologyDrivenPool(t *testing.T) {
	mockSvc := mocks.NewService(t)
	app := &appImpl{ngService: mockSvc}

	_, err := app.CreateSharedPool(context.Background(), &apiv1.CreateSharedPoolRequest{
		SharedPool:            &apiv1.SharedPool{Id: "100"},
		InitSharedPoolConfig:  &apiv1.CreateSharedPoolRequest_InitSharedPoolConfig{},
		CapacityMode:          apiv1.CapacityMode_TOPOLOGY_DRIVEN,
		EnableCapacityPlanner: true,
	})

	require.Equal(t, codes.InvalidArgument, status.Code(err))
	mockSvc.AssertNotCalled(t, "CreateSharedPool", mock.Anything, mock.Anything, mock.Anything)
}

func TestValidateSharedPoolPauseReturnsStructuredFailures(t *testing.T) {
	ctx := context.Background()
	mockSvc := mocks.NewService(t)
	app := &appImpl{ngService: mockSvc}
	want := &entity.ValidateSharedPoolPauseResponse{Failures: []*entity.SharedPoolPauseValidationFailure{
		{Code: "RESOURCE_MUTATION_IN_PROGRESS", Message: "a conflicting resource mutation is in progress"},
	}}
	mockSvc.EXPECT().ValidateSharedPoolPause(ctx, "100").Return(want, nil).Once()

	response, err := app.ValidateSharedPoolPause(ctx, &apiv1.ValidateSharedPoolPauseRequest{SharedPoolId: "100"})
	require.NoError(t, err)
	require.Same(t, want, response)
}

func TestCreateExportSnapshotTSO(t *testing.T) {
	ctx := context.Background()

	t.Run("passes a specified snapshot TSO to the domain service", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}
		const snapshotTSO = uint64(448266862624768042)
		req := &apiv1.CreateExportRequest{
			ClusterId: "100",
			Export: &apiv1.Export{
				ExportId:    "export-1",
				SnapshotTso: snapshotTSO,
			},
		}

		mockSvc.EXPECT().
			CreateExport(ctx, mock.MatchedBy(func(export *entity.Export) bool {
				return export.GetClusterId() == "100" && export.GetSnapshotTso() == snapshotTSO
			})).
			Return(req.Export, nil).
			Once()

		_, err := app.CreateExport(ctx, req)
		require.NoError(t, err)
	})

	t.Run("keeps the legacy zero value when snapshot TSO is omitted", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}
		req := &apiv1.CreateExportRequest{
			ClusterId: "100",
			Export:    &apiv1.Export{ExportId: "export-1"},
		}

		mockSvc.EXPECT().
			CreateExport(ctx, mock.MatchedBy(func(export *entity.Export) bool {
				return export.GetClusterId() == "100" && export.GetSnapshotTso() == 0
			})).
			Return(req.Export, nil).
			Once()

		_, err := app.CreateExport(ctx, req)
		require.NoError(t, err)
	})

}

func TestCreateSharedPoolPassesEnabledTiKVRemoteWorkers(t *testing.T) {
	ctx := context.Background()
	mockSvc := mocks.NewService(t)
	app := &appImpl{ngService: mockSvc}

	req := &apiv1.CreateSharedPoolRequest{
		SharedPool:           &apiv1.SharedPool{Id: "100"},
		InitSharedPoolConfig: &apiv1.CreateSharedPoolRequest_InitSharedPoolConfig{},
		EnabledTikvRemoteWorkers: []apiv1.ComponentType{
			apiv1.ComponentType_TIKV_WORKER,
		},
	}

	mockSvc.On("CreateSharedPool", ctx, mock.MatchedBy(func(sharedPool *entity.SharedPool) bool {
		return sharedPool.GetId() == "100" &&
			assert.ObjectsAreEqual(
				[]apiv1.ComponentType{apiv1.ComponentType_TIKV_WORKER},
				sharedPool.GetEnabledTikvRemoteWorkers(),
			)
	}), (*entity.RequestUnit)(nil), req.GetInitSharedPoolConfig()).Return(&entity.SharedPool{
		SharedPool: apiv1.SharedPool{
			Id: "100",
			EnabledTikvRemoteWorkers: []apiv1.ComponentType{
				apiv1.ComponentType_TIKV_WORKER,
			},
			ObjectStorage:   &apiv1.SharedPool_ObjectStorage{},
			InternalStorage: &apiv1.SharedPool_ObjectStorage{},
		},
	}, nil).Once()

	resp, err := app.CreateSharedPool(ctx, req)
	require.NoError(t, err)
	assert.Equal(
		t,
		[]apiv1.ComponentType{apiv1.ComponentType_TIKV_WORKER},
		resp.GetEnabledTikvRemoteWorkers(),
	)
	mockSvc.AssertExpectations(t)
}

func TestCreateSharedPoolPassesRequestedTiflashMode(t *testing.T) {
	for _, requestedMode := range []apiv1.TiflashMode{
		apiv1.TiflashMode_TIFLASH_MODE_UNSPECIFIED,
		apiv1.TiflashMode_TIFLASH_MODE_AUTO,
		apiv1.TiflashMode_TIFLASH_MODE_LEGACY,
		apiv1.TiflashMode_TIFLASH_MODE_COLUMNAR,
	} {
		t.Run(requestedMode.String(), func(t *testing.T) {
			ctx := context.Background()
			mockSvc := mocks.NewService(t)
			app := &appImpl{ngService: mockSvc}
			req := &apiv1.CreateSharedPoolRequest{
				SharedPool:           &apiv1.SharedPool{Id: "100"},
				InitSharedPoolConfig: &apiv1.CreateSharedPoolRequest_InitSharedPoolConfig{},
				TiflashMode:          requestedMode,
			}

			mockSvc.On("CreateSharedPool", ctx, mock.MatchedBy(func(sharedPool *entity.SharedPool) bool {
				return sharedPool.RequestedTiflashMode == requestedMode
			}), (*entity.RequestUnit)(nil), req.GetInitSharedPoolConfig()).Return(&entity.SharedPool{
				SharedPool: apiv1.SharedPool{
					Id:              "100",
					ObjectStorage:   &apiv1.SharedPool_ObjectStorage{},
					InternalStorage: &apiv1.SharedPool_ObjectStorage{},
				},
			}, nil).Once()

			_, err := app.CreateSharedPool(ctx, req)
			require.NoError(t, err)
			mockSvc.AssertExpectations(t)
		})
	}
}

func TestCreateSharedPoolForwardsResourcePoolFields(t *testing.T) {
	ctx := context.Background()
	mockSvc := mocks.NewService(t)
	app := &appImpl{ngService: mockSvc}
	tikvSpec := &apiv1.NodeSpec{Cpu: 8, MemoryGi: 32}
	computeSpec := &apiv1.NodeSpec{Cpu: 16, MemoryGi: 64}
	req := &apiv1.CreateSharedPoolRequest{
		SharedPool:           &apiv1.SharedPool{Id: "100"},
		InitSharedPoolConfig: &apiv1.CreateSharedPoolRequest_InitSharedPoolConfig{},
		BootstrapCapacity: &apiv1.SharedPoolBootstrapCapacity{
			Tikv:           &apiv1.BootstrapComponentCapacity{NodeSpec: tikvSpec, NodeCount: 6},
			TiflashCompute: &apiv1.BootstrapComponentCapacity{NodeSpec: computeSpec, NodeCount: 1},
		},
		EnforceO11YNetworkNonOverlap: true,
	}

	mockSvc.On(
		"CreateSharedPool",
		ctx,
		mock.MatchedBy(func(sharedPool *entity.SharedPool) bool { return sharedPool.GetId() == "100" }),
		(*entity.RequestUnit)(nil),
		req.GetInitSharedPoolConfig(),
		mock.MatchedBy(func(options []*server.CreateSharedPoolOptions) bool {
			if len(options) != 1 || options[0] == nil || options[0].BootstrapCapacity == nil {
				return false
			}
			capacity := options[0].BootstrapCapacity
			return options[0].EnforceO11yNetworkNonOverlap &&
				capacity.TiKV.NodeSpec == tikvSpec && capacity.TiKV.NodeCount == 6 &&
				capacity.TiFlashCompute.NodeSpec == computeSpec && capacity.TiFlashCompute.NodeCount == 1
		}),
	).Return(&entity.SharedPool{SharedPool: apiv1.SharedPool{
		Id:              "100",
		ObjectStorage:   &apiv1.SharedPool_ObjectStorage{},
		InternalStorage: &apiv1.SharedPool_ObjectStorage{},
	}}, nil).Once()

	resp, err := app.CreateSharedPool(ctx, req)
	require.NoError(t, err)
	assert.Equal(t, "100", resp.GetId())
}

// TestCreateSharedPoolDerivesProviderPlacementGroupEnabledFromInitConfigNodeZones is
// the regression test for the public ingress gap (PR #3040 review): the registered
// app adapter constructs entity.SharedPool without Provider/NodeZones (those are
// OUTPUT_ONLY), so the durable Provider.PlacementGroupEnabled gate must be derived
// from initConfig.NodeZones (the canonical public input), not from the empty
// Provider.NodeZones that would silently pin a single-AZ AWS request to false.
// User-supplied SharedPool.placement_group_enabled is accepted but ignored: PR #3040
// makes the "is AWS single-AZ" check the source of truth.
func TestCreateSharedPoolDerivesProviderPlacementGroupEnabledFromInitConfigNodeZones(t *testing.T) {
	ctx := context.Background()

	tests := []struct {
		name          string
		userOptIn     bool // what the caller puts in SharedPool.placement_group_enabled
		nodeZones     []string
		wantPlacement bool
	}{
		{
			name:          "AWS single-AZ without user opt-in still enables PG",
			userOptIn:     false,
			nodeZones:     []string{"us-west-2a"},
			wantPlacement: true,
		},
		{
			name:          "AWS single-AZ with user opt-in still enables PG",
			userOptIn:     true,
			nodeZones:     []string{"us-west-2a"},
			wantPlacement: true,
		},
		{
			name:          "AWS multi-AZ leaves PG off even with user opt-in (ignored)",
			userOptIn:     true,
			nodeZones:     []string{"us-west-2a", "us-west-2b", "us-west-2c"},
			wantPlacement: false,
		},
		{
			name:          "empty node zones leaves PG off",
			userOptIn:     false,
			nodeZones:     nil,
			wantPlacement: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			mockSvc := mocks.NewService(t)
			app := &appImpl{ngService: mockSvc}

			req := &apiv1.CreateSharedPoolRequest{
				SharedPool: &apiv1.SharedPool{
					Id:                    "100",
					PlacementGroupEnabled: tt.userOptIn,
				},
				InitSharedPoolConfig: &apiv1.CreateSharedPoolRequest_InitSharedPoolConfig{
					NetworkCidr: "10.0.0.0/16",
					NodeZones:   tt.nodeZones,
				},
				EnforceO11YNetworkNonOverlap: true,
			}

			mockSvc.On(
				"CreateSharedPool",
				ctx,
				mock.MatchedBy(func(sharedPool *entity.SharedPool) bool {
					// Provider should be left nil at the appImpl ingress (it is
					// OUTPUT_ONLY on the public request); serverImpl.CreateSharedPool
					// is responsible for populating it after deriving the gate.
					return sharedPool.GetProvider() == nil &&
						sharedPool.GetPlacementGroupEnabled() == tt.userOptIn
				}),
				(*entity.RequestUnit)(nil),
				req.GetInitSharedPoolConfig(),
				mock.Anything,
			).Return(&entity.SharedPool{SharedPool: apiv1.SharedPool{
				Id: "100",
				Provider: &apiv1.Provider{
					Type:                  apiv1.ProviderType_AWS,
					Region:                "us-west-2",
					PlacementGroupEnabled: tt.wantPlacement,
				},
				ObjectStorage:   &apiv1.SharedPool_ObjectStorage{},
				InternalStorage: &apiv1.SharedPool_ObjectStorage{},
			}}, nil).Once()

			resp, err := app.CreateSharedPool(ctx, req)
			require.NoError(t, err)
			require.NotNil(t, resp)
			assert.Equal(t, tt.wantPlacement, resp.GetProvider().GetPlacementGroupEnabled(),
				"public response must carry the server-derived gate value")
		})
	}
}

func TestResourcePoolManagementRPCsForward(t *testing.T) {
	ctx := context.Background()
	mockSvc := mocks.NewService(t)
	app := &appImpl{ngService: mockSvc}
	nodeSpecs := &entity.ListNodeSpecsResponse{NodeSpecOptions: []*apiv1.NodeSpecOptions{{
		NodeSpec: &apiv1.NodeSpec{Cpu: 8, MemoryGi: 32},
	}}}
	provision := &entity.SharedPoolProvisionConfiguration{
		O11YClusterCidr:    "10.0.0.0/16",
		AvailableNodeZones: []string{"us-west-2a", "us-west-2b", "us-west-2c"},
	}
	prefix := int32(20)
	capacityRequest := &apiv1.EvaluateSharedPoolNetworkCapacityRequest{
		CidrPrefixLength:     &prefix,
		HighAvailabilityType: apiv1.CreateSharedPoolRequest_InitSharedPoolConfig_ZONAL,
	}
	capacity := &entity.EvaluateSharedPoolNetworkCapacityResponse{Supported: true, SupportedNodeCount: 60}
	mockSvc.On("ListNodeSpecs", ctx, "tenant-1", apiv1.ComponentType_TIKV).Return(nodeSpecs, nil).Once()
	mockSvc.On("GetSharedPoolProvisionConfiguration", ctx).Return(provision, nil).Once()
	mockSvc.On("EvaluateSharedPoolNetworkCapacity", ctx, capacityRequest).Return(capacity, nil).Once()

	listed, err := app.ListNodeSpecs(ctx, &apiv1.ListNodeSpecsRequest{TenantId: "tenant-1", ComponentType: apiv1.ComponentType_TIKV})
	require.NoError(t, err)
	assert.Same(t, nodeSpecs, listed)

	configured, err := app.GetSharedPoolProvisionConfiguration(ctx, &apiv1.GetSharedPoolProvisionConfigurationRequest{})
	require.NoError(t, err)
	assert.Same(t, provision, configured)
	assert.Equal(t, []string{"us-west-2a", "us-west-2b", "us-west-2c"}, configured.GetAvailableNodeZones())

	evaluated, err := app.EvaluateSharedPoolNetworkCapacity(ctx, capacityRequest)
	require.NoError(t, err)
	assert.Same(t, capacity, evaluated)
}

func TestCreateSharedPoolRejectsResourceManagerWithMonolithicPD(t *testing.T) {
	ctx := context.Background()
	mockSvc := mocks.NewService(t)
	app := &appImpl{ngService: mockSvc}

	req := &apiv1.CreateSharedPoolRequest{
		PdDeploymentMode: apiv1.PDDeploymentMode_PD_DEPLOYMENT_MODE_MONOLITHIC,
		SharedPool: &apiv1.SharedPool{
			Id: "100",
			ComponentGroups: []*apiv1.Component{
				{ComponentType: apiv1.ComponentType_RESOURCE_MANAGER},
			},
		},
	}

	_, err := app.CreateSharedPool(ctx, req)
	require.Error(t, err)
	assert.Equal(t, codes.InvalidArgument, status.Code(err))
	mockSvc.AssertNotCalled(t, "CreateSharedPool", mock.Anything, mock.Anything, mock.Anything, mock.Anything)
}

func TestUpdateSharedPoolVcpusLimit(t *testing.T) {
	ctx := context.Background()
	mockSvc := mocks.NewService(t)
	app := &appImpl{ngService: mockSvc}

	req := &apiv1.UpdateSharedPoolVcpusLimitRequest{SharedPoolId: "100", VcpusLimit: 106}
	mockSvc.On("UpdateSharedPoolVcpusLimit", ctx, "100", uint64(106)).Return(&entity.SharedPool{
		SharedPool: apiv1.SharedPool{Id: "100", VcpusLimit: 106},
	}, nil).Once()

	resp, err := app.UpdateSharedPoolVcpusLimit(ctx, req)
	require.NoError(t, err)
	assert.EqualValues(t, 106, resp.GetVcpusLimit())
	mockSvc.AssertExpectations(t)
}

func TestUpdateSharedPoolCloudResourceTagsPassesRemovals(t *testing.T) {
	ctx := context.Background()
	mockSvc := mocks.NewService(t)
	app := &appImpl{ngService: mockSvc}
	req := &apiv1.UpdateSharedPoolCloudResourceTagsRequest{
		SharedPoolId:               "100",
		CloudResourceTags:          map[string]string{"owner": "platform"},
		RemoveCloudResourceTagKeys: []string{"cost-center"},
	}

	mockSvc.On(
		"UpdateSharedPoolCloudResourceTags",
		ctx,
		"100",
		req.GetCloudResourceTags(),
		req.GetRemoveCloudResourceTagKeys(),
	).Return(&entity.SharedPool{
		SharedPool: apiv1.SharedPool{
			Id:                "100",
			CloudResourceTags: map[string]string{"owner": "platform"},
		},
	}, nil).Once()

	resp, err := app.UpdateSharedPoolCloudResourceTags(ctx, req)
	require.NoError(t, err)
	assert.Equal(t, map[string]string{"owner": "platform"}, resp.GetCloudResourceTags())
	mockSvc.AssertExpectations(t)
}

func TestGetSharedPoolVcpusUsage(t *testing.T) {
	ctx := context.Background()
	mockSvc := mocks.NewService(t)
	app := &appImpl{ngService: mockSvc}

	req := &apiv1.GetSharedPoolVcpusUsageRequest{SharedPoolId: "100"}
	expected := &apiv1.GetSharedPoolVcpusUsageResponse{
		SharedPoolId: "100",
		VcpusUsed:    8,
		CountedNodes: 1,
		DroppedNodes: 2,
		DropReasons: &apiv1.SharedPoolVcpusUsageDropReasons{
			ConflictingOwnership: 1,
			MissingIdentity:      1,
		},
	}
	mockSvc.On("GetSharedPoolVcpusUsage", ctx, "100").Return(expected, nil).Once()

	resp, err := app.GetSharedPoolVcpusUsage(ctx, req)
	require.NoError(t, err)
	assert.Equal(t, expected, resp)
	mockSvc.AssertExpectations(t)
}

func TestScaleOperationHandlers(t *testing.T) {
	ctx := context.Background()
	mockSvc := mocks.NewService(t)
	app := &appImpl{ngService: mockSvc}
	operation := &entity.ScaleOperation{
		ScaleOperation: apiv1.ScaleOperation{
			SharedPoolId:     "100",
			ComponentId:      "tikv-1",
			ScaleOperationId: "op-1",
			Status:           apiv1.ScaleOperationStatus_SCALE_OPERATION_STATUS_RUNNING,
		},
	}

	mockSvc.On("GetScaleOperation", ctx, "100", "op-1").Return(operation, nil).Once()
	got, err := app.GetScaleOperation(ctx, &apiv1.GetScaleOperationRequest{
		SharedPoolId:     "100",
		ScaleOperationId: "op-1",
	})
	require.NoError(t, err)
	assert.Equal(t, "op-1", got.GetScaleOperationId())

	mockSvc.On(
		"CancelScaleOperation",
		ctx,
		"100",
		"op-1",
		"capacity is stuck",
	).Return(operation, nil).Once()
	canceled, err := app.CancelScaleOperation(ctx, &apiv1.CancelScaleOperationRequest{
		SharedPoolId:     "100",
		ScaleOperationId: "op-1",
		Reason:           "capacity is stuck",
	})
	require.NoError(t, err)
	assert.Equal(t, "op-1", canceled.GetScaleOperationId())
	mockSvc.AssertExpectations(t)
}

// A non-empty scale_operation_id is an optional, additive field. Until the
// guarded scale-out path is wired, it must be ignored and the mutation must run
// its normal legacy behavior, so setting it never breaks a working request.
func TestScaleOperationIdIgnoredRunsLegacyMutation(t *testing.T) {
	ctx := context.Background()

	t.Run("ScaleSharedPool", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}

		mockSvc.On("ScaleSharedPool", ctx, "100", mock.MatchedBy(func(comp *entity.Component) bool {
			return comp.GetComponentType() == apiv1.ComponentType_TIKV
		}), false).Return(&entity.SharedPool{SharedPool: apiv1.SharedPool{Id: "100"}}, nil).Once()

		resp, err := app.ScaleSharedPool(ctx, &apiv1.ScaleSharedPoolRequest{
			SharedPoolId:     "100",
			ScaleOperationId: "op-1",
			Component:        &apiv1.ScaleSharedPoolRequest_ScaleComponent{ComponentType: apiv1.ComponentType_TIKV},
		})

		require.NoError(t, err)
		require.NotNil(t, resp)
		mockSvc.AssertExpectations(t)
	})

	t.Run("SetSharedPoolComponentQuota", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}

		mockSvc.On("SetSharedPoolComponentQuota", ctx, "100", mock.MatchedBy(func(comp *entity.Component) bool {
			return comp.GetComponentType() == apiv1.ComponentType_TIKV
		})).Return(&entity.Component{ComponentType: apiv1.ComponentType_TIKV}, nil).Once()

		resp, err := app.SetSharedPoolComponentQuota(ctx, &apiv1.SetSharedPoolComponentQuotaRequest{
			SharedPoolId:     "100",
			ComponentType:    apiv1.ComponentType_TIKV,
			ScaleOperationId: "op-1",
		})

		require.NoError(t, err)
		require.NotNil(t, resp)
		mockSvc.AssertExpectations(t)
	})
}

func TestGetSharedPoolScaleConstraints(t *testing.T) {
	ctx := context.Background()
	mockSvc := mocks.NewService(t)
	app := &appImpl{ngService: mockSvc}

	req := &apiv1.GetSharedPoolScaleConstraintsRequest{SharedPoolId: "100"}
	expected := &apiv1.GetSharedPoolScaleConstraintsResponse{
		SharedPoolId: "100",
		Components: []*apiv1.SharedPoolScaleConstraintComponent{
			{
				ComponentId:   "default",
				ComponentType: apiv1.ComponentType_TIKV,
			},
		},
	}
	mockSvc.On("GetSharedPoolScaleConstraints", ctx, "100").Return(expected, nil).Once()

	resp, err := app.GetSharedPoolScaleConstraints(ctx, req)
	require.NoError(t, err)
	assert.Equal(t, expected, resp)
	mockSvc.AssertExpectations(t)
}

func TestGetImportRoleConfig(t *testing.T) {
	ctx := context.Background()

	t.Run("success", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}

		expected := &entity.ImportRoleConfig{
			RoleArn:         "arn:aws:iam::123456789012:role/import-role-123",
			ExternalId:      "123",
			Bucket:          "physical-internal-bucket",
			AllowedPrefix:   "internal/prefix/123",
			CloudStorageUri: "s3://physical-internal-bucket/internal/prefix/123",
		}
		mockSvc.On("GetImportRoleConfig", ctx, "123").Return(expected, nil).Once()

		resp, err := app.GetImportRoleConfig(ctx, &apiv1.GetImportRoleConfigRequest{ClusterId: "123"})

		require.NoError(t, err)
		assert.Equal(t, expected, resp)
		mockSvc.AssertExpectations(t)
	})

	t.Run("empty cluster id", func(t *testing.T) {
		app := &appImpl{}

		resp, err := app.GetImportRoleConfig(ctx, &apiv1.GetImportRoleConfigRequest{})

		require.Error(t, err)
		assert.Nil(t, resp)
		assert.Equal(t, codes.InvalidArgument, status.Code(err))
		assert.Contains(t, err.Error(), "cluster id is empty")
	})

	t.Run("service error", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}

		expectedErr := status.Error(codes.FailedPrecondition, "keyspace role is not ready")
		mockSvc.On("GetImportRoleConfig", ctx, "123").Return((*entity.ImportRoleConfig)(nil), expectedErr).Once()

		resp, err := app.GetImportRoleConfig(ctx, &apiv1.GetImportRoleConfigRequest{ClusterId: "123"})

		require.Error(t, err)
		assert.Nil(t, resp)
		assert.Equal(t, codes.FailedPrecondition, status.Code(err))
		assert.Contains(t, err.Error(), "keyspace role is not ready")
		mockSvc.AssertExpectations(t)
	})

}

func TestListLightweightBackups(t *testing.T) {
	ctx := context.Background()

	t.Run("success", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}

		req := &apiv1.ListLightweightBackupsRequest{}
		mockInput := &server.ListLightweightBackupsInput{
			ListLightweightBackupsRequest: req,
		}
		mockOutput := &server.ListLightweightBackupsOutput{
			DfsPrefix: "cse",
			DfsBucket: "s3://dummy-bucket",
			DfsRegion: "us-west-2",
			Keyspace2ClusterId: map[string]string{
				"keyspace-name": "cluster_id",
			},
			ListLightweightBackupsResponse: clustergateway.ListLightweightBackupsResponse{
				ClusterID: uint64(123),
				Data: []clustergateway.KernelLightweightBackupItem{
					{
						ID:             1,
						Name:           "backup-name",
						LastModifyTime: "2025-06-11 11:00:00.000",
						Details: clustergateway.BackupItemDetails{
							BackupTs: 111,
							SafeTs:   100,
							Keyspaces: []clustergateway.KeyspaceBackupInfo{
								{
									KeyspaceId:    222,
									KeyspaceName:  "keyspace-name",
									KeyspaceState: "ENABLED",
									Size:          12345,
								},
							},
						},
					},
				},
			},
		}

		mockSvc.On("ListLightweightBackups", ctx, mockInput).Return(mockOutput, nil)

		resp, err := app.ListLightweightBackups(ctx, req)

		require.NoError(t, err)
		assert.NotNil(t, resp)
		assert.Equal(t, "cse", resp.DfsPrefix)
		assert.Equal(t, "us-west-2", resp.DfsRegion)
		assert.Equal(t, "s3://dummy-bucket", resp.DfsBucket)
		assert.Equal(t, uint32(1), resp.TotalCount)
		assert.Len(t, resp.Backups, 1)
		assert.Equal(t, &apiv1.LightweightBackup{
			Id:         1,
			Name:       "backup-name",
			BackupTime: "2025-06-11 11:00:00.000",
			BackupTs:   111,
			ClusterBackupInfo: []*apiv1.ClusterBackupInfo{
				{
					ClusterId:       "cluster_id",
					BackupSizeBytes: 12345,
				},
			},
		}, resp.Backups[0])

		mockSvc.AssertExpectations(t)
	})

	t.Run("error from ngService", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}

		req := &apiv1.ListLightweightBackupsRequest{
			SharedPoolId: "123",
		}
		mockInput := &server.ListLightweightBackupsInput{
			ListLightweightBackupsRequest: req,
		}
		expectedErr := errors.New("service unavailable")

		mockSvc.On("ListLightweightBackups", ctx, mockInput).Return((*server.ListLightweightBackupsOutput)(nil), expectedErr)

		resp, err := app.ListLightweightBackups(ctx, req)

		assert.Nil(t, resp)
		require.Error(t, err)

		st, ok := status.FromError(err)
		assert.True(t, ok)
		assert.Equal(t, codes.Internal, st.Code())
		assert.Contains(t, st.Message(), "service unavailable")

		mockSvc.AssertExpectations(t)
	})
}

func TestCreateLightweightRestore(t *testing.T) {
	ctx := context.Background()

	t.Run("validation fails", func(t *testing.T) {
		app := &appImpl{}

		req := &apiv1.CreateLightweightRestoreRequest{
			KernelBackupId:   123,
			KernelBackupName: "dsds",
			PointInTime:      "dsds",
		}

		resp, err := app.CreateLightweightRestore(ctx, req)

		assert.Nil(t, resp)
		require.Error(t, err)

		st, ok := status.FromError(err)
		assert.True(t, ok)
		assert.Equal(t, codes.InvalidArgument, st.Code())
		assert.Contains(t, st.Message(), "validate request failed")
	})

	t.Run("ngService fails", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}

		req := &apiv1.CreateLightweightRestoreRequest{
			KernelBackupId:   123,
			KernelBackupName: "dsds",
		}

		mockInput := &br.CreateLightweightRestoreInput{
			CreateLightweightRestoreRequest: req,
		}

		mockSvc.On("CreateLightweightRestore", ctx, mockInput).Return(errors.New("create failed"))

		resp, err := app.CreateLightweightRestore(ctx, req)

		assert.Nil(t, resp)
		require.Error(t, err)
		st, ok := status.FromError(err)
		assert.True(t, ok)
		assert.Equal(t, codes.Internal, st.Code())
		assert.Contains(t, st.Message(), "create failed")

		mockSvc.AssertExpectations(t)
	})

	t.Run("preserves lifecycle guard status", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}
		req := &apiv1.CreateLightweightRestoreRequest{
			RestoreId:        "restore-1",
			KernelBackupId:   123,
			KernelBackupName: "snapshot-1",
		}
		mockSvc.On("CreateLightweightRestore", ctx, mock.Anything).
			Return(status.Error(codes.FailedPrecondition, "shared pool mutation is blocked")).
			Once()

		resp, err := app.CreateLightweightRestore(ctx, req)
		require.Error(t, err)
		require.Nil(t, resp)
		require.Equal(t, codes.FailedPrecondition, status.Code(err))
	})

	t.Run("success", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}

		req := &apiv1.CreateLightweightRestoreRequest{
			KernelBackupId:   123,
			KernelBackupName: "dsds",
		}

		mockInput := &br.CreateLightweightRestoreInput{
			CreateLightweightRestoreRequest: req,
		}

		mockSvc.On("CreateLightweightRestore", ctx, mockInput).Return(nil)

		resp, err := app.CreateLightweightRestore(ctx, req)

		require.NoError(t, err)
		assert.Equal(t, &emptypb.Empty{}, resp)

		mockSvc.AssertExpectations(t)
	})
}

func TestSetTiProxyExternalDNS(t *testing.T) {
	ctx := context.Background()

	t.Run("validation fails when cluster id is empty", func(t *testing.T) {
		app := &appImpl{}

		resp, err := app.SetTiProxyExternalDNS(ctx, &apiv1.SetTiProxyExternalDNSRequest{})

		assert.Nil(t, resp)
		require.Error(t, err)
		st, ok := status.FromError(err)
		assert.True(t, ok)
		assert.Equal(t, codes.InvalidArgument, st.Code())
	})

	t.Run("validation fails when tidb group id is empty", func(t *testing.T) {
		app := &appImpl{}

		resp, err := app.SetTiProxyExternalDNS(ctx, &apiv1.SetTiProxyExternalDNSRequest{ClusterId: "100"})

		assert.Nil(t, resp)
		require.Error(t, err)
		st, ok := status.FromError(err)
		assert.True(t, ok)
		assert.Equal(t, codes.InvalidArgument, st.Code())
	})

	t.Run("delegates to ngService", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}

		req := &apiv1.SetTiProxyExternalDNSRequest{ClusterId: "100", TidbGroupId: "200", Enable: true}
		mockSvc.On("SetTiProxyExternalDNS", ctx, req).Return(nil).Once()

		resp, err := app.SetTiProxyExternalDNS(ctx, req)

		require.NoError(t, err)
		assert.Equal(t, &emptypb.Empty{}, resp)
		mockSvc.AssertExpectations(t)
	})
}

func TestGetLightweightRestore(t *testing.T) {
	ctx := context.Background()

	t.Run("restore not found", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}

		req := &apiv1.GetLightweightRestoreRequest{RestoreId: "1"}
		downReq := &br.GetLightweightRestoreInput{GetLightweightRestoreRequest: req}

		mockSvc.On("GetLightweightRestore", ctx, downReq).Return(nil, server.ErrLightweightRestoreNotFound)

		resp, err := app.GetLightweightRestore(ctx, req)
		assert.Nil(t, resp)
		require.Error(t, err)
		st, ok := status.FromError(err)
		assert.True(t, ok)
		assert.Equal(t, codes.NotFound, st.Code())
		mockSvc.AssertExpectations(t)
	})

	t.Run("ngService error", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}

		req := &apiv1.GetLightweightRestoreRequest{RestoreId: "2"}
		downReq := &br.GetLightweightRestoreInput{GetLightweightRestoreRequest: req}

		mockSvc.On("GetLightweightRestore", ctx, downReq).Return(nil, errors.New("ng error"))

		resp, err := app.GetLightweightRestore(ctx, req)
		assert.Nil(t, resp)
		require.Error(t, err)
		st, ok := status.FromError(err)
		assert.True(t, ok)
		assert.Equal(t, codes.Internal, st.Code())
		mockSvc.AssertExpectations(t)
	})

	t.Run("success", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}

		req := &apiv1.GetLightweightRestoreRequest{RestoreId: "4"}
		downReq := &br.GetLightweightRestoreInput{GetLightweightRestoreRequest: req}
		downResp := &br.GetLightweightRestoreOutput{
			Status: apiv1.LightweightRestoreStatus_SUCCEED,
		}

		mockSvc.On("GetLightweightRestore", ctx, downReq).Return(downResp, nil)

		resp, err := app.GetLightweightRestore(ctx, req)
		require.NoError(t, err)
		assert.NotNil(t, resp)
		assert.Equal(t, apiv1.LightweightRestoreStatus_SUCCEED, resp.Status)
		assert.Empty(t, resp.Message)
		mockSvc.AssertExpectations(t)
	})
}

func TestDeleteLightweightRestore(t *testing.T) {
	ctx := context.TODO()
	upReq := &apiv1.DeleteLightweightRestoreRequest{RestoreId: "123"}

	t.Run("success", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		mockSvc.On("DeleteLightweightRestore", mock.Anything, &br.DeleteLightweightRestoreInput{
			DeleteLightweightRestoreRequest: upReq,
		}).Return(nil)

		s := &appImpl{
			ngService: mockSvc,
		}

		resp, err := s.DeleteLightweightRestore(ctx, upReq)

		require.NoError(t, err)
		assert.Equal(t, &emptypb.Empty{}, resp)
		mockSvc.AssertExpectations(t)
	})

	t.Run("failure", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		mockSvc.On("DeleteLightweightRestore", mock.Anything, mock.Anything).Return(errors.New("delete failed"))

		s := &appImpl{
			ngService: mockSvc,
		}

		resp, err := s.DeleteLightweightRestore(ctx, upReq)

		assert.Nil(t, resp)
		require.Error(t, err)
		assert.Contains(t, err.Error(), "DeleteLightweightRestore failed")
		st, ok := status.FromError(err)
		assert.True(t, ok)
		assert.Equal(t, codes.Internal, st.Code())
		mockSvc.AssertExpectations(t)
	})
}

func TestForceUpdateGatewayIdentity(t *testing.T) {
	ctx := context.Background()

	t.Run("validation fails on empty cluster id", func(t *testing.T) {
		app := &appImpl{}

		resp, err := app.ForceUpdateGatewayIdentity(ctx, &apiv1.ForceUpdateGatewayIdentityRequest{
			ClusterId:       "",
			GatewayIdentity: "valid-gateway",
		})

		assert.Nil(t, resp)
		require.Error(t, err)
		assert.Equal(t, codes.InvalidArgument, status.Code(err))
		assert.Contains(t, err.Error(), "cluster id is empty")
	})

	t.Run("validation fails on empty gateway identity", func(t *testing.T) {
		app := &appImpl{}

		resp, err := app.ForceUpdateGatewayIdentity(ctx, &apiv1.ForceUpdateGatewayIdentityRequest{
			ClusterId:       "test-cluster-id",
			GatewayIdentity: "",
		})

		assert.Nil(t, resp)
		require.Error(t, err)
		assert.Equal(t, codes.InvalidArgument, status.Code(err))
		assert.Contains(t, err.Error(), "gateway identity is empty")
	})

	t.Run("success", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}

		req := &apiv1.ForceUpdateGatewayIdentityRequest{
			ClusterId:       "test-cluster-id",
			GatewayIdentity: "valid-gateway",
		}

		mockSvc.On("ForceUpdateGatewayIdentity", ctx, "test-cluster-id", "valid-gateway").Return(nil).Once()

		resp, err := app.ForceUpdateGatewayIdentity(ctx, req)

		require.NoError(t, err)
		assert.Equal(t, &emptypb.Empty{}, resp)
		mockSvc.AssertExpectations(t)
	})
}

func TestTiKVCapacityReservation(t *testing.T) {
	ctx := context.Background()

	t.Run("create validates request", func(t *testing.T) {
		app := &appImpl{}

		resp, err := app.CreateTiKVCapacityReservation(ctx, &apiv1.CreateTiKVCapacityReservationRequest{
			SharedPoolId:  "10001",
			ReservationId: "job-1",
		})

		assert.Nil(t, resp)
		require.Error(t, err)
		assert.Equal(t, codes.InvalidArgument, status.Code(err))
	})

	t.Run("create forwards request", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}

		req := &apiv1.CreateTiKVCapacityReservationRequest{
			SharedPoolId:      "10001",
			BizType:           apiv1.TiKVCapacityReservationBizType_TIKV_CAPACITY_RESERVATION_BIZ_TYPE_IMPORT,
			ReservationId:     "job-1",
			RequiredDataBytes: 1024,
		}
		downReq := &server.CreateTiKVCapacityReservationInput{CreateTiKVCapacityReservationRequest: req}
		expected := &apiv1.TiKVCapacityReservation{
			SharedPoolId:  "10001",
			BizType:       apiv1.TiKVCapacityReservationBizType_TIKV_CAPACITY_RESERVATION_BIZ_TYPE_IMPORT,
			ReservationId: "job-1",
		}

		mockSvc.On("CreateTiKVCapacityReservation", ctx, downReq).Return(expected, nil).Once()

		resp, err := app.CreateTiKVCapacityReservation(ctx, req)

		require.NoError(t, err)
		assert.Equal(t, expected, resp)
	})

	t.Run("create auto_scaler denied when flags disabled", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}

		req := &apiv1.CreateTiKVCapacityReservationRequest{
			SharedPoolId:      "10001",
			BizType:           apiv1.TiKVCapacityReservationBizType_TIKV_CAPACITY_RESERVATION_BIZ_TYPE_AUTO_SCALER,
			ReservationId:     "cpu-as-1",
			RequiredDataBytes: 1024,
		}

		// Both flags disabled
		ruConfig := &entity.RuConfig{RuConfig: apiv1.RuConfig{
			FeatureFlags: &apiv1.RuConfig_FeatureFlags{
				EnableAutoscaleTikvCpuScaleOut: proto.Bool(false),
				EnableAutoscaleDiskExpand:      proto.Bool(false),
			},
		}}
		mockSvc.On("GetRuConfig", ctx, "10001", apiv1.RuConfig_SHARED_POOL).Return(ruConfig, nil).Once()

		resp, err := app.CreateTiKVCapacityReservation(ctx, req)

		assert.Nil(t, resp)
		require.Error(t, err)
		assert.Equal(t, codes.PermissionDenied, status.Code(err))
	})

	t.Run("create auto_scaler allowed when cpu flag enabled", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}
		ctx := metadata.NewIncomingContext(context.Background(), metadata.Pairs("x-autoscale-trigger", "tikv_cpu_scale_out"))

		req := &apiv1.CreateTiKVCapacityReservationRequest{
			SharedPoolId:      "10001",
			BizType:           apiv1.TiKVCapacityReservationBizType_TIKV_CAPACITY_RESERVATION_BIZ_TYPE_AUTO_SCALER,
			ReservationId:     "cpu-as-2",
			RequiredDataBytes: 2048,
		}

		ruConfig := &entity.RuConfig{RuConfig: apiv1.RuConfig{
			FeatureFlags: &apiv1.RuConfig_FeatureFlags{
				EnableAutoscaleTikvCpuScaleOut: proto.Bool(true),
				EnableAutoscaleDiskExpand:      proto.Bool(false),
			},
		}}
		mockSvc.On("GetRuConfig", ctx, "10001", apiv1.RuConfig_SHARED_POOL).Return(ruConfig, nil).Once()

		downReq := &server.CreateTiKVCapacityReservationInput{CreateTiKVCapacityReservationRequest: req}
		expected := &apiv1.TiKVCapacityReservation{
			SharedPoolId:  "10001",
			BizType:       apiv1.TiKVCapacityReservationBizType_TIKV_CAPACITY_RESERVATION_BIZ_TYPE_AUTO_SCALER,
			ReservationId: "as-2",
		}
		mockSvc.On("CreateTiKVCapacityReservation", ctx, downReq).Return(expected, nil).Once()

		resp, err := app.CreateTiKVCapacityReservation(ctx, req)

		require.NoError(t, err)
		assert.Equal(t, expected, resp)
	})

	t.Run("create auto_scaler denied when trigger-specific flag disabled", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}
		ctx := metadata.NewIncomingContext(context.Background(), metadata.Pairs("x-autoscale-trigger", "tikv_cpu_scale_out"))

		req := &apiv1.CreateTiKVCapacityReservationRequest{
			SharedPoolId:      "10001",
			BizType:           apiv1.TiKVCapacityReservationBizType_TIKV_CAPACITY_RESERVATION_BIZ_TYPE_AUTO_SCALER,
			ReservationId:     "cpu-as-3",
			RequiredDataBytes: 2048,
		}

		ruConfig := &entity.RuConfig{RuConfig: apiv1.RuConfig{
			FeatureFlags: &apiv1.RuConfig_FeatureFlags{
				EnableAutoscaleTikvCpuScaleOut: proto.Bool(false),
				EnableAutoscaleDiskExpand:      proto.Bool(true),
			},
		}}
		mockSvc.On("GetRuConfig", ctx, "10001", apiv1.RuConfig_SHARED_POOL).Return(ruConfig, nil).Once()

		resp, err := app.CreateTiKVCapacityReservation(ctx, req)

		assert.Nil(t, resp)
		require.Error(t, err)
		assert.Equal(t, codes.PermissionDenied, status.Code(err))
	})

	t.Run("create auto_scaler denied when trigger header missing even if reservation id has prefix", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}

		req := &apiv1.CreateTiKVCapacityReservationRequest{
			SharedPoolId:      "10001",
			BizType:           apiv1.TiKVCapacityReservationBizType_TIKV_CAPACITY_RESERVATION_BIZ_TYPE_AUTO_SCALER,
			ReservationId:     "cpu-as-4",
			RequiredDataBytes: 2048,
		}

		ruConfig := &entity.RuConfig{RuConfig: apiv1.RuConfig{
			FeatureFlags: &apiv1.RuConfig_FeatureFlags{
				EnableAutoscaleTikvCpuScaleOut: proto.Bool(true),
				EnableAutoscaleDiskExpand:      proto.Bool(true),
			},
		}}
		mockSvc.On("GetRuConfig", ctx, "10001", apiv1.RuConfig_SHARED_POOL).Return(ruConfig, nil).Once()

		resp, err := app.CreateTiKVCapacityReservation(ctx, req)

		assert.Nil(t, resp)
		require.Error(t, err)
		assert.Equal(t, codes.PermissionDenied, status.Code(err))
	})

	t.Run("get forwards request", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}

		req := &apiv1.GetTiKVCapacityReservationRequest{
			SharedPoolId:  "10001",
			BizType:       apiv1.TiKVCapacityReservationBizType_TIKV_CAPACITY_RESERVATION_BIZ_TYPE_AUTO_SCALER,
			ReservationId: "as-1",
		}
		downReq := &server.GetTiKVCapacityReservationInput{GetTiKVCapacityReservationRequest: req}
		expected := &apiv1.TiKVCapacityReservation{
			SharedPoolId:  "10001",
			BizType:       apiv1.TiKVCapacityReservationBizType_TIKV_CAPACITY_RESERVATION_BIZ_TYPE_AUTO_SCALER,
			ReservationId: "as-1",
		}

		mockSvc.On("GetTiKVCapacityReservation", ctx, downReq).Return(expected, nil).Once()

		resp, err := app.GetTiKVCapacityReservation(ctx, req)

		require.NoError(t, err)
		assert.Equal(t, expected, resp)
	})

	t.Run("delete forwards request", func(t *testing.T) {
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}

		req := &apiv1.DeleteTiKVCapacityReservationRequest{
			SharedPoolId:  "10001",
			BizType:       apiv1.TiKVCapacityReservationBizType_TIKV_CAPACITY_RESERVATION_BIZ_TYPE_IMPORT,
			ReservationId: "job-1",
		}
		downReq := &server.DeleteTiKVCapacityReservationInput{DeleteTiKVCapacityReservationRequest: req}

		mockSvc.On("DeleteTiKVCapacityReservation", ctx, downReq).Return(nil).Once()

		resp, err := app.DeleteTiKVCapacityReservation(ctx, req)

		require.NoError(t, err)
		assert.Equal(t, &emptypb.Empty{}, resp)
	})
}

func TestScaleSharedPoolAutoscaleSafetyNet(t *testing.T) {
	baseReq := &apiv1.ScaleSharedPoolRequest{
		SharedPoolId: "10001",
		Component: &apiv1.ScaleSharedPoolRequest_ScaleComponent{
			Id:            "comp-1",
			ComponentType: apiv1.ComponentType_TIKV,
			Storages: []*apiv1.Component_Storage{
				{
					Name:   apiv1.Component_Storage_DATA,
					SizeGi: 2500,
				},
			},
		},
	}

	t.Run("autoscale disk expand denied when flag disabled", func(t *testing.T) {
		ctx := metadata.NewIncomingContext(context.Background(), metadata.Pairs("x-autoscale-trigger", "disk_expand"))
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}

		ruConfig := &entity.RuConfig{RuConfig: apiv1.RuConfig{
			FeatureFlags: &apiv1.RuConfig_FeatureFlags{
				EnableAutoscaleDiskExpand: proto.Bool(false),
			},
		}}
		mockSvc.On("GetRuConfig", ctx, "10001", apiv1.RuConfig_SHARED_POOL).Return(ruConfig, nil).Once()

		resp, err := app.ScaleSharedPool(ctx, baseReq)

		assert.Nil(t, resp)
		require.Error(t, err)
		assert.Equal(t, codes.PermissionDenied, status.Code(err))
	})

	t.Run("autoscale disk expand allowed when flag enabled", func(t *testing.T) {
		ctx := metadata.NewIncomingContext(context.Background(), metadata.Pairs("x-autoscale-trigger", "disk_expand"))
		mockSvc := mocks.NewService(t)
		app := &appImpl{ngService: mockSvc}

		ruConfig := &entity.RuConfig{RuConfig: apiv1.RuConfig{
			FeatureFlags: &apiv1.RuConfig_FeatureFlags{
				EnableAutoscaleDiskExpand: proto.Bool(true),
			},
		}}
		mockSvc.On("GetRuConfig", ctx, "10001", apiv1.RuConfig_SHARED_POOL).Return(ruConfig, nil).Once()
		mockSvc.On("ScaleSharedPool", ctx, "10001", mock.MatchedBy(func(comp *entity.Component) bool {
			return comp.GetId() == "comp-1" &&
				comp.GetComponentType() == apiv1.ComponentType_TIKV &&
				len(comp.GetStorages()) == 1 &&
				comp.GetStorages()[0].GetSizeGi() == 2500
		}), false).Return(&entity.SharedPool{SharedPool: apiv1.SharedPool{Id: "10001"}}, nil).Once()

		resp, err := app.ScaleSharedPool(ctx, baseReq)

		require.NoError(t, err)
		require.NotNil(t, resp)
		assert.Equal(t, "10001", resp.GetId())
	})
}

func TestPartitionComponents(t *testing.T) {
	t.Run("empty input", func(t *testing.T) {
		core, brComponents, _ := partitionComponents([]*apiv1.TargetComponent{})

		assert.Empty(t, core)
		assert.Empty(t, brComponents)
	})

	t.Run("only BR components", func(t *testing.T) {
		components := []*apiv1.TargetComponent{
			{Id: "1", ComponentType: apiv1.ComponentType_TIBR_AUTOBACKUP},
			{Id: "2", ComponentType: apiv1.ComponentType_TIBRGC_T2_DFSGC},
			{Id: "3", ComponentType: apiv1.ComponentType_TIBRGC_T3_ARCHIVE},
		}

		core, brComponents, _ := partitionComponents(components)

		assert.Empty(t, core)
		assert.Len(t, brComponents, 3)
		assert.Equal(t, components, brComponents)
	})

	t.Run("only core components", func(t *testing.T) {
		components := []*apiv1.TargetComponent{
			{Id: "1", ComponentType: apiv1.ComponentType_TIKV},
			{Id: "2", ComponentType: apiv1.ComponentType_TIDB},
			{Id: "3", ComponentType: apiv1.ComponentType_PD},
		}

		core, brComponents, _ := partitionComponents(components)

		assert.Len(t, core, 3)
		assert.Empty(t, brComponents)
		assert.Equal(t, components, core)
	})

	t.Run("mixed components", func(t *testing.T) {
		components := []*apiv1.TargetComponent{
			{Id: "1", ComponentType: apiv1.ComponentType_TIKV},
			{Id: "2", ComponentType: apiv1.ComponentType_TIBR_AUTOBACKUP},
			{Id: "3", ComponentType: apiv1.ComponentType_TIDB},
			{Id: "4", ComponentType: apiv1.ComponentType_TIBRGC_T2_DFSGC},
			{Id: "5", ComponentType: apiv1.ComponentType_PD},
			{Id: "6", ComponentType: apiv1.ComponentType_TIBRGC_T3_ARCHIVE},
		}

		core, brComponents, _ := partitionComponents(components)

		assert.Len(t, core, 3)
		assert.Len(t, brComponents, 3)

		// Check core components
		expectedCoreIds := []string{"1", "3", "5"}
		actualCoreIds := make([]string, len(core))
		for i, c := range core {
			actualCoreIds[i] = c.Id
		}
		assert.ElementsMatch(t, expectedCoreIds, actualCoreIds)

		// Check BR components
		expectedBRIds := []string{"2", "4", "6"}
		actualBRIds := make([]string, len(brComponents))
		for i, c := range brComponents {
			actualBRIds[i] = c.Id
		}
		assert.ElementsMatch(t, expectedBRIds, actualBRIds)
	})

	t.Run("single BR component", func(t *testing.T) {
		components := []*apiv1.TargetComponent{
			{Id: "1", ComponentType: apiv1.ComponentType_TIBR_AUTOBACKUP},
		}

		core, brComponents, _ := partitionComponents(components)

		assert.Empty(t, core)
		assert.Len(t, brComponents, 1)
		assert.Equal(t, components[0], brComponents[0])
	})

	t.Run("single core component", func(t *testing.T) {
		components := []*apiv1.TargetComponent{
			{Id: "1", ComponentType: apiv1.ComponentType_TIKV},
		}

		core, brComponents, _ := partitionComponents(components)

		assert.Len(t, core, 1)
		assert.Empty(t, brComponents)
		assert.Equal(t, components[0], core[0])
	})
}