package server

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"
	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/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"
)

func TestApplySharedPool(t *testing.T) {
	tests := []struct {
		name           string
		sharedPoolId   string
		applyComponent *entity.TargetComponent
		force          bool
		setupMocks     func(*repomnocks.SharedPoolRepo, *repomnocks.LockRepo, *executormocks.RemoteExecutor, *bractionmocks.ScaleManager)
		expectedError  string
		expectedCode   codes.Code
	}{
		{
			name:         "successful apply with force=false",
			sharedPoolId: "testsharedpoolid",
			applyComponent: &entity.TargetComponent{
				Id:            "tidb-1",
				ComponentType: apiv1.ComponentType_TIDB,
			},
			force: false,
			setupMocks: func(sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) {
				sharedPool := &entity.SharedPool{
					SharedPool: apiv1.SharedPool{
						Id: "testsharedpoolid",
						ComponentGroups: []*apiv1.Component{
							{
								Id:            "tidb-1",
								ComponentType: apiv1.ComponentType_TIDB,
								Phase:         apiv1.Component_AVAILABLE,
							},
						},
					},
				}
				sharedPoolRepo.On("Load", mock.Anything, "testsharedpoolid", false).Return(sharedPool, nil)

				// Mock merge status (no running DAGs)
				dagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil)
				// Mock no restore
				scaleManager.On("AnyFullfillingDecision", mock.Anything, mock.Anything).Return(false, nil)

				// Mock lock/unlock operations
				lockRepo.On("Lock", mock.Anything, mock.MatchedBy(func(lock *entity.Lock) bool {
					return lock.Key == "testsharedpoolid"
				})).Return(func() error { return nil }, nil)

				// Mock DAG submit
				dagExecutor.On("Submit", mock.Anything, dag.PlanApplyComponent, mock.Anything, mock.Anything).Return(nil)
			},
		},
		{
			name:         "successful apply with force=true",
			sharedPoolId: "testsharedpoolid",
			applyComponent: &entity.TargetComponent{
				Id:            "tidb-1",
				ComponentType: apiv1.ComponentType_TIDB,
			},
			force: true,
			setupMocks: func(sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) {
				sharedPool := &entity.SharedPool{
					SharedPool: apiv1.SharedPool{
						Id: "testsharedpoolid",
						ComponentGroups: []*apiv1.Component{
							{
								Id:            "tidb-1",
								ComponentType: apiv1.ComponentType_TIDB,
								Phase:         apiv1.Component_MODIFYING, // Not available, but force=true
							},
						},
					},
				}
				sharedPoolRepo.On("Load", mock.Anything, "testsharedpoolid", false).Return(sharedPool, nil)

				// Mock merge status (no running DAGs)
				dagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil)
				// Mock no restore
				scaleManager.On("AnyFullfillingDecision", mock.Anything, mock.Anything).Return(false, nil)

				// Mock lock/unlock operations
				lockRepo.On("Lock", mock.Anything, mock.MatchedBy(func(lock *entity.Lock) bool {
					return lock.Key == "testsharedpoolid"
				})).Return(func() error { return nil }, nil)

				// Mock DAG submit
				dagExecutor.On("Submit", mock.Anything, dag.PlanApplyComponent, mock.Anything, mock.Anything).Return(nil)
			},
		},
		{
			name:           "empty shared pool id",
			sharedPoolId:   "",
			applyComponent: nil,
			force:          false,
			setupMocks: func(sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) {
				// No mocks needed for this test case
			},
			expectedError: "sharedPool id is required",
			expectedCode:  codes.InvalidArgument,
		},
		{
			name:           "nil apply component",
			sharedPoolId:   "testsharedpoolid",
			applyComponent: nil,
			force:          false,
			setupMocks: func(sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) {
				// No mocks needed for this test case
			},
			expectedError: "apply component is required",
			expectedCode:  codes.InvalidArgument,
		},
		{
			name:         "unspecified component type",
			sharedPoolId: "testsharedpoolid",
			applyComponent: &entity.TargetComponent{
				Id:            "tidb-1",
				ComponentType: apiv1.ComponentType_COMPONENT_TYPE_UNSPECIFIED,
			},
			force: false,
			setupMocks: func(sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) {
				// No mocks needed for this test case
			},
			expectedError: "apply component type is required",
			expectedCode:  codes.InvalidArgument,
		},
		{
			name:         "shared pool not found",
			sharedPoolId: "nonexistentid",
			applyComponent: &entity.TargetComponent{
				Id:            "tidb-1",
				ComponentType: apiv1.ComponentType_TIDB,
			},
			force: false,
			setupMocks: func(sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) {
				lockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil)
				sharedPoolRepo.On("Load", mock.Anything, "nonexistentid", false).Return(nil, status.Error(codes.NotFound, "not found"))
			},
			expectedError: "sharedPool nonexistentid not found",
			expectedCode:  codes.NotFound,
		},
		{
			name:         "repository load error",
			sharedPoolId: "testsharedpoolid",
			applyComponent: &entity.TargetComponent{
				Id:            "tidb-1",
				ComponentType: apiv1.ComponentType_TIDB,
			},
			force: false,
			setupMocks: func(sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) {
				lockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil)
				sharedPoolRepo.On("Load", mock.Anything, "testsharedpoolid", false).Return(nil, errors.New("database error"))
			},
			expectedError: "failed to load sharedPool: database error",
			expectedCode:  codes.Internal,
		},
		{
			name:         "component not available and force=false",
			sharedPoolId: "testsharedpoolid",
			applyComponent: &entity.TargetComponent{
				Id:            "tidb-1",
				ComponentType: apiv1.ComponentType_TIDB,
			},
			force: false,
			setupMocks: func(sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) {
				sharedPool := &entity.SharedPool{
					SharedPool: apiv1.SharedPool{
						Id: "testsharedpoolid",
						ComponentGroups: []*apiv1.Component{
							{
								Id:            "tidb-1",
								ComponentType: apiv1.ComponentType_TIDB,
								Phase:         apiv1.Component_MODIFYING, // Not available
							},
						},
					},
				}
				sharedPoolRepo.On("Load", mock.Anything, "testsharedpoolid", false).Return(sharedPool, nil)

				// Mock merge status (no running DAGs)
				dagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil)
				// Mock no restore
				scaleManager.On("AnyFullfillingDecision", mock.Anything, mock.Anything).Return(false, nil)

				// Mock lock/unlock operations
				lockRepo.On("Lock", mock.Anything, mock.MatchedBy(func(lock *entity.Lock) bool {
					return lock.Key == "testsharedpoolid"
				})).Return(func() error { return nil }, nil)
			},
			expectedError: "target component tidb tidb-1 is not available",
			expectedCode:  codes.FailedPrecondition,
		},
		{
			name:         "no component to apply",
			sharedPoolId: "testsharedpoolid",
			applyComponent: &entity.TargetComponent{
				Id:            "non-existent",
				ComponentType: apiv1.ComponentType_TIDB,
			},
			force: false,
			setupMocks: func(sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) {
				sharedPool := &entity.SharedPool{
					SharedPool: apiv1.SharedPool{
						Id: "testsharedpoolid",
						ComponentGroups: []*apiv1.Component{
							{
								Id:            "tidb-1", // Different ID
								ComponentType: apiv1.ComponentType_TIDB,
								Phase:         apiv1.Component_AVAILABLE,
							},
						},
					},
				}
				sharedPoolRepo.On("Load", mock.Anything, "testsharedpoolid", false).Return(sharedPool, nil)

				// Mock merge status (no running DAGs)
				dagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil)
				// Mock no restore
				scaleManager.On("AnyFullfillingDecision", mock.Anything, mock.Anything).Return(false, nil)

				// Mock lock/unlock operations
				lockRepo.On("Lock", mock.Anything, mock.MatchedBy(func(lock *entity.Lock) bool {
					return lock.Key == "testsharedpoolid"
				})).Return(func() error { return nil }, nil)
			},
			expectedError: "component tidb non-existent not found",
			expectedCode:  codes.NotFound,
		},
		{
			name:         "lock failure",
			sharedPoolId: "testsharedpoolid",
			applyComponent: &entity.TargetComponent{
				Id:            "tidb-1",
				ComponentType: apiv1.ComponentType_TIDB,
			},
			force: false,
			setupMocks: func(sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) {
				// Mock lock failure - since lock fails, other methods won't be called
				lockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, errors.New("lock failed"))
			},
			expectedError: "lock failed",
		},
		{
			name:         "dag submit failure",
			sharedPoolId: "testsharedpoolid",
			applyComponent: &entity.TargetComponent{
				Id:            "tidb-1",
				ComponentType: apiv1.ComponentType_TIDB,
			},
			force: false,
			setupMocks: func(sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) {
				sharedPool := &entity.SharedPool{
					SharedPool: apiv1.SharedPool{
						Id: "testsharedpoolid",
						ComponentGroups: []*apiv1.Component{
							{
								Id:            "tidb-1",
								ComponentType: apiv1.ComponentType_TIDB,
								Phase:         apiv1.Component_AVAILABLE,
							},
						},
					},
				}
				sharedPoolRepo.On("Load", mock.Anything, "testsharedpoolid", false).Return(sharedPool, nil)

				// Mock merge status (no running DAGs)
				dagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil)
				// Mock no restore
				scaleManager.On("AnyFullfillingDecision", mock.Anything, mock.Anything).Return(false, nil)

				// Mock lock/unlock operations
				lockRepo.On("Lock", mock.Anything, mock.MatchedBy(func(lock *entity.Lock) bool {
					return lock.Key == "testsharedpoolid"
				})).Return(func() error { return nil }, nil)

				// Mock DAG submit failure
				dagExecutor.On("Submit", mock.Anything, dag.PlanApplyComponent, mock.Anything, mock.Anything).Return(errors.New("submit failed"))
			},
			expectedError: "failed to submit apply sharedpool dag: submit failed",
			expectedCode:  codes.Internal,
		},
	}

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

			// Setup mocks
			tt.setupMocks(sharedPoolRepo, lockRepo, dagExecutor, scaleManager)

			// Create server implementation
			server := &serverImpl{
				sharedPoolRepo: sharedPoolRepo,
				lockRepo:       lockRepo,
				dagExecutor:    dagExecutor,
				transactioner:  &mockTransactioner{},
				scaleManager:   scaleManager,
			}

			// Execute test
			err := server.ApplySharedPool(context.Background(), tt.sharedPoolId, tt.applyComponent, tt.force)

			if 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))
				}
			} else {
				require.NoError(t, err)
			}
		})
	}
}

func TestApplySharedPoolForceRejectsRunningDeleteTiKVGroup(t *testing.T) {
	ctx := context.Background()
	sharedPool := deleteTiKVGroupTestSharedPool()
	deletingSharedPool := deleteTiKVGroupTestSharedPool()
	deletingSharedPool.Id = sharedPool.GetId()
	deletingSharedPool.Phase = apiv1.SharedPool_MODIFYING
	findComponentByTypeAndId(deletingSharedPool, apiv1.ComponentType_TIKV, "blue").Phase = apiv1.Component_DELETING

	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	lockRepo := repomnocks.NewLockRepo(t)
	dagExecutor := executormocks.NewRemoteExecutor(t)
	scaleManager := bractionmocks.NewScaleManager(t)
	lockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil)
	sharedPoolRepo.On("Load", mock.Anything, sharedPool.GetId(), false).Return(sharedPool, nil)
	setupPlacementPolicyDAGMocks(dagExecutor, map[string]executor.ExecTaskIdStatus{
		dag.PlanDeleteTiKVGroup: placementPolicyRunningSharedPoolStatus(t, deletingSharedPool),
	})
	scaleManager.On("AnyFullfillingDecision", mock.Anything, sharedPool.GetId()).Return(false, nil).Maybe()

	server := &serverImpl{
		sharedPoolRepo: sharedPoolRepo,
		lockRepo:       lockRepo,
		dagExecutor:    dagExecutor,
		transactioner:  &mockTransactioner{},
		scaleManager:   scaleManager,
	}

	err := server.ApplySharedPool(ctx, sharedPool.GetId(), &entity.TargetComponent{
		Id:            "blue",
		ComponentType: apiv1.ComponentType_TIKV,
	}, true)
	require.Error(t, err)
	assert.Equal(t, codes.FailedPrecondition, status.Code(err))
	assert.Contains(t, err.Error(), "target component tikv blue is deleting")
	dagExecutor.AssertNotCalled(t, "Submit", mock.Anything, dag.PlanApplyComponent, mock.Anything, mock.Anything)
}

func TestApplySharedPool_GuardBlocksOnRunningTiKVRemoteWorkerDAG(t *testing.T) {
	for _, runningPlan := range tikvRemoteWorkersDAGPlans() {
		t.Run(runningPlan, func(t *testing.T) {
			sharedPool := newTKRWSharedPool(tkrwEnabledBoth(t), 3)
			sharedPoolRepo, lockRepo, dagExecutor, scaleManager := newTKRWHandlerMocksWithRunningDag(
				t, sharedPool, runningPlan, runningDagStatus(tkrwKey(t, tkrwEnabledBoth(t))),
			)
			server := &serverImpl{sharedPoolRepo: sharedPoolRepo, lockRepo: lockRepo, dagExecutor: dagExecutor, scaleManager: scaleManager, transactioner: &mockTKRWTransactioner{}}

			err := server.ApplySharedPool(context.Background(), testTKRWSharedPoolID, &entity.TargetComponent{Id: testTKRWComponentID, ComponentType: apiv1.ComponentType_TIKV}, false)
			require.Error(t, err)
			assert.Equal(t, codes.FailedPrecondition, status.Code(err))
			dagExecutor.AssertNotCalled(t, "Submit", mock.Anything, dag.PlanApplyComponent, mock.Anything, mock.Anything)
		})
	}
}

func TestApplySharedPool_GuardBlocksOnTerminalTiKVRemoteWorkerDrift(t *testing.T) {
	sharedPool := newTKRWSharedPool(tkrwEnabledBoth(t), 3)
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	lockRepo := repomnocks.NewLockRepo(t)
	dagExecutor := executormocks.NewRemoteExecutor(t)
	scaleManager := bractionmocks.NewScaleManager(t)
	sharedPoolRepo.On("Load", mock.Anything, testTKRWSharedPoolID, false).Return(sharedPool, nil)
	lockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil)
	setupAllDAGMocks(dagExecutor)
	scaleManager.On("AnyFullfillingDecision", mock.Anything, mock.Anything).Return(false, nil).Maybe()

	server := &serverImpl{sharedPoolRepo: sharedPoolRepo, lockRepo: lockRepo, dagExecutor: dagExecutor, scaleManager: scaleManager, dataplaneManager: newTKRWDataplaneManager(t), transactioner: &mockTKRWTransactioner{}}
	err := server.ApplySharedPool(context.Background(), testTKRWSharedPoolID, &entity.TargetComponent{Id: testTKRWComponentID, ComponentType: apiv1.ComponentType_TIKV}, false)
	require.Error(t, err)
	assert.Equal(t, codes.FailedPrecondition, status.Code(err))
	dagExecutor.AssertNotCalled(t, "Submit", mock.Anything, dag.PlanApplyComponent, mock.Anything, mock.Anything)
}