package server

import (
	"context"
	"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/workflow"
	bractionmocks "github.com/tidbcloud/cluster-service-ng/domain/workflow/action/br/mocks"
	"github.com/tidbcloud/cluster-service-ng/domain/workflow/dag"
	"github.com/tidbcloud/cluster-service-ng/pkg/runtime/actions"
	"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"
	"google.golang.org/protobuf/encoding/protojson"
)

func TestGetSharedPool(t *testing.T) {
	tests := []struct {
		name          string
		sharedPoolId  string
		setupMocks    func(*repomnocks.SharedPoolRepo, *executormocks.RemoteExecutor, *bractionmocks.ScaleManager)
		expectedError string
		expectedCode  codes.Code
		validate      func(*testing.T, *entity.SharedPool, error)
	}{
		{
			name:         "successful get existing shared pool",
			sharedPoolId: "testsharedpool-id",
			setupMocks: func(sharedPoolRepo *repomnocks.SharedPoolRepo, dagExecutor *executormocks.RemoteExecutor, scalemgr *bractionmocks.ScaleManager) {
				sharedPool := &entity.SharedPool{
					SharedPool: apiv1.SharedPool{
						Id:    "testsharedpool-id",
						Phase: apiv1.SharedPool_AVAILABLE,
						ComponentGroups: []*apiv1.Component{
							{
								Id:            "comp-1",
								ComponentType: apiv1.ComponentType_TIKV,
								Phase:         apiv1.Component_AVAILABLE,
							},
						},
					},
				}
				sharedPoolRepo.On("Load", mock.Anything, "testsharedpool-id", false).Return(sharedPool, nil)

				// Mock merge status from DAG task (no running DAGs)
				dagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil)
				// Mock no restore
				scalemgr.On("AnyFullfillingDecision", mock.Anything, sharedPool.Id).Return(false, nil)
			},
			validate: func(t *testing.T, result *entity.SharedPool, err error) {
				t.Helper()
				require.NoError(t, err)
				require.NotNil(t, result)
				assert.Equal(t, "testsharedpool-id", result.GetId())
				assert.Equal(t, apiv1.SharedPool_AVAILABLE, result.GetPhase())
				assert.Len(t, result.GetComponentGroups(), 1)
			},
		},
		{
			name:         "successful get shared pool with running DAG",
			sharedPoolId: "testsharedpool-id",
			setupMocks: func(sharedPoolRepo *repomnocks.SharedPoolRepo, dagExecutor *executormocks.RemoteExecutor, scalemgr *bractionmocks.ScaleManager) {
				sharedPool := &entity.SharedPool{
					SharedPool: apiv1.SharedPool{
						Id:    "testsharedpool-id",
						Phase: apiv1.SharedPool_AVAILABLE,
						ComponentGroups: []*apiv1.Component{
							{
								Id:            "comp-1",
								ComponentType: apiv1.ComponentType_TIKV,
								Phase:         apiv1.Component_AVAILABLE,
							},
						},
					},
				}
				sharedPoolRepo.On("Load", mock.Anything, "testsharedpool-id", false).Return(sharedPool, nil)

				// Mock running DAG that updates the shared pool status
				tempSharedPool := &entity.SharedPool{
					SharedPool: apiv1.SharedPool{
						Id:    "testsharedpool-id",
						Phase: apiv1.SharedPool_MODIFYING,
						ComponentGroups: []*apiv1.Component{
							{
								Id:            "comp-1",
								ComponentType: apiv1.ComponentType_TIKV,
								Phase:         apiv1.Component_MODIFYING,
							},
						},
					},
				}
				sharedPoolStr, _ := protojson.Marshal(tempSharedPool)

				runningDags := executor.NewExecTaskIdStatus()
				runningDags["dag-1"] = &executor.ExecTaskStatusEx{
					ExecTaskStatus: executor.ExecTaskStatus{
						ExecTask: &executor.ExecutorTask{
							RootName: dag.PlanUpdateTiDBGroup,
							Args: actions.ActionArgs{
								workflow.SharedPoolKey:     actions.ActionArgString(string(sharedPoolStr)),
								workflow.ComponentIdsKey:   actions.ActionArgString("comp-1"),
								workflow.ComponentTypesKey: actions.ActionArgString("tikv"),
							},
						},
					},
				}

				// Mock all DAG plans that are checked in mergeSharedPoolStatusFromDagTask
				dagExecutor.On("QueryRunnings", mock.Anything, "create-sharedpool", mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil)
				dagExecutor.On("QueryRunnings", mock.Anything, "delete-sharedpool", mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil)
				dagExecutor.On("QueryRunnings", mock.Anything, "update-component", mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil)
				dagExecutor.On("QueryRunnings", mock.Anything, "apply-component", mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil)
				dagExecutor.On("QueryRunnings", mock.Anything, "create-component", mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil)
				dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanDeleteTiKVGroup, mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil)
				dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanSetSharedPoolPDDeploymentModeToMonolithic, mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil)
				dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanSetSharedPoolPDDeploymentModeToMicroservice, mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil)
				dagExecutor.On("QueryRunnings", mock.Anything, "create-tidbcluster", mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil)
				dagExecutor.On("QueryRunnings", mock.Anything, "update-tidbcluster", mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil)
				dagExecutor.On("QueryRunnings", mock.Anything, "delete-tidbcluster", mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil)
				dagExecutor.On("QueryRunnings", mock.Anything, "create-tidbgroup", mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil)
				// This call returns the running DAG for update-tidbgroup
				dagExecutor.On("QueryRunnings", mock.Anything, "update-tidbgroup", mock.AnythingOfType("executor.MatchFilter")).Return(runningDags, nil)
				dagExecutor.On("QueryRunnings", mock.Anything, "delete-tidbgroup", mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil)
				dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanEnableCoprocessorWorker, mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil)
				dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanDisableCoprocessorWorker, mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil)
				// Mock no restore
				scalemgr.On("AnyFullfillingDecision", mock.Anything, sharedPool.Id).Return(false, nil)
			},
			validate: func(t *testing.T, result *entity.SharedPool, err error) {
				t.Helper()
				require.NoError(t, err)
				require.NotNil(t, result)
				assert.Equal(t, "testsharedpool-id", result.GetId())
				assert.Equal(t, apiv1.SharedPool_MODIFYING, result.GetPhase())
				assert.Len(t, result.GetComponentGroups(), 1)
				assert.Equal(t, apiv1.Component_MODIFYING, result.GetComponentGroups()[0].GetPhase())
			},
		},
		{
			name:         "shared pool not found",
			sharedPoolId: "nonexistingpool",
			setupMocks: func(sharedPoolRepo *repomnocks.SharedPoolRepo, dagExecutor *executormocks.RemoteExecutor, scalemgr *bractionmocks.ScaleManager) {
				// Mock shared pool not found in repo
				sharedPoolRepo.On("Load", mock.Anything, "nonexistingpool", false).Return(nil, status.Error(codes.NotFound, "not found"))
			},
			expectedError: "sharedPool nonexistingpool not found",
			expectedCode:  codes.NotFound,
		},
		{
			name:         "empty shared pool id",
			sharedPoolId: "",
			setupMocks: func(sharedPoolRepo *repomnocks.SharedPoolRepo, dagExecutor *executormocks.RemoteExecutor, scalemgr *bractionmocks.ScaleManager) {
				// No mocks needed for this test case
			},
			expectedError: "sharedPoolId is empty",
			expectedCode:  codes.InvalidArgument,
		},
		{
			name:         "repo load error",
			sharedPoolId: "error-pool",
			setupMocks: func(sharedPoolRepo *repomnocks.SharedPoolRepo, dagExecutor *executormocks.RemoteExecutor, scalemgr *bractionmocks.ScaleManager) {
				// Mock repo error (not NotFound)
				sharedPoolRepo.On("Load", mock.Anything, "error-pool", false).Return(nil, status.Error(codes.Internal, "database error"))
			},
			expectedError: "failed to load sharedPool: rpc error: code = Internal desc = database error",
			expectedCode:  codes.Internal,
		},
		{
			name:         "DAG executor error during merge",
			sharedPoolId: "dag-error-pool",
			setupMocks: func(sharedPoolRepo *repomnocks.SharedPoolRepo, dagExecutor *executormocks.RemoteExecutor, scalemgr *bractionmocks.ScaleManager) {
				sharedPool := &entity.SharedPool{
					SharedPool: apiv1.SharedPool{
						Id:    "dag-error-pool",
						Phase: apiv1.SharedPool_AVAILABLE,
					},
				}
				sharedPoolRepo.On("Load", mock.Anything, "dag-error-pool", false).Return(sharedPool, nil)

				// Mock DAG executor error
				dagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(nil, status.Error(codes.Internal, "DAG query failed"))
			},
			expectedError: "failed to merge sharedPool status from dag task",
			expectedCode:  codes.Internal,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Create mocks
			sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
			dagExecutor := executormocks.NewRemoteExecutor(t)
			scalemgr := bractionmocks.NewScaleManager(t)

			// Setup mocks
			tt.setupMocks(sharedPoolRepo, dagExecutor, scalemgr)

			// Create server implementation
			server := &serverImpl{
				sharedPoolRepo: sharedPoolRepo,
				dagExecutor:    dagExecutor,
				scaleManager:   scalemgr,
			}

			// Execute test
			result, err := server.GetSharedPool(context.Background(), tt.sharedPoolId)

			switch {
			case tt.expectedError != "":
				require.Error(t, err)
				assert.Contains(t, err.Error(), tt.expectedError)
				if tt.expectedCode != codes.OK {
					assert.Equal(t, tt.expectedCode, status.Code(err))
				}
				assert.Nil(t, result)

			case tt.validate != nil:
				tt.validate(t, result, err)

			default:
				require.NoError(t, err)
				require.NotNil(t, result)
			}
		})
	}
}

func TestMergeCassandraProxyStatusFromRepo_UpdatesExistingComponent(t *testing.T) {
	sharedPool := &entity.SharedPool{
		SharedPool: apiv1.SharedPool{
			Id:    "pool-1",
			Phase: apiv1.SharedPool_AVAILABLE,
			ComponentGroups: []*apiv1.Component{
				{
					Id:            "c1",
					ComponentType: apiv1.ComponentType_CASSANDRA_PROXY,
					Phase:         apiv1.Component_CREATING,
					Labels: map[string]string{
						cassandraProxyHostIdLabelKey:   "old-host",
						cassandraProxyNodeSpecLabelKey: "old-spec",
					},
				},
			},
		},
	}
	cassandraProxyStateRepo := repomnocks.NewCassandraProxyStateRepo(t)
	cassandraProxyStateRepo.On("ListBySharedPoolID", mock.Anything, "pool-1").Return([]*entity.CassandraProxyState{{
		ClusterID:    "c1",
		SharedPoolID: "pool-1",
		HostID:       "stable-host-id",
		NodeSpec:     "8C16G",
		Phase:        entity.ComponentPhaseToCassandraProxyPhase(apiv1.Component_AVAILABLE),
		NodeQuotaEntity: &apiv1.NodeQuota{
			NodeSpec:  &apiv1.NodeSpec{Cpu: 8, MemoryGi: 16, AvailableMemoryMi: 14000},
			NodeCount: 2,
		},
	}}, nil).Once()
	server := &serverImpl{
		cassandraProxyStateRepo: cassandraProxyStateRepo,
	}

	err := server.mergeCassandraProxyStatusFromRepo(context.Background(), sharedPool)
	require.NoError(t, err)
	require.Len(t, sharedPool.GetComponentGroups(), 1)
	component := sharedPool.GetComponentGroups()[0]
	assert.Equal(t, apiv1.Component_AVAILABLE, component.GetPhase())
	assert.Equal(t, "stable-host-id", component.GetLabels()[cassandraProxyHostIdLabelKey])
	assert.Equal(t, "8C16G", component.GetLabels()[cassandraProxyNodeSpecLabelKey])
	assert.Equal(t, uint32(2), component.GetNodeQuota().GetNodeCount())
}

func TestMergeCassandraProxyStatusFromRepo_SynthesizesMissingComponent(t *testing.T) {
	sharedPool := &entity.SharedPool{
		SharedPool: apiv1.SharedPool{
			Id:    "pool-1",
			Phase: apiv1.SharedPool_AVAILABLE,
		},
	}
	cassandraProxyStateRepo := repomnocks.NewCassandraProxyStateRepo(t)
	cassandraProxyStateRepo.On("ListBySharedPoolID", mock.Anything, "pool-1").Return([]*entity.CassandraProxyState{{
		ClusterID:    "c1",
		SharedPoolID: "pool-1",
		HostID:       "stable-host-id",
		NodeSpec:     "8C16G",
		Phase:        entity.ComponentPhaseToCassandraProxyPhase(apiv1.Component_AVAILABLE),
		NodeQuotaEntity: &apiv1.NodeQuota{
			NodeSpec:  &apiv1.NodeSpec{Cpu: 8, MemoryGi: 16, AvailableMemoryMi: 14000},
			NodeCount: 2,
		},
	}}, nil).Once()
	server := &serverImpl{
		cassandraProxyStateRepo: cassandraProxyStateRepo,
	}

	err := server.mergeCassandraProxyStatusFromRepo(context.Background(), sharedPool)
	require.NoError(t, err)
	require.Len(t, sharedPool.GetComponentGroups(), 1)
	component := sharedPool.GetComponentGroups()[0]
	assert.Equal(t, "c1", component.GetId())
	assert.Equal(t, apiv1.ComponentType_CASSANDRA_PROXY, component.GetComponentType())
	assert.Equal(t, apiv1.Component_AVAILABLE, component.GetPhase())
	assert.Equal(t, "stable-host-id", component.GetLabels()[cassandraProxyHostIdLabelKey])
	assert.Equal(t, "8C16G", component.GetLabels()[cassandraProxyNodeSpecLabelKey])
}

func TestServerImpl_mergeSharedPoolStatusFromRestore(t *testing.T) {
	tests := []struct {
		name          string
		sharedPool    *entity.SharedPool
		setupMocks    func(*bractionmocks.ScaleManager)
		expectedPhase apiv1.SharedPool_Phase
		expectedError bool
		errorContains string
		description   string
	}{
		{
			name: "no scaling decisions",
			sharedPool: &entity.SharedPool{
				SharedPool: apiv1.SharedPool{
					Id: "pool-1",
					ComponentGroups: []*apiv1.Component{
						{
							ComponentType: apiv1.ComponentType_TIKV,
							Phase:         apiv1.Component_AVAILABLE,
						},
					},
					Phase: apiv1.SharedPool_AVAILABLE,
				},
			},
			setupMocks: func(scaleManager *bractionmocks.ScaleManager) {
				scaleManager.On("AnyFullfillingDecision", mock.Anything, "pool-1").Return(false, nil)
			},
			expectedPhase: apiv1.SharedPool_AVAILABLE,
			expectedError: false,
			description:   "Should not modify phase when no scaling decisions exist",
		},
		{
			name: "has scaling decision",
			sharedPool: &entity.SharedPool{
				SharedPool: apiv1.SharedPool{
					Id: "pool-1",
					ComponentGroups: []*apiv1.Component{
						{
							ComponentType: apiv1.ComponentType_TIKV,
							Phase:         apiv1.Component_AVAILABLE,
						},
						{
							ComponentType: apiv1.ComponentType_TIDB,
							Phase:         apiv1.Component_AVAILABLE,
						},
					},
					Phase: apiv1.SharedPool_AVAILABLE,
				},
			},
			setupMocks: func(scaleManager *bractionmocks.ScaleManager) {
				scaleManager.On("AnyFullfillingDecision", mock.Anything, "pool-1").Return(true, nil)
			},
			expectedPhase: apiv1.SharedPool_MODIFYING,
			expectedError: false,
			description:   "Should set phase to MODIFYING when scaling decision exists",
		},
		{
			name: "scaling decision check error",
			sharedPool: &entity.SharedPool{
				SharedPool: apiv1.SharedPool{
					Id: "pool-1",
					ComponentGroups: []*apiv1.Component{
						{
							ComponentType: apiv1.ComponentType_TIKV,
							Phase:         apiv1.Component_AVAILABLE,
						},
					},
					Phase: apiv1.SharedPool_AVAILABLE,
				},
			},
			setupMocks: func(scaleManager *bractionmocks.ScaleManager) {
				scaleManager.On("AnyFullfillingDecision", mock.Anything, "pool-1").Return(false, assert.AnError)
			},
			expectedPhase: apiv1.SharedPool_AVAILABLE,
			expectedError: true,
			errorContains: "failed to check if there are fullfilling scale decisions",
			description:   "Should return error when scaling decision check fails",
		},
		{
			name: "tikv component phase updated to modifying",
			sharedPool: &entity.SharedPool{
				SharedPool: apiv1.SharedPool{
					Id: "pool-1",
					ComponentGroups: []*apiv1.Component{
						{
							ComponentType: apiv1.ComponentType_TIKV,
							Phase:         apiv1.Component_AVAILABLE,
						},
						{
							ComponentType: apiv1.ComponentType_PD,
							Phase:         apiv1.Component_AVAILABLE,
						},
						{
							ComponentType: apiv1.ComponentType_TIDB,
							Phase:         apiv1.Component_AVAILABLE,
						},
					},
					Phase: apiv1.SharedPool_AVAILABLE,
				},
			},
			setupMocks: func(scaleManager *bractionmocks.ScaleManager) {
				scaleManager.On("AnyFullfillingDecision", mock.Anything, "pool-1").Return(true, nil)
			},
			expectedPhase: apiv1.SharedPool_MODIFYING,
			expectedError: false,
			description:   "Should update only TIKV component phase to MODIFYING while leaving others unchanged",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Create mock
			scaleManager := bractionmocks.NewScaleManager(t)

			// Setup mock expectations
			tt.setupMocks(scaleManager)

			// Create server
			server := &serverImpl{
				scaleManager: scaleManager,
			}

			// Call the method
			ctx := context.Background()
			err := server.mergeSharedPoolStatusFromRestore(ctx, tt.sharedPool)

			// Check expectations
			if tt.expectedError {
				assert.Error(t, err, tt.description)
				if tt.errorContains != "" {
					assert.Contains(t, err.Error(), tt.errorContains, tt.description)
				}
				return
			}

			require.NoError(t, err, tt.description)

			// Check shared pool phase
			assert.Equal(t, tt.expectedPhase, tt.sharedPool.Phase, "Shared pool phase should match expected")

			// If expected phase is MODIFYING, check that TIKV component phase is also MODIFYING
			if tt.expectedPhase == apiv1.SharedPool_MODIFYING {
				tikvComponentFound := false
				for _, component := range tt.sharedPool.GetComponentGroups() {
					if component.GetComponentType() == apiv1.ComponentType_TIKV {
						tikvComponentFound = true
						assert.Equal(t, apiv1.Component_MODIFYING, component.Phase, "TIKV component phase should be MODIFYING")
						break
					}
				}
				assert.True(t, tikvComponentFound, "TIKV component should exist in shared pool")
			}
		})
	}
}

func TestGetSharedPool_OverlayAffectedTypesPriority(t *testing.T) {
	tests := []struct {
		name         string
		sharedPoolId string
		setupMocks   func(*repomnocks.SharedPoolRepo, *executormocks.RemoteExecutor, *bractionmocks.ScaleManager)
		validate     func(*testing.T, *entity.SharedPool, error)
	}{
		{
			name:         "affected types priority - both TIKV and COPROCESSOR_WORKER exposed",
			sharedPoolId: "testsharedpool-overlay",
			setupMocks: func(sharedPoolRepo *repomnocks.SharedPoolRepo, dagExecutor *executormocks.RemoteExecutor, scalemgr *bractionmocks.ScaleManager) {
				sharedPool := &entity.SharedPool{
					SharedPool: apiv1.SharedPool{
						Id:    "testsharedpool-overlay",
						Phase: apiv1.SharedPool_AVAILABLE,
						ComponentGroups: []*apiv1.Component{
							{
								Id:            "tikv-1",
								ComponentType: apiv1.ComponentType_TIKV,
								Phase:         apiv1.Component_AVAILABLE,
							},
							{
								Id:            "copr-1",
								ComponentType: apiv1.ComponentType_COPROCESSOR_WORKER,
								Phase:         apiv1.Component_AVAILABLE,
							},
						},
					},
				}
				sharedPoolRepo.On("Load", mock.Anything, "testsharedpool-overlay", false).Return(sharedPool, nil)

				// Mock running enable-coprocessor-worker DAG with AffectedComponentTypesKey
				tempSharedPool := &entity.SharedPool{
					SharedPool: apiv1.SharedPool{
						Id:    "testsharedpool-overlay",
						Phase: apiv1.SharedPool_MODIFYING,
						ComponentGroups: []*apiv1.Component{
							{
								Id:            "tikv-1",
								ComponentType: apiv1.ComponentType_TIKV,
								Phase:         apiv1.Component_MODIFYING,
							},
							{
								Id:            "copr-1",
								ComponentType: apiv1.ComponentType_COPROCESSOR_WORKER,
								Phase:         apiv1.Component_CREATING,
							},
						},
					},
				}
				sharedPoolStr, _ := protojson.Marshal(tempSharedPool)

				runningDags := executor.NewExecTaskIdStatus()
				runningDags["dag-1"] = &executor.ExecTaskStatusEx{
					ExecTaskStatus: executor.ExecTaskStatus{
						ExecTask: &executor.ExecutorTask{
							RootName: dag.PlanEnableCoprocessorWorker,
							Args: actions.ActionArgs{
								workflow.SharedPoolKey:             actions.ActionArgString(string(sharedPoolStr)),
								workflow.ComponentIdsKey:           actions.ActionArgString("copr-1"),
								workflow.ComponentTypesKey:         actions.ActionArgString("coprocessor-tikv-worker"),
								workflow.AffectedComponentTypesKey: actions.ActionArgString("tikv,coprocessor-tikv-worker"),
							},
						},
					},
				}

				dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanEnableCoprocessorWorker, mock.AnythingOfType("executor.MatchFilter")).Return(runningDags, nil)
				dagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil)
				scalemgr.On("AnyFullfillingDecision", mock.Anything, sharedPool.Id).Return(false, nil)
			},
			validate: func(t *testing.T, result *entity.SharedPool, err error) {
				t.Helper()
				require.NoError(t, err)
				require.NotNil(t, result)
				assert.Equal(t, apiv1.SharedPool_MODIFYING, result.GetPhase())

				// Both TIKV and COPROCESSOR_WORKER should be exposed as MODIFYING/CREATING
				tikvFound := false
				coprFound := false
				for _, comp := range result.GetComponentGroups() {
					if comp.GetComponentType() == apiv1.ComponentType_TIKV {
						tikvFound = true
						assert.Equal(t, apiv1.Component_MODIFYING, comp.GetPhase(), "TIKV should be MODIFYING via affected types")
					}
					if comp.GetComponentType() == apiv1.ComponentType_COPROCESSOR_WORKER {
						coprFound = true
						assert.Equal(t, apiv1.Component_CREATING, comp.GetPhase(), "COPROCESSOR_WORKER should be CREATING")
					}
				}
				assert.True(t, tikvFound, "TIKV should be exposed via AffectedComponentTypesKey")
				assert.True(t, coprFound, "COPROCESSOR_WORKER should be exposed")
			},
		},
		{
			name:         "old DAG fallback - no AffectedComponentTypesKey uses ComponentTypesKey",
			sharedPoolId: "testsharedpool-fallback",
			setupMocks: func(sharedPoolRepo *repomnocks.SharedPoolRepo, dagExecutor *executormocks.RemoteExecutor, scalemgr *bractionmocks.ScaleManager) {
				sharedPool := &entity.SharedPool{
					SharedPool: apiv1.SharedPool{
						Id:    "testsharedpool-fallback",
						Phase: apiv1.SharedPool_AVAILABLE,
						ComponentGroups: []*apiv1.Component{
							{
								Id:            "tidb-1",
								ComponentType: apiv1.ComponentType_TIDB,
								Phase:         apiv1.Component_AVAILABLE,
							},
							{
								Id:            "tikv-1",
								ComponentType: apiv1.ComponentType_TIKV,
								Phase:         apiv1.Component_AVAILABLE,
							},
						},
					},
				}
				sharedPoolRepo.On("Load", mock.Anything, "testsharedpool-fallback", false).Return(sharedPool, nil)

				// Mock running update-component DAG without AffectedComponentTypesKey (old DAG)
				tempSharedPool := &entity.SharedPool{
					SharedPool: apiv1.SharedPool{
						Id:    "testsharedpool-fallback",
						Phase: apiv1.SharedPool_MODIFYING,
						ComponentGroups: []*apiv1.Component{
							{
								Id:            "tidb-1",
								ComponentType: apiv1.ComponentType_TIDB,
								Phase:         apiv1.Component_MODIFYING,
							},
						},
					},
				}
				sharedPoolStr, _ := protojson.Marshal(tempSharedPool)

				runningDags := executor.NewExecTaskIdStatus()
				runningDags["dag-1"] = &executor.ExecTaskStatusEx{
					ExecTaskStatus: executor.ExecTaskStatus{
						ExecTask: &executor.ExecutorTask{
							RootName: dag.PlanUpdateComponent,
							Args: actions.ActionArgs{
								workflow.SharedPoolKey:     actions.ActionArgString(string(sharedPoolStr)),
								workflow.ComponentIdsKey:   actions.ActionArgString("tidb-1"),
								workflow.ComponentTypesKey: actions.ActionArgString("tidb"),
								// No AffectedComponentTypesKey - old DAG fallback
							},
						},
					},
				}

				dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanUpdateComponent, mock.AnythingOfType("executor.MatchFilter")).Return(runningDags, nil)
				dagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil)
				scalemgr.On("AnyFullfillingDecision", mock.Anything, sharedPool.Id).Return(false, nil)
			},
			validate: func(t *testing.T, result *entity.SharedPool, err error) {
				t.Helper()
				require.NoError(t, err)
				require.NotNil(t, result)
				assert.Equal(t, apiv1.SharedPool_MODIFYING, result.GetPhase())

				// Only TIDB should be exposed (mutation target), TIKV should remain AVAILABLE
				tidbFound := false
				tikvFound := false
				for _, comp := range result.GetComponentGroups() {
					if comp.GetComponentType() == apiv1.ComponentType_TIDB {
						tidbFound = true
						assert.Equal(t, apiv1.Component_MODIFYING, comp.GetPhase(), "TIDB should be MODIFYING")
					}
					if comp.GetComponentType() == apiv1.ComponentType_TIKV {
						tikvFound = true
						assert.Equal(t, apiv1.Component_AVAILABLE, comp.GetPhase(), "TIKV should remain AVAILABLE (not in mutation target)")
					}
				}
				assert.True(t, tidbFound, "TIDB should be exposed")
				assert.True(t, tikvFound, "TIKV should exist but remain AVAILABLE")
			},
		},
		{
			name:         "disable direction - TIKV MODIFYING and COPROCESSOR_WORKER DELETING, TIKV_WORKER untouched",
			sharedPoolId: "testsharedpool-disable-overlay",
			setupMocks: func(sharedPoolRepo *repomnocks.SharedPoolRepo, dagExecutor *executormocks.RemoteExecutor, scalemgr *bractionmocks.ScaleManager) {
				sharedPool := &entity.SharedPool{
					SharedPool: apiv1.SharedPool{
						Id:    "testsharedpool-disable-overlay",
						Phase: apiv1.SharedPool_AVAILABLE,
						ComponentGroups: []*apiv1.Component{
							{
								Id:            "tikv-1",
								ComponentType: apiv1.ComponentType_TIKV,
								Phase:         apiv1.Component_AVAILABLE,
							},
							{
								Id:            "worker-1",
								ComponentType: apiv1.ComponentType_TIKV_WORKER,
								Phase:         apiv1.Component_AVAILABLE,
							},
							{
								Id:            "copr-1",
								ComponentType: apiv1.ComponentType_COPROCESSOR_WORKER,
								Phase:         apiv1.Component_AVAILABLE,
							},
						},
					},
				}
				sharedPoolRepo.On("Load", mock.Anything, "testsharedpool-disable-overlay", false).Return(sharedPool, nil)

				// Mock running disable-coprocessor-worker DAG: TIKV is being
				// repointed (MODIFYING) and the coprocessor worker is being
				// torn down (DELETING); the default tikv worker is not part of
				// the switch and must not be overlaid.
				tempSharedPool := &entity.SharedPool{
					SharedPool: apiv1.SharedPool{
						Id:    "testsharedpool-disable-overlay",
						Phase: apiv1.SharedPool_MODIFYING,
						ComponentGroups: []*apiv1.Component{
							{
								Id:            "tikv-1",
								ComponentType: apiv1.ComponentType_TIKV,
								Phase:         apiv1.Component_MODIFYING,
							},
							{
								Id:            "worker-1",
								ComponentType: apiv1.ComponentType_TIKV_WORKER,
								Phase:         apiv1.Component_AVAILABLE,
							},
							{
								Id:            "copr-1",
								ComponentType: apiv1.ComponentType_COPROCESSOR_WORKER,
								Phase:         apiv1.Component_DELETING,
							},
						},
					},
				}
				sharedPoolStr, _ := protojson.Marshal(tempSharedPool)

				runningDags := executor.NewExecTaskIdStatus()
				runningDags["dag-1"] = &executor.ExecTaskStatusEx{
					ExecTaskStatus: executor.ExecTaskStatus{
						ExecTask: &executor.ExecutorTask{
							RootName: dag.PlanDisableCoprocessorWorker,
							Args: actions.ActionArgs{
								workflow.SharedPoolKey:             actions.ActionArgString(string(sharedPoolStr)),
								workflow.ComponentIdsKey:           actions.ActionArgString("copr-1"),
								workflow.ComponentTypesKey:         actions.ActionArgString("coprocessor-tikv-worker"),
								workflow.AffectedComponentTypesKey: actions.ActionArgString("tikv,coprocessor-tikv-worker"),
							},
						},
					},
				}

				dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanDisableCoprocessorWorker, mock.AnythingOfType("executor.MatchFilter")).Return(runningDags, nil)
				dagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil)
				scalemgr.On("AnyFullfillingDecision", mock.Anything, sharedPool.Id).Return(false, nil)
			},
			validate: func(t *testing.T, result *entity.SharedPool, err error) {
				t.Helper()
				require.NoError(t, err)
				require.NotNil(t, result)
				assert.Equal(t, apiv1.SharedPool_MODIFYING, result.GetPhase())

				tikvFound := false
				workerFound := false
				coprFound := false
				for _, comp := range result.GetComponentGroups() {
					switch comp.GetComponentType() {
					case apiv1.ComponentType_TIKV:
						tikvFound = true
						assert.Equal(t, apiv1.Component_MODIFYING, comp.GetPhase(), "TIKV should be MODIFYING via affected types")
					case apiv1.ComponentType_TIKV_WORKER:
						workerFound = true
						assert.Equal(t, apiv1.Component_AVAILABLE, comp.GetPhase(), "TIKV_WORKER must not be overlaid by the switch")
					case apiv1.ComponentType_COPROCESSOR_WORKER:
						coprFound = true
						assert.Equal(t, apiv1.Component_DELETING, comp.GetPhase(), "COPROCESSOR_WORKER should be DELETING")
					}
				}
				assert.True(t, tikvFound, "TIKV should be exposed via AffectedComponentTypesKey")
				assert.True(t, workerFound, "TIKV_WORKER should exist but stay AVAILABLE")
				assert.True(t, coprFound, "COPROCESSOR_WORKER should be exposed")
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
			dagExecutor := executormocks.NewRemoteExecutor(t)
			scalemgr := bractionmocks.NewScaleManager(t)

			tt.setupMocks(sharedPoolRepo, dagExecutor, scalemgr)

			server := &serverImpl{
				sharedPoolRepo: sharedPoolRepo,
				dagExecutor:    dagExecutor,
				scaleManager:   scalemgr,
			}

			ctx := context.Background()
			result, err := server.GetSharedPool(ctx, tt.sharedPoolId)

			tt.validate(t, result, err)
		})
	}
}