package server

import (
	"context"
	"fmt"
	"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"
	clusteragentv1 "github.com/tidbcloud/cluster-service-ng/api/tidbcloud/cluster-agent/v1"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	bractionmocks "github.com/tidbcloud/cluster-service-ng/domain/workflow/action/br/mocks"
	"github.com/tidbcloud/cluster-service-ng/domain/workflow/dag"
	clustergatewaymocks "github.com/tidbcloud/cluster-service-ng/pkg/clustergateway/mocks"
	dataplanemocks "github.com/tidbcloud/cluster-service-ng/pkg/dataplane/mocks"
	kubemocks "github.com/tidbcloud/cluster-service-ng/pkg/externalmocks/mocks"
	"github.com/tidbcloud/cluster-service-ng/pkg/runtime/executor"
	executormocks "github.com/tidbcloud/cluster-service-ng/pkg/runtime/executor/mocks"
	repomnocks "github.com/tidbcloud/cluster-service-ng/pkg/storage/repo/mocks"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	autoscalingv2 "k8s.io/api/autoscaling/v2"
)

func TestUpdateTiDBGroup_BasicValidation(t *testing.T) {
	testCases := []struct {
		name          string
		tidbGroup     *entity.TiDBGroup
		force         bool
		expectedError string
		expectedCode  codes.Code
	}{
		{
			name: "empty tidb group ID",
			tidbGroup: &entity.TiDBGroup{
				TiDBGroup: apiv1.TiDBGroup{
					Id: "",
				},
			},
			force:         false,
			expectedError: "tidb group id is empty",
			expectedCode:  codes.InvalidArgument,
		},
		{
			name:          "nil tidb group",
			tidbGroup:     nil,
			force:         false,
			expectedError: "tidb group id is empty",
			expectedCode:  codes.InvalidArgument,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			// Create server instance with minimal dependencies for validation tests
			server := &serverImpl{
				transactioner: &mockTransactioner{},
			}

			// Execute test
			result, err := server.UpdateTiDBGroup(context.Background(), tc.tidbGroup, tc.force)

			// Validate results
			require.Error(t, err)
			assert.Contains(t, err.Error(), tc.expectedError)
			statusErr, ok := status.FromError(err)
			require.True(t, ok, "error should be a gRPC status error")
			assert.Equal(t, tc.expectedCode, statusErr.Code())
			assert.Nil(t, result)
		})
	}
}

func TestUpdateTiDBGroup_HPAGateBehavior(t *testing.T) {
	tests := []struct {
		name                string
		request             *entity.TiDBGroup
		force               bool
		setupHPA            func(*kubemocks.Client)
		setupClusterGateway func(*clustergatewaymocks.Client)
		expectedNodeCount   uint32
	}{
		{
			name: "force bypasses HPA scale-out gate",
			request: &entity.TiDBGroup{
				TiDBGroup: apiv1.TiDBGroup{
					Id:        "default",
					ClusterId: "cluster-1",
					NodeQuota: &apiv1.NodeQuota{
						NodeCount: 3,
					},
				},
			},
			force: true,
			setupHPA: func(kubeclient *kubemocks.Client) {
				kubeclient.On("List", mock.Anything, mock.IsType(&autoscalingv2.HorizontalPodAutoscalerList{}), mock.Anything).
					Maybe().
					Run(func(args mock.Arguments) {
						hpaList := args.Get(1).(*autoscalingv2.HorizontalPodAutoscalerList)
						selectPolicy := autoscalingv2.MaxChangePolicySelect
						hpaList.Items = []autoscalingv2.HorizontalPodAutoscaler{
							{
								Spec: autoscalingv2.HorizontalPodAutoscalerSpec{
									Behavior: &autoscalingv2.HorizontalPodAutoscalerBehavior{
										ScaleDown: &autoscalingv2.HPAScalingRules{
											SelectPolicy: &selectPolicy,
										},
									},
								},
							},
						}
					}).
					Return(nil)
			},
			expectedNodeCount: 3,
		},
		{
			name: "scale in is allowed without checking HPA scale-up state",
			request: &entity.TiDBGroup{
				TiDBGroup: apiv1.TiDBGroup{
					Id:        "default",
					ClusterId: "cluster-1",
					NodeQuota: &apiv1.NodeQuota{
						NodeCount: 1,
					},
				},
			},
			force: false,
			setupHPA: func(kubeclient *kubemocks.Client) {
				kubeclient.On("List", mock.Anything, mock.IsType(&autoscalingv2.HorizontalPodAutoscalerList{}), mock.Anything).
					Maybe().
					Run(func(args mock.Arguments) {
						hpaList := args.Get(1).(*autoscalingv2.HorizontalPodAutoscalerList)
						selectPolicy := autoscalingv2.MaxChangePolicySelect
						hpaList.Items = []autoscalingv2.HorizontalPodAutoscaler{
							{
								Spec: autoscalingv2.HorizontalPodAutoscalerSpec{
									Behavior: &autoscalingv2.HorizontalPodAutoscalerBehavior{
										ScaleUp: &autoscalingv2.HPAScalingRules{
											SelectPolicy: &selectPolicy,
										},
									},
								},
							},
						}
					}).
					Return(nil)
			},
			setupClusterGateway: func(clusterGatewayMock *clustergatewaymocks.Client) {
				clusterGatewayMock.On("GetComponentResourceMetrics", mock.Anything, mock.Anything, mock.Anything, mock.Anything).
					Return(&clusteragentv1.GetComponentResourceMetricsResp{
						MetricsServerAvailable: true,
						Components: []*clusteragentv1.ComponentResourceMetrics{
							{CpuCoresUsed: 1.0, MemoryBytesUsed: 1 * 1024 * 1024 * 1024},
							{CpuCoresUsed: 1.0, MemoryBytesUsed: 1 * 1024 * 1024 * 1024},
						},
					}, nil)
			},
			expectedNodeCount: 1,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
			tidbGroupRepo := repomnocks.NewTiDBGroupRepo(t)
			tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t)
			lockRepo := repomnocks.NewLockRepo(t)
			dagExecutor := executormocks.NewRemoteExecutor(t)
			scaleManager := bractionmocks.NewScaleManager(t)
			dataplaneManager := dataplanemocks.NewManager(t)
			clusterGatewayMock := clustergatewaymocks.NewClient(t)
			kubeclient := kubemocks.NewClient(t)

			currentTiDBGroup := &entity.TiDBGroup{
				TiDBGroup: apiv1.TiDBGroup{
					Id:           "default",
					ClusterId:    "cluster-1",
					SharedPoolId: "testsharedpoolid",
					TenantId:     "tenant-1",
					Phase:        apiv1.TiDBGroup_AVAILABLE,
					NodeQuota: &apiv1.NodeQuota{
						NodeCount:    2,
						MinNodeCount: 1,
						MaxNodeCount: 10,
						NodeSpec: &apiv1.NodeSpec{
							Cpu:      4,
							MemoryGi: 8,
						},
					},
				},
			}

			tidbClusterRepo.On("Load", mock.Anything, "cluster-1", false).Return(&entity.TiDBCluster{
				Id: "cluster-1",
			}, nil)
			tidbGroupRepo.On("Load", mock.Anything, "default", false).Return(currentTiDBGroup, nil)
			tidbGroupRepo.On("Save", mock.Anything, mock.AnythingOfType("*entity.TiDBGroup")).Return(nil)
			sharedPoolRepo.On("Load", mock.Anything, "testsharedpoolid", false).Return(createTestSharedPool(), nil)
			lockRepo.On("Lock", mock.Anything, mock.MatchedBy(func(lock *entity.Lock) bool {
				return lock.Key == "testsharedpoolid"
			})).Return(func() error { return nil }, nil)

			dagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).
				Maybe().
				Return(executor.NewExecTaskIdStatus(), nil)
			dagExecutor.On("Submit", mock.Anything, dag.PlanUpdateTiDBGroup, mock.Anything, mock.Anything).
				Return(nil)
			scaleManager.On("AnyFullfillingDecision", mock.Anything, "testsharedpoolid").Maybe().Return(false, nil)

			dataplaneManager.On("GetClient", mock.AnythingOfType("string")).Maybe().Return(kubeclient, nil)
			if tt.setupHPA != nil {
				tt.setupHPA(kubeclient)
			}
			if tt.setupClusterGateway != nil {
				tt.setupClusterGateway(clusterGatewayMock)
			}

			server := &serverImpl{
				sharedPoolRepo:    sharedPoolRepo,
				tidbGroupRepo:     tidbGroupRepo,
				tidbClusterRepo:   tidbClusterRepo,
				lockRepo:          lockRepo,
				transactioner:     &mockTransactioner{},
				dagExecutor:       dagExecutor,
				scaleManager:      scaleManager,
				dataplaneManager:  dataplaneManager,
				clustergatewayCli: clusterGatewayMock,
			}

			result, err := server.UpdateTiDBGroup(context.Background(), tt.request, tt.force)

			require.NoError(t, err)
			require.NotNil(t, result)
			assert.Equal(t, tt.expectedNodeCount, result.GetNodeQuota().GetNodeCount())
			assert.Equal(t, apiv1.TiDBGroup_MODIFYING, result.GetPhase())
			dataplaneManager.AssertNotCalled(t, "GetClient", mock.Anything)
		})
	}
}

func TestUpdateTiDBGroup_EndpointValidation(t *testing.T) {
	tests := []struct {
		name          string
		sqlEndpoint   *apiv1.TiDBGroup_SqlEndpoint
		expectedError string
		description   string
	}{
		{
			name: "valid public lb configuration",
			sqlEndpoint: &apiv1.TiDBGroup_SqlEndpoint{
				PublicLb: &apiv1.TiDBGroup_SqlEndpoint_LoadBalancer{
					Enabled: true,
					Port:    4000,
					IpAllowList: []*entity.IpAllowList{
						{CidrBlock: "192.168.1.0/24"},
						{CidrBlock: "10.0.0.0/8"},
					},
				},
			},
			expectedError: "",
			description:   "Should accept valid public LB configuration",
		},
		{
			name: "missing port when public lb enabled",
			sqlEndpoint: &apiv1.TiDBGroup_SqlEndpoint{
				PublicLb: &apiv1.TiDBGroup_SqlEndpoint_LoadBalancer{
					Enabled: true,
					Port:    0,
					IpAllowList: []*entity.IpAllowList{
						{CidrBlock: "192.168.1.0/24"},
					},
				},
			},
			expectedError: "tidb group public lb port is required when enabled",
			description:   "Should reject missing port when public LB is enabled",
		},
		{
			name: "invalid CIDR block",
			sqlEndpoint: &apiv1.TiDBGroup_SqlEndpoint{
				PublicLb: &apiv1.TiDBGroup_SqlEndpoint_LoadBalancer{
					Enabled: true,
					Port:    4000,
					IpAllowList: []*entity.IpAllowList{
						{CidrBlock: "invalid-cidr"},
					},
				},
			},
			expectedError: "tidb group public lb ip allow list cidr block invalid-cidr is invalid",
			description:   "Should reject invalid CIDR block format",
		},
		{
			name: "empty CIDR block",
			sqlEndpoint: &apiv1.TiDBGroup_SqlEndpoint{
				PublicLb: &apiv1.TiDBGroup_SqlEndpoint_LoadBalancer{
					Enabled: true,
					Port:    4000,
					IpAllowList: []*entity.IpAllowList{
						{CidrBlock: ""},
					},
				},
			},
			expectedError: "tidb group public lb ip allow list cidr block is required when enabled",
			description:   "Should reject empty CIDR block",
		},
		{
			name: "valid private lb configuration",
			sqlEndpoint: &apiv1.TiDBGroup_SqlEndpoint{
				PrivateLb: &apiv1.TiDBGroup_SqlEndpoint_LoadBalancer{
					Enabled: true,
					Port:    4000,
					IpAllowList: []*entity.IpAllowList{
						{CidrBlock: "10.0.0.0/16"},
					},
				},
			},
			expectedError: "",
			description:   "Should accept valid private LB configuration",
		},
		{
			name: "private lb cidr with surrounding whitespace",
			sqlEndpoint: &apiv1.TiDBGroup_SqlEndpoint{
				PrivateLb: &apiv1.TiDBGroup_SqlEndpoint_LoadBalancer{
					Enabled: true,
					Port:    4000,
					IpAllowList: []*entity.IpAllowList{
						{CidrBlock: " 10.0.0.0/16 "},
					},
				},
			},
			expectedError: "",
			description:   "Should normalize surrounding whitespace before validating private LB CIDR block",
		},
		{
			name: "missing port when private lb enabled",
			sqlEndpoint: &apiv1.TiDBGroup_SqlEndpoint{
				PrivateLb: &apiv1.TiDBGroup_SqlEndpoint_LoadBalancer{
					Enabled: true,
					Port:    0,
				},
			},
			expectedError: "tidb group private lb port is required when enabled",
			description:   "Should reject missing port when private LB is enabled",
		},
		{
			name: "invalid private lb CIDR block",
			sqlEndpoint: &apiv1.TiDBGroup_SqlEndpoint{
				PrivateLb: &apiv1.TiDBGroup_SqlEndpoint_LoadBalancer{
					Enabled: true,
					Port:    4000,
					IpAllowList: []*entity.IpAllowList{
						{CidrBlock: "invalid-cidr"},
					},
				},
			},
			expectedError: "tidb group private lb ip allow list cidr block invalid-cidr is invalid",
			description:   "Should reject invalid private LB CIDR block format",
		},
		{
			name: "empty private lb CIDR block",
			sqlEndpoint: &apiv1.TiDBGroup_SqlEndpoint{
				PrivateLb: &apiv1.TiDBGroup_SqlEndpoint_LoadBalancer{
					Enabled: true,
					Port:    4000,
					IpAllowList: []*entity.IpAllowList{
						{CidrBlock: ""},
					},
				},
			},
			expectedError: "tidb group private lb ip allow list cidr block is required when enabled",
			description:   "Should reject empty private LB CIDR block",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tidbGroup := &entity.TiDBGroup{
				TiDBGroup: apiv1.TiDBGroup{
					SqlEndpoint: tt.sqlEndpoint,
				},
			}

			err := checkTidbGroupEndpoint(tidbGroup)

			if tt.expectedError != "" {
				require.Error(t, err, tt.description)
				assert.Contains(t, err.Error(), tt.expectedError, tt.description)
			} else {
				assert.NoError(t, err, tt.description)
			}
		})
	}
}

func TestUpdateTiDBGroupEndpoint_ReplacesPrivateLBWithRequestPayload(t *testing.T) {
	server := &serverImpl{}
	sharedPool := createTestSharedPool()
	currTiDBGroup := &entity.TiDBGroup{
		TiDBGroup: apiv1.TiDBGroup{
			Id:           "default",
			ClusterId:    "cluster-1",
			SharedPoolId: sharedPool.GetId(),
			SqlEndpoint: &entity.SQLEndpoint{
				PrivateLb: &entity.LoadBalancer{
					Enabled: true,
					Port:    4000,
					IpAllowList: []*entity.IpAllowList{
						{CidrBlock: "172.16.0.0/16"},
					},
				},
			},
		},
	}
	tidbGroup := &entity.TiDBGroup{
		TiDBGroup: apiv1.TiDBGroup{
			Id:           "default",
			ClusterId:    "cluster-1",
			SharedPoolId: sharedPool.GetId(),
			SqlEndpoint: &entity.SQLEndpoint{
				PrivateLb: &apiv1.TiDBGroup_SqlEndpoint_LoadBalancer{
					Enabled: false,
					Port:    4000,
				},
			},
		},
	}
	currTiDBComp := &entity.Component{
		Id:            "default",
		ComponentType: apiv1.ComponentType_TIDB,
		Phase:         apiv1.Component_AVAILABLE,
	}
	optionsMap := map[string][]string{}

	lbUpdated, err := server.updateTiDBGroupEndpoint(context.Background(), tidbGroup, currTiDBGroup, sharedPool, currTiDBComp, optionsMap)

	require.NoError(t, err)
	assert.True(t, lbUpdated)
	assert.False(t, currTiDBGroup.GetSqlEndpoint().GetPrivateLb().GetEnabled())
	assert.Empty(t, currTiDBGroup.GetSqlEndpoint().GetPrivateLb().GetIpAllowList())
}

func TestUpdateTiDBGroup_NodeQuotaUpdate(t *testing.T) {
	tests := []struct {
		name          string
		tidbGroup     *entity.TiDBGroup
		currentComp   *entity.Component
		force         bool
		expectedError string
		description   string
	}{
		{
			name: "successful node count update",
			tidbGroup: &entity.TiDBGroup{
				TiDBGroup: apiv1.TiDBGroup{
					Id: "test-group-1",
					NodeQuota: &apiv1.NodeQuota{
						NodeCount: 3,
					},
				},
			},
			currentComp: &entity.Component{
				Id:            "test-group-1",
				ComponentType: apiv1.ComponentType_TIDB,
				Phase:         apiv1.Component_AVAILABLE,
				NodeQuota: &apiv1.NodeQuota{
					NodeCount:    1,
					MinNodeCount: 1,
					MaxNodeCount: 5,
					NodeSpec: &apiv1.NodeSpec{
						Cpu:      2,
						MemoryGi: 4,
					},
				},
			},
			force:         false,
			expectedError: "",
			description:   "Should successfully update node count within allowed range",
		},
		{
			name: "node count exceeds maximum",
			tidbGroup: &entity.TiDBGroup{
				TiDBGroup: apiv1.TiDBGroup{
					Id: "test-group-1",
					NodeQuota: &apiv1.NodeQuota{
						NodeCount: 10,
					},
				},
			},
			currentComp: &entity.Component{
				Id:            "test-group-1",
				ComponentType: apiv1.ComponentType_TIDB,
				Phase:         apiv1.Component_AVAILABLE,
				NodeQuota: &apiv1.NodeQuota{
					NodeCount:    1,
					MinNodeCount: 1,
					MaxNodeCount: 5,
				},
			},
			force:         false,
			expectedError: "tidb group node count is larger than max node count",
			description:   "Should reject node count that exceeds maximum",
		},
		{
			name: "node count below minimum",
			tidbGroup: &entity.TiDBGroup{
				TiDBGroup: apiv1.TiDBGroup{
					Id: "test-group-1",
					NodeQuota: &apiv1.NodeQuota{
						NodeCount: 0,
					},
				},
			},
			currentComp: &entity.Component{
				Id:            "test-group-1",
				ComponentType: apiv1.ComponentType_TIDB,
				Phase:         apiv1.Component_AVAILABLE,
				NodeQuota: &apiv1.NodeQuota{
					NodeCount:    2,
					MinNodeCount: 1,
					MaxNodeCount: 5,
				},
			},
			force:         false,
			expectedError: "tidb group node count is less than min node count",
			description:   "Should reject node count that is below minimum",
		},
		{
			name: "update when component not available without force",
			tidbGroup: &entity.TiDBGroup{
				TiDBGroup: apiv1.TiDBGroup{
					Id: "test-group-1",
					NodeQuota: &apiv1.NodeQuota{
						NodeCount: 3,
					},
				},
			},
			currentComp: &entity.Component{
				Id:            "test-group-1",
				ComponentType: apiv1.ComponentType_TIDB,
				Phase:         apiv1.Component_CREATING,
				NodeQuota: &apiv1.NodeQuota{
					NodeCount:    1,
					MinNodeCount: 1,
					MaxNodeCount: 5,
				},
			},
			force:         false,
			expectedError: "tidb group phase is CREATING, can not update node count",
			description:   "Should reject update when component is not available and force is false",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// This is a unit test for the business logic validation
			// We're not testing the full UpdateTiDBGroup method here
			// but rather the specific validation logic for node quota updates

			var validationError error

			// Test the node count validation logic
			if tt.tidbGroup.GetNodeQuota() != nil {
				// Check if there's a node count update request
				if tt.tidbGroup.GetNodeQuota().GetNodeCount() != tt.currentComp.GetNodeQuota().GetNodeCount() {
					switch {
					case tt.currentComp.GetNodeQuota().GetMaxNodeCount() < tt.tidbGroup.GetNodeQuota().GetNodeCount():
						validationError = status.Errorf(codes.InvalidArgument, "tidb group node count is larger than max node count")
					case tt.currentComp.GetNodeQuota().GetMinNodeCount() > tt.tidbGroup.GetNodeQuota().GetNodeCount():
						validationError = status.Errorf(codes.InvalidArgument, "tidb group node count is less than min node count")
					case tt.currentComp.GetPhase() != apiv1.Component_AVAILABLE && !tt.force:
						validationError = status.Errorf(codes.FailedPrecondition, "tidb group phase is %s, can not update node count", tt.currentComp.GetPhase().String())
					}
				}
			}

			// Verify results
			if tt.expectedError != "" {
				require.Error(t, validationError, tt.description)
				assert.Contains(t, validationError.Error(), tt.expectedError, tt.description)
			} else {
				assert.NoError(t, validationError, tt.description)
			}
		})
	}
}

func TestUpdateTiDBGroup_TiProxyComponentUpdate(t *testing.T) {
	tests := []struct {
		name               string
		sharedPoolComps    []*entity.Component
		tidbGroupClusterId string
		updateLb           bool
		expectTiproxyInIds bool
		expectError        bool
		description        string
	}{
		{
			name: "tiproxy component found when updating lb",
			sharedPoolComps: []*entity.Component{
				{
					Id:            "cluster-1",
					ComponentType: apiv1.ComponentType_TIPROXY,
					Phase:         apiv1.Component_AVAILABLE,
				},
				{
					Id:            "tidb-group-1",
					ComponentType: apiv1.ComponentType_TIDB,
					Phase:         apiv1.Component_AVAILABLE,
				},
			},
			tidbGroupClusterId: "cluster-1",
			updateLb:           true,
			expectTiproxyInIds: true,
			expectError:        false,
			description:        "Should find tiproxy component and add to component ids when updating lb",
		},
		{
			name: "tiproxy component not found when updating lb",
			sharedPoolComps: []*entity.Component{
				{
					Id:            "tidb-group-1",
					ComponentType: apiv1.ComponentType_TIDB,
					Phase:         apiv1.Component_AVAILABLE,
				},
			},
			tidbGroupClusterId: "cluster-1",
			updateLb:           true,
			expectTiproxyInIds: false,
			expectError:        true,
			description:        "Should return error when tiproxy component not found",
		},
		{
			name: "no tiproxy update when lb not changed",
			sharedPoolComps: []*entity.Component{
				{
					Id:            "cluster-1",
					ComponentType: apiv1.ComponentType_TIPROXY,
					Phase:         apiv1.Component_AVAILABLE,
				},
			},
			tidbGroupClusterId: "cluster-1",
			updateLb:           false,
			expectTiproxyInIds: false,
			expectError:        false,
			description:        "Should not add tiproxy to component ids when lb is not updated",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Simulate finding tiproxy component
			var tiproxyComp *entity.Component
			var err error

			if tt.updateLb {
				for _, comp := range tt.sharedPoolComps {
					if comp.GetComponentType() == apiv1.ComponentType_TIPROXY && comp.GetId() == tt.tidbGroupClusterId {
						tiproxyComp = comp
						break
					}
				}
				if tiproxyComp == nil {
					err = status.Errorf(codes.Internal, "tiproxy component not found in shared pool")
				}
			}

			// Build component ids
			componentIds := []string{"tidb-group-1"}
			if tt.updateLb && tiproxyComp != nil {
				componentIds = append(componentIds, tt.tidbGroupClusterId)
			}

			// Verify results
			if tt.expectError {
				require.Error(t, err, tt.description)
				assert.Contains(t, err.Error(), "tiproxy component not found", tt.description)
			} else {
				assert.NoError(t, err, tt.description)
				if tt.expectTiproxyInIds {
					assert.Contains(t, componentIds, tt.tidbGroupClusterId, tt.description)
					assert.Len(t, componentIds, 2, tt.description)
				} else {
					assert.Len(t, componentIds, 1, tt.description)
				}
			}
		})
	}
}

// TestValidateNodeSpecAgainstRUConfig_TiDBComponent tests RU validation for TiDB component updates
func TestValidateNodeSpecAgainstRUConfig_TiDBComponent(t *testing.T) {
	tests := []struct {
		name            string
		requestNodeSpec *entity.NodeSpec
		ruConfig        *entity.RuConfig
		minRU           uint32
		force           bool
		expectedError   string
		expectedCode    codes.Code
		description     string
	}{
		{
			name: "valid node spec meets RU requirement",
			requestNodeSpec: &entity.NodeSpec{
				Cpu:      8,
				MemoryGi: 16,
			},
			ruConfig: &entity.RuConfig{
				RuConfig: apiv1.RuConfig{
					Components: map[string]*apiv1.RuConfig_ComponentsMap{
						"tidb": {NodeSpecRuRange: map[string]*apiv1.RuConfig_RuRange{
							"4c8g":  {Min: 0, Max: 5000},
							"8c16g": {Min: 5001, Max: 20000},
						}},
					},
				},
			},
			minRU:         10000,
			force:         false,
			expectedError: "",
			expectedCode:  codes.OK,
			description:   "8c16g should support 10000 RU",
		},
		{
			name: "node spec below RU requirement - reject",
			requestNodeSpec: &entity.NodeSpec{
				Cpu:      4,
				MemoryGi: 8,
			},
			ruConfig: &entity.RuConfig{
				RuConfig: apiv1.RuConfig{
					Components: map[string]*apiv1.RuConfig_ComponentsMap{
						"tidb": {NodeSpecRuRange: map[string]*apiv1.RuConfig_RuRange{
							"4c8g":  {Min: 0, Max: 5000},
							"8c16g": {Min: 5001, Max: 20000},
						}},
					},
				},
			},
			minRU:         10000,
			force:         false,
			expectedError: "cannot support min RU",
			expectedCode:  codes.InvalidArgument,
			description:   "4c8g should not support 10000 RU, should reject",
		},
		{
			name: "node spec below RU requirement - force bypass",
			requestNodeSpec: &entity.NodeSpec{
				Cpu:      4,
				MemoryGi: 8,
			},
			ruConfig: &entity.RuConfig{
				RuConfig: apiv1.RuConfig{
					Components: map[string]*apiv1.RuConfig_ComponentsMap{
						"tidb": {NodeSpecRuRange: map[string]*apiv1.RuConfig_RuRange{
							"4c8g":  {Min: 0, Max: 5000},
							"8c16g": {Min: 5001, Max: 20000},
						}},
					},
				},
			},
			minRU:         10000,
			force:         true,
			expectedError: "",
			expectedCode:  codes.OK,
			description:   "4c8g with force=true should bypass RU validation",
		},
		{
			name: "minRU is zero - skip validation",
			requestNodeSpec: &entity.NodeSpec{
				Cpu:      2,
				MemoryGi: 4,
			},
			ruConfig: &entity.RuConfig{
				RuConfig: apiv1.RuConfig{
					Components: map[string]*apiv1.RuConfig_ComponentsMap{
						"tidb": {NodeSpecRuRange: map[string]*apiv1.RuConfig_RuRange{
							"2c4g": {Min: 0, Max: 1000},
						}},
					},
				},
			},
			minRU:         0,
			force:         false,
			expectedError: "",
			expectedCode:  codes.OK,
			description:   "minRU=0 should skip validation entirely",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			err := validateNodeSpecAgainstRUConfig(
				context.Background(),
				apiv1.ComponentType_TIDB,
				tt.requestNodeSpec,
				tt.ruConfig,
				tt.minRU,
				"test-tidb-component",
				tt.force,
			)

			if tt.expectedError != "" {
				require.Error(t, err, tt.description)
				assert.Contains(t, err.Error(), tt.expectedError, tt.description)
				statusErr, ok := status.FromError(err)
				require.True(t, ok, "error should be a gRPC status error")
				assert.Equal(t, tt.expectedCode, statusErr.Code(), tt.description)
			} else {
				assert.NoError(t, err, tt.description)
			}
		})
	}
}

// TestUpdateTiDBGroup_ScaleProtection tests the scale protection logic for UpdateTiDBGroup
func TestUpdateTiDBGroup_ScaleProtection(t *testing.T) {
	tests := []struct {
		name            string
		force           bool
		currentCount    uint32
		currentSpec     *entity.NodeSpec
		targetCount     uint32
		targetSpec      *entity.NodeSpec
		metricsResponse *clusteragentv1.GetComponentResourceMetricsResp
		metricsError    error
		expectedError   bool
		expectedCode    codes.Code
		description     string
	}{
		{
			name:          "scale protection skipped with force flag",
			force:         true,
			currentCount:  3,
			currentSpec:   &entity.NodeSpec{Cpu: 8, MemoryGi: 16},
			targetCount:   2,
			targetSpec:    &entity.NodeSpec{Cpu: 8, MemoryGi: 16},
			expectedError: false,
			description:   "Force flag should bypass scale protection",
		},
		{
			name:          "scale up allowed without force",
			force:         false,
			currentCount:  2,
			currentSpec:   &entity.NodeSpec{Cpu: 8, MemoryGi: 16},
			targetCount:   3,
			targetSpec:    &entity.NodeSpec{Cpu: 8, MemoryGi: 16},
			expectedError: false,
			description:   "Scale up should be allowed without force flag",
		},
		{
			name:          "spec upgrade allowed without force",
			force:         false,
			currentCount:  2,
			currentSpec:   &entity.NodeSpec{Cpu: 8, MemoryGi: 16},
			targetCount:   2,
			targetSpec:    &entity.NodeSpec{Cpu: 16, MemoryGi: 32},
			expectedError: false,
			description:   "Spec upgrade should be allowed without force flag",
		},
		{
			name:         "safe scale in - low usage",
			force:        false,
			currentCount: 3,
			currentSpec:  &entity.NodeSpec{Cpu: 8, MemoryGi: 16},
			targetCount:  2,
			targetSpec:   &entity.NodeSpec{Cpu: 8, MemoryGi: 16},
			metricsResponse: &clusteragentv1.GetComponentResourceMetricsResp{
				MetricsServerAvailable: true,
				Components: []*clusteragentv1.ComponentResourceMetrics{
					{CpuCoresUsed: 1.0, MemoryBytesUsed: 2 * 1024 * 1024 * 1024},
					{CpuCoresUsed: 1.0, MemoryBytesUsed: 2 * 1024 * 1024 * 1024},
					{CpuCoresUsed: 1.0, MemoryBytesUsed: 2 * 1024 * 1024 * 1024},
				},
			},
			expectedError: false,
			description:   "Scale in should be allowed when load is low",
		},
		{
			name:         "unsafe scale in - high cpu usage",
			force:        false,
			currentCount: 3,
			currentSpec:  &entity.NodeSpec{Cpu: 8, MemoryGi: 16},
			targetCount:  2,
			targetSpec:   &entity.NodeSpec{Cpu: 8, MemoryGi: 16},
			metricsResponse: &clusteragentv1.GetComponentResourceMetricsResp{
				MetricsServerAvailable: true,
				Components: []*clusteragentv1.ComponentResourceMetrics{
					{CpuCoresUsed: 6.0, MemoryBytesUsed: 2 * 1024 * 1024 * 1024},
					{CpuCoresUsed: 6.0, MemoryBytesUsed: 2 * 1024 * 1024 * 1024},
					{CpuCoresUsed: 6.0, MemoryBytesUsed: 2 * 1024 * 1024 * 1024},
				},
			},
			expectedError: true,
			expectedCode:  codes.FailedPrecondition,
			description:   "Scale in should be blocked when CPU load is high",
		},
		{
			name:         "unsafe scale in - high memory usage",
			force:        false,
			currentCount: 3,
			currentSpec:  &entity.NodeSpec{Cpu: 8, MemoryGi: 16},
			targetCount:  2,
			targetSpec:   &entity.NodeSpec{Cpu: 8, MemoryGi: 16},
			metricsResponse: &clusteragentv1.GetComponentResourceMetricsResp{
				MetricsServerAvailable: true,
				Components: []*clusteragentv1.ComponentResourceMetrics{
					{CpuCoresUsed: 1.0, MemoryBytesUsed: 12 * 1024 * 1024 * 1024},
					{CpuCoresUsed: 1.0, MemoryBytesUsed: 12 * 1024 * 1024 * 1024},
					{CpuCoresUsed: 1.0, MemoryBytesUsed: 12 * 1024 * 1024 * 1024},
				},
			},
			expectedError: true,
			expectedCode:  codes.FailedPrecondition,
			description:   "Scale in should be blocked when memory load is high",
		},
		{
			name:         "metrics server unavailable",
			force:        false,
			currentCount: 3,
			currentSpec:  &entity.NodeSpec{Cpu: 8, MemoryGi: 16},
			targetCount:  2,
			targetSpec:   &entity.NodeSpec{Cpu: 8, MemoryGi: 16},
			metricsResponse: &clusteragentv1.GetComponentResourceMetricsResp{
				MetricsServerAvailable: false,
			},
			expectedError: true,
			expectedCode:  codes.FailedPrecondition,
			description:   "Scale in should be blocked when metrics server is unavailable",
		},
		{
			name:          "metrics fetch error",
			force:         false,
			currentCount:  3,
			currentSpec:   &entity.NodeSpec{Cpu: 8, MemoryGi: 16},
			targetCount:   2,
			targetSpec:    &entity.NodeSpec{Cpu: 8, MemoryGi: 16},
			metricsError:  fmt.Errorf("connection refused"),
			expectedError: true,
			description:   "Scale in should be blocked when metrics fetch fails",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Setup mocks
			clusterGatewayMock := clustergatewaymocks.NewClient(t)
			if tt.metricsResponse != nil || tt.metricsError != nil {
				clusterGatewayMock.On("GetComponentResourceMetrics", mock.Anything, mock.Anything, mock.Anything, mock.Anything).
					Return(tt.metricsResponse, tt.metricsError)
			}

			server := &serverImpl{
				clustergatewayCli: clusterGatewayMock,
			}

			sharedPool := &entity.SharedPool{
				SharedPool: apiv1.SharedPool{
					Id:            "test-pool",
					DataplaneName: "test-dp",
				},
			}
			component := &entity.Component{
				Id:            "test-tidb-group",
				ComponentType: apiv1.ComponentType_TIDB,
			}

			err := server.checkScaleInSafety(
				context.Background(),
				sharedPool,
				component,
				tt.currentCount,
				tt.currentSpec,
				tt.targetCount,
				tt.targetSpec,
				tt.force,
			)

			if tt.expectedError {
				require.Error(t, err, tt.description)
				if tt.expectedCode != 0 {
					statusErr, ok := status.FromError(err)
					require.True(t, ok, "error should be a gRPC status error")
					assert.Equal(t, tt.expectedCode, statusErr.Code(), tt.description)
				}
			} else {
				require.NoError(t, err, tt.description)
			}

			clusterGatewayMock.AssertExpectations(t)
		})
	}
}

// TestValidateUpdateTiDBGroupRequest tests the request validation function
func TestValidateUpdateTiDBGroupRequest(t *testing.T) {
	tests := []struct {
		name          string
		tidbGroup     *entity.TiDBGroup
		expectedError bool
		expectedCode  codes.Code
	}{
		{
			name: "valid request",
			tidbGroup: &entity.TiDBGroup{
				TiDBGroup: apiv1.TiDBGroup{
					Id: "test-group",
				},
			},
			expectedError: false,
		},
		{
			name:          "nil tidb group",
			tidbGroup:     nil,
			expectedError: true,
			expectedCode:  codes.InvalidArgument,
		},
		{
			name: "empty id",
			tidbGroup: &entity.TiDBGroup{
				TiDBGroup: apiv1.TiDBGroup{
					Id: "",
				},
			},
			expectedError: true,
			expectedCode:  codes.InvalidArgument,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			server := &serverImpl{}
			err := server.validateUpdateTiDBGroupRequest(tt.tidbGroup)

			if tt.expectedError {
				require.Error(t, err)
				statusErr, ok := status.FromError(err)
				require.True(t, ok)
				assert.Equal(t, tt.expectedCode, statusErr.Code())
			} else {
				assert.NoError(t, err)
			}
		})
	}
}

// TestFindTiDBComponent tests the findTiDBComponent helper function
func TestFindTiDBComponent(t *testing.T) {
	tests := []struct {
		name          string
		components    []*entity.Component
		tidbGroupId   string
		expectedFound bool
	}{
		{
			name: "component found",
			components: []*entity.Component{
				{
					Id:            "group-1",
					ComponentType: apiv1.ComponentType_TIDB,
				},
				{
					Id:            "group-2",
					ComponentType: apiv1.ComponentType_TIDB,
				},
			},
			tidbGroupId:   "group-1",
			expectedFound: true,
		},
		{
			name: "component not found - wrong type",
			components: []*entity.Component{
				{
					Id:            "group-1",
					ComponentType: apiv1.ComponentType_TIKV,
				},
			},
			tidbGroupId:   "group-1",
			expectedFound: false,
		},
		{
			name: "component not found - wrong id",
			components: []*entity.Component{
				{
					Id:            "group-1",
					ComponentType: apiv1.ComponentType_TIDB,
				},
			},
			tidbGroupId:   "group-2",
			expectedFound: false,
		},
		{
			name:          "empty component list",
			components:    []*entity.Component{},
			tidbGroupId:   "group-1",
			expectedFound: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			server := &serverImpl{}
			sharedPool := &entity.SharedPool{
				SharedPool: apiv1.SharedPool{
					ComponentGroups: tt.components,
				},
			}

			comp, err := server.findTiDBComponent(sharedPool, tt.tidbGroupId)

			if tt.expectedFound {
				require.NoError(t, err)
				assert.NotNil(t, comp)
				assert.Equal(t, tt.tidbGroupId, comp.GetId())
				assert.Equal(t, apiv1.ComponentType_TIDB, comp.GetComponentType())
			} else {
				require.Error(t, err)
				assert.Equal(t, codes.NotFound, status.Code(err))
			}
		})
	}
}

// TestFindTiProxyComponent tests the findTiProxyComponent helper function
func TestFindTiProxyComponent(t *testing.T) {
	tests := []struct {
		name        string
		components  []*entity.Component
		clusterId   string
		expectedNil bool
	}{
		{
			name: "tiproxy found",
			components: []*entity.Component{
				{
					Id:            "cluster-1",
					ComponentType: apiv1.ComponentType_TIPROXY,
				},
			},
			clusterId:   "cluster-1",
			expectedNil: false,
		},
		{
			name: "tiproxy not found - wrong cluster",
			components: []*entity.Component{
				{
					Id:            "cluster-1",
					ComponentType: apiv1.ComponentType_TIPROXY,
				},
			},
			clusterId:   "cluster-2",
			expectedNil: true,
		},
		{
			name: "tiproxy not found - wrong type",
			components: []*entity.Component{
				{
					Id:            "cluster-1",
					ComponentType: apiv1.ComponentType_TIDB,
				},
			},
			clusterId:   "cluster-1",
			expectedNil: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			server := &serverImpl{}
			sharedPool := &entity.SharedPool{
				SharedPool: apiv1.SharedPool{
					ComponentGroups: tt.components,
				},
			}

			comp := server.findTiProxyComponent(sharedPool, tt.clusterId)

			if tt.expectedNil {
				assert.Nil(t, comp)
			} else {
				assert.NotNil(t, comp)
				assert.Equal(t, tt.clusterId, comp.GetId())
				assert.Equal(t, apiv1.ComponentType_TIPROXY, comp.GetComponentType())
			}
		})
	}
}

// TestValidateNodeCountUpdate tests the node count validation logic
func TestValidateNodeCountUpdate(t *testing.T) {
	tests := []struct {
		name          string
		phase         apiv1.TiDBGroup_Phase
		force         bool
		currentCount  uint32
		targetCount   uint32
		minCount      uint32
		maxCount      uint32
		expectedError bool
		expectedCode  codes.Code
	}{
		{
			name:          "valid update",
			phase:         apiv1.TiDBGroup_AVAILABLE,
			force:         false,
			currentCount:  2,
			targetCount:   3,
			minCount:      1,
			maxCount:      5,
			expectedError: false,
		},
		{
			name:          "not available without force",
			phase:         apiv1.TiDBGroup_CREATING,
			force:         false,
			currentCount:  2,
			targetCount:   3,
			minCount:      1,
			maxCount:      5,
			expectedError: true,
			expectedCode:  codes.FailedPrecondition,
		},
		{
			name:          "not available with force",
			phase:         apiv1.TiDBGroup_CREATING,
			force:         true,
			currentCount:  2,
			targetCount:   3,
			minCount:      1,
			maxCount:      5,
			expectedError: false,
		},
		{
			name:          "exceeds max",
			phase:         apiv1.TiDBGroup_AVAILABLE,
			force:         false,
			currentCount:  2,
			targetCount:   10,
			minCount:      1,
			maxCount:      5,
			expectedError: true,
			expectedCode:  codes.InvalidArgument,
		},
		{
			name:          "below min",
			phase:         apiv1.TiDBGroup_AVAILABLE,
			force:         false,
			currentCount:  3,
			targetCount:   0,
			minCount:      1,
			maxCount:      5,
			expectedError: true,
			expectedCode:  codes.InvalidArgument,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			server := &serverImpl{}
			tidbGroup := &entity.TiDBGroup{
				TiDBGroup: apiv1.TiDBGroup{
					Phase: tt.phase,
					NodeQuota: &apiv1.NodeQuota{
						NodeCount: tt.targetCount,
					},
				},
			}
			currTiDBGroup := &entity.TiDBGroup{
				TiDBGroup: apiv1.TiDBGroup{
					Phase: tt.phase,
				},
			}
			currTiDBComp := &entity.Component{
				NodeQuota: &apiv1.NodeQuota{
					NodeCount:    tt.currentCount,
					MinNodeCount: tt.minCount,
					MaxNodeCount: tt.maxCount,
				},
			}

			err := server.validateNodeCountUpdate(context.Background(), tidbGroup, currTiDBGroup, currTiDBComp, tt.force)

			if tt.expectedError {
				require.Error(t, err)
				assert.Equal(t, tt.expectedCode, status.Code(err))
			} else {
				assert.NoError(t, err)
			}
		})
	}
}