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"
	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"
	repomocks "github.com/tidbcloud/cluster-service-ng/pkg/storage/repo/mocks"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

// Mock transactioner for testing
type mockQuotaTransactioner struct{}

func (m *mockQuotaTransactioner) Transaction(ctx context.Context, fn func(context.Context) error) error {
	return fn(ctx)
}

func TestSetSharedPoolComponentQuota_BasicValidation(t *testing.T) {
	testCases := []struct {
		name            string
		sharedPoolId    string
		targetComponent *entity.Component
		expectedError   string
		expectedCode    codes.Code
	}{
		{
			name:            "empty shared pool ID",
			sharedPoolId:    "",
			targetComponent: &entity.Component{},
			expectedError:   "sharedPool id is required",
			expectedCode:    codes.InvalidArgument,
		},
		{
			name:            "nil target component",
			sharedPoolId:    "testpool",
			targetComponent: nil,
			expectedError:   "target components is empty",
			expectedCode:    codes.InvalidArgument,
		},
		{
			name:         "unspecified component type",
			sharedPoolId: "testpool",
			targetComponent: &entity.Component{
				ComponentType: apiv1.ComponentType_COMPONENT_TYPE_UNSPECIFIED,
			},
			expectedError: "target component type is required",
			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: &mockQuotaTransactioner{},
			}

			// Execute test
			result, err := server.SetSharedPoolComponentQuota(context.Background(), tc.sharedPoolId, tc.targetComponent)

			// 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 TestSetSharedPoolComponentQuota_SuccessfulUpdate(t *testing.T) {
	// Create mocks
	mockSharedPoolRepo := repomocks.NewSharedPoolRepo(t)
	mockLockRepo := repomocks.NewLockRepo(t)
	mockDagExecutor := executormocks.NewRemoteExecutor(t)
	mockScaleManager := bractionmocks.NewScaleManager(t)

	// Test data
	targetComponent := &entity.Component{
		Id:            "test-comp",
		ComponentType: apiv1.ComponentType_TIDB,
		NodeQuota: &apiv1.NodeQuota{
			MaxNodeCount: 10,
			MinNodeCount: 2,
		},
	}

	sharedPool := &entity.SharedPool{
		SharedPool: apiv1.SharedPool{
			Id:    "testpool",
			Phase: apiv1.SharedPool_AVAILABLE,
			ComponentGroups: []*entity.Component{
				{
					Id:            "test-comp",
					ComponentType: apiv1.ComponentType_TIDB,
					Phase:         apiv1.Component_AVAILABLE,
					NodeQuota: &apiv1.NodeQuota{
						NodeCount:    3,
						MinNodeCount: 1,
						MaxNodeCount: 5,
					},
				},
			},
		},
	}

	// Setup mocks
	mockSharedPoolRepo.On("Load", mock.Anything, "testpool", false).Return(sharedPool, nil)
	mockLockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil)
	mockDagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil)
	mockScaleManager.On("AnyFullfillingDecision", mock.Anything, mock.Anything).Return(false, nil)
	mockDagExecutor.On("Submit", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)

	// Create server instance
	server := &serverImpl{
		sharedPoolRepo: mockSharedPoolRepo,
		lockRepo:       mockLockRepo,
		dagExecutor:    mockDagExecutor,
		transactioner:  &mockQuotaTransactioner{},
		scaleManager:   mockScaleManager,
	}

	// Execute test
	ctx := context.Background()
	result, err := server.SetSharedPoolComponentQuota(ctx, "testpool", targetComponent)

	// Verify results
	require.NoError(t, err)
	assert.NotNil(t, result)
	assert.Equal(t, targetComponent.Id, result.Id)
	assert.Equal(t, targetComponent.ComponentType, result.ComponentType)
	assert.Equal(t, targetComponent.NodeQuota.MaxNodeCount, result.NodeQuota.MaxNodeCount)
	assert.Equal(t, targetComponent.NodeQuota.MinNodeCount, result.NodeQuota.MinNodeCount)

	// Verify all expectations were met
	mockSharedPoolRepo.AssertExpectations(t)
	mockLockRepo.AssertExpectations(t)
	mockDagExecutor.AssertExpectations(t)
}

func TestSetSharedPoolComponentQuota_ValidationErrors(t *testing.T) {
	t.Run("missing node quota", func(t *testing.T) {
		mockSharedPoolRepo := repomocks.NewSharedPoolRepo(t)
		mockLockRepo := repomocks.NewLockRepo(t)
		mockDagExecutor := executormocks.NewRemoteExecutor(t)
		mockScaleManager := bractionmocks.NewScaleManager(t)

		// Mock shared pool - needed since Load is called before validation
		sharedPool := &entity.SharedPool{
			SharedPool: apiv1.SharedPool{
				Id:    "testpool",
				Phase: apiv1.SharedPool_AVAILABLE,
			},
		}
		mockSharedPoolRepo.On("Load", mock.Anything, "testpool", false).Return(sharedPool, nil)
		mockLockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil)
		mockDagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil)
		mockScaleManager.On("AnyFullfillingDecision", mock.Anything, "testpool").Return(false, nil)

		server := &serverImpl{
			sharedPoolRepo: mockSharedPoolRepo,
			lockRepo:       mockLockRepo,
			dagExecutor:    mockDagExecutor,
			scaleManager:   mockScaleManager,
			transactioner:  &mockQuotaTransactioner{},
		}

		targetComponent := &entity.Component{
			Id:            "test-comp",
			ComponentType: apiv1.ComponentType_TIDB,
			NodeQuota:     nil,
		}

		result, err := server.SetSharedPoolComponentQuota(context.Background(), "testpool", targetComponent)

		require.Error(t, err)
		assert.Contains(t, err.Error(), "target component node quota is required")
		assert.Nil(t, result)
	})

	t.Run("missing max node count", func(t *testing.T) {
		mockSharedPoolRepo := repomocks.NewSharedPoolRepo(t)
		mockLockRepo := repomocks.NewLockRepo(t)
		mockDagExecutor := executormocks.NewRemoteExecutor(t)
		mockScaleManager := bractionmocks.NewScaleManager(t)

		// Mock shared pool - needed since Load is called before validation
		sharedPool := &entity.SharedPool{
			SharedPool: apiv1.SharedPool{
				Id:    "testpool",
				Phase: apiv1.SharedPool_AVAILABLE,
			},
		}
		mockSharedPoolRepo.On("Load", mock.Anything, "testpool", false).Return(sharedPool, nil)
		mockLockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil)
		mockDagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil)
		mockScaleManager.On("AnyFullfillingDecision", mock.Anything, "testpool").Return(false, nil)

		server := &serverImpl{
			sharedPoolRepo: mockSharedPoolRepo,
			lockRepo:       mockLockRepo,
			dagExecutor:    mockDagExecutor,
			scaleManager:   mockScaleManager,
			transactioner:  &mockQuotaTransactioner{},
		}

		targetComponent := &entity.Component{
			Id:            "test-comp",
			ComponentType: apiv1.ComponentType_TIDB,
			NodeQuota: &apiv1.NodeQuota{
				MaxNodeCount: 0,
				MinNodeCount: 1,
			},
		}

		result, err := server.SetSharedPoolComponentQuota(context.Background(), "testpool", targetComponent)

		require.Error(t, err)
		assert.Contains(t, err.Error(), "target component max node count is required")
		assert.Nil(t, result)
	})

	t.Run("non-tiflash component with zero min node count", func(t *testing.T) {
		mockSharedPoolRepo := repomocks.NewSharedPoolRepo(t)
		mockLockRepo := repomocks.NewLockRepo(t)
		mockDagExecutor := executormocks.NewRemoteExecutor(t)
		mockScaleManager := bractionmocks.NewScaleManager(t)

		// Mock shared pool - needed since Load is called before validation
		sharedPool := &entity.SharedPool{
			SharedPool: apiv1.SharedPool{
				Id:    "testpool",
				Phase: apiv1.SharedPool_AVAILABLE,
			},
		}
		mockSharedPoolRepo.On("Load", mock.Anything, "testpool", false).Return(sharedPool, nil)
		mockLockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil)
		mockDagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil)
		mockScaleManager.On("AnyFullfillingDecision", mock.Anything, "testpool").Return(false, nil)

		server := &serverImpl{
			sharedPoolRepo: mockSharedPoolRepo,
			lockRepo:       mockLockRepo,
			dagExecutor:    mockDagExecutor,
			scaleManager:   mockScaleManager,
			transactioner:  &mockQuotaTransactioner{},
		}

		targetComponent := &entity.Component{
			Id:            "test-comp",
			ComponentType: apiv1.ComponentType_TIDB,
			NodeQuota: &apiv1.NodeQuota{
				MaxNodeCount: 5,
				MinNodeCount: 0,
			},
		}

		result, err := server.SetSharedPoolComponentQuota(context.Background(), "testpool", targetComponent)

		require.Error(t, err)
		assert.Contains(t, err.Error(), "target component min node count is required")
		assert.Nil(t, result)
	})

	t.Run("min node count greater than max node count", func(t *testing.T) {
		mockSharedPoolRepo := repomocks.NewSharedPoolRepo(t)
		mockLockRepo := repomocks.NewLockRepo(t)
		mockDagExecutor := executormocks.NewRemoteExecutor(t)
		mockScaleManager := bractionmocks.NewScaleManager(t)

		// Mock shared pool - needed since Load is called before validation
		sharedPool := &entity.SharedPool{
			SharedPool: apiv1.SharedPool{
				Id:    "testpool",
				Phase: apiv1.SharedPool_AVAILABLE,
			},
		}
		mockSharedPoolRepo.On("Load", mock.Anything, "testpool", false).Return(sharedPool, nil)
		mockLockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil)
		mockDagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil)
		mockScaleManager.On("AnyFullfillingDecision", mock.Anything, "testpool").Return(false, nil)

		server := &serverImpl{
			sharedPoolRepo: mockSharedPoolRepo,
			lockRepo:       mockLockRepo,
			dagExecutor:    mockDagExecutor,
			scaleManager:   mockScaleManager,
			transactioner:  &mockQuotaTransactioner{},
		}

		targetComponent := &entity.Component{
			Id:            "test-comp",
			ComponentType: apiv1.ComponentType_TIDB,
			NodeQuota: &apiv1.NodeQuota{
				MaxNodeCount: 2,
				MinNodeCount: 5,
			},
		}

		result, err := server.SetSharedPoolComponentQuota(context.Background(), "testpool", targetComponent)

		require.Error(t, err)
		assert.Contains(t, err.Error(), "target component min node count cannot be greater than max node count")
		assert.Nil(t, result)
	})
}

func TestSetSharedPoolComponentQuota_BusinessLogicValidation(t *testing.T) {
	t.Run("missing target component", func(t *testing.T) {
		mockSharedPoolRepo := repomocks.NewSharedPoolRepo(t)
		mockLockRepo := repomocks.NewLockRepo(t)
		mockDagExecutor := executormocks.NewRemoteExecutor(t)
		mockScaleManager := bractionmocks.NewScaleManager(t)

		targetComponent := &entity.Component{
			Id:            "test-comp",
			ComponentType: apiv1.ComponentType_COPROCESSOR_WORKER,
			NodeQuota: &apiv1.NodeQuota{
				MaxNodeCount: 10,
				MinNodeCount: 1,
			},
		}

		sharedPool := &entity.SharedPool{
			SharedPool: apiv1.SharedPool{
				Id:    "testpool",
				Phase: apiv1.SharedPool_AVAILABLE,
				ComponentGroups: []*entity.Component{
					{
						Id:            "test-comp",
						ComponentType: apiv1.ComponentType_TIKV_WORKER,
						Phase:         apiv1.Component_AVAILABLE,
						NodeQuota: &apiv1.NodeQuota{
							NodeCount:    3,
							MinNodeCount: 1,
							MaxNodeCount: 10,
						},
					},
				},
			},
		}

		mockSharedPoolRepo.On("Load", mock.Anything, "testpool", false).Return(sharedPool, nil)
		mockLockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil)
		mockDagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil)
		mockScaleManager.On("AnyFullfillingDecision", mock.Anything, mock.Anything).Return(false, nil)

		server := &serverImpl{
			sharedPoolRepo: mockSharedPoolRepo,
			lockRepo:       mockLockRepo,
			dagExecutor:    mockDagExecutor,
			transactioner:  &mockQuotaTransactioner{},
			scaleManager:   mockScaleManager,
		}

		result, err := server.SetSharedPoolComponentQuota(context.Background(), "testpool", targetComponent)

		require.Error(t, err)
		assert.Equal(t, codes.NotFound, status.Code(err))
		assert.Contains(t, err.Error(), "component coprocessor-tikv-worker test-comp not found in shared pool testpool")
		assert.Nil(t, result)
		mockDagExecutor.AssertNotCalled(t, "Submit", mock.Anything, mock.Anything, mock.Anything, mock.Anything)
		mockSharedPoolRepo.AssertExpectations(t)
		mockLockRepo.AssertExpectations(t)
		mockDagExecutor.AssertExpectations(t)
	})

	t.Run("current node count exceeds new max", func(t *testing.T) {
		// Create mocks
		mockSharedPoolRepo := repomocks.NewSharedPoolRepo(t)
		mockLockRepo := repomocks.NewLockRepo(t)
		mockDagExecutor := executormocks.NewRemoteExecutor(t)
		mockScaleManager := bractionmocks.NewScaleManager(t)

		// Test data - current node count is 5, new max is 2
		targetComponent := &entity.Component{
			Id:            "test-comp",
			ComponentType: apiv1.ComponentType_TIDB,
			NodeQuota: &apiv1.NodeQuota{
				MaxNodeCount: 2,
				MinNodeCount: 1,
			},
		}

		sharedPool := &entity.SharedPool{
			SharedPool: apiv1.SharedPool{
				Id:    "testpool",
				Phase: apiv1.SharedPool_AVAILABLE,
				ComponentGroups: []*entity.Component{
					{
						Id:            "test-comp",
						ComponentType: apiv1.ComponentType_TIDB,
						Phase:         apiv1.Component_AVAILABLE,
						NodeQuota: &apiv1.NodeQuota{
							NodeCount:    5, // Current node count is 5
							MinNodeCount: 1,
							MaxNodeCount: 10,
						},
					},
				},
			},
		}

		// Setup mocks
		mockSharedPoolRepo.On("Load", mock.Anything, "testpool", false).Return(sharedPool, nil)
		mockLockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil)
		mockDagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil)
		mockScaleManager.On("AnyFullfillingDecision", mock.Anything, mock.Anything).Return(false, nil)

		server := &serverImpl{
			sharedPoolRepo: mockSharedPoolRepo,
			lockRepo:       mockLockRepo,
			dagExecutor:    mockDagExecutor,
			transactioner:  &mockQuotaTransactioner{},
			scaleManager:   mockScaleManager,
		}

		// Execute test
		result, err := server.SetSharedPoolComponentQuota(context.Background(), "testpool", targetComponent)

		// Verify results
		require.Error(t, err)
		assert.Contains(t, err.Error(), "target component node count 5 is greater than new max node count 2")
		assert.Nil(t, result)

		// Verify expectations
		mockSharedPoolRepo.AssertExpectations(t)
		mockLockRepo.AssertExpectations(t)
		mockDagExecutor.AssertExpectations(t)
	})

	t.Run("current node count below new min", func(t *testing.T) {
		// Create mocks
		mockSharedPoolRepo := repomocks.NewSharedPoolRepo(t)
		mockLockRepo := repomocks.NewLockRepo(t)
		mockDagExecutor := executormocks.NewRemoteExecutor(t)
		mockScaleManager := bractionmocks.NewScaleManager(t)

		// Test data - current node count is 2, new min is 5
		targetComponent := &entity.Component{
			Id:            "test-comp",
			ComponentType: apiv1.ComponentType_TIDB,
			NodeQuota: &apiv1.NodeQuota{
				MaxNodeCount: 10,
				MinNodeCount: 5,
			},
		}

		sharedPool := &entity.SharedPool{
			SharedPool: apiv1.SharedPool{
				Id:    "testpool",
				Phase: apiv1.SharedPool_AVAILABLE,
				ComponentGroups: []*entity.Component{
					{
						Id:            "test-comp",
						ComponentType: apiv1.ComponentType_TIDB,
						Phase:         apiv1.Component_AVAILABLE,
						NodeQuota: &apiv1.NodeQuota{
							NodeCount:    2, // Current node count is 2
							MinNodeCount: 1,
							MaxNodeCount: 10,
						},
					},
				},
			},
		}

		// Setup mocks
		mockSharedPoolRepo.On("Load", mock.Anything, "testpool", false).Return(sharedPool, nil)
		mockLockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil)
		mockDagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil)
		mockScaleManager.On("AnyFullfillingDecision", mock.Anything, mock.Anything).Return(false, nil)
		mockDagExecutor.On("Submit", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)

		server := &serverImpl{
			sharedPoolRepo: mockSharedPoolRepo,
			lockRepo:       mockLockRepo,
			dagExecutor:    mockDagExecutor,
			transactioner:  &mockQuotaTransactioner{},
			scaleManager:   mockScaleManager,
		}

		// Execute test
		result, err := server.SetSharedPoolComponentQuota(context.Background(), "testpool", targetComponent)

		// Verify results
		require.NoError(t, err)
		assert.NotNil(t, result)
		assert.Equal(t, targetComponent.Id, result.Id)
		assert.Equal(t, targetComponent.ComponentType, result.ComponentType)
		assert.Equal(t, targetComponent.NodeQuota.MaxNodeCount, result.NodeQuota.MaxNodeCount)
		assert.Equal(t, targetComponent.NodeQuota.MinNodeCount, result.NodeQuota.MinNodeCount)

		// Verify expectations
		mockSharedPoolRepo.AssertExpectations(t)
		mockLockRepo.AssertExpectations(t)
		mockDagExecutor.AssertExpectations(t)
	})
}

func TestSetSharedPoolComponentQuota_GuardBlocksOnRunningTiKVRemoteWorkerDAG(t *testing.T) {
	sharedPoolRepo := repomocks.NewSharedPoolRepo(t)
	lockRepo := repomocks.NewLockRepo(t)
	dagExecutor := executormocks.NewRemoteExecutor(t)
	scaleManager := bractionmocks.NewScaleManager(t)

	sharedPool := newTKRWSharedPool(tkrwEnabledBoth(t), 3)
	sharedPoolRepo.On("Load", mock.Anything, testTKRWSharedPoolID, false).Return(sharedPool, nil)
	lockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil)
	setupAllDAGMocksWithConflict(dagExecutor, dag.PlanEnableCoprocessorWorker, runningDagStatus(tkrwKey(t, tkrwEnabledBoth(t))))
	scaleManager.On("AnyFullfillingDecision", mock.Anything, mock.Anything).Return(false, nil).Maybe()

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

	targetComponent := &entity.Component{
		Id:            testTKRWComponentID,
		ComponentType: apiv1.ComponentType_TIKV,
		NodeQuota: &apiv1.NodeQuota{
			MaxNodeCount: 9,
			MinNodeCount: 3,
		},
	}

	result, err := server.SetSharedPoolComponentQuota(context.Background(), testTKRWSharedPoolID, targetComponent)
	require.Error(t, err)
	assert.Equal(t, codes.FailedPrecondition, status.Code(err))
	assert.Nil(t, result)
	dagExecutor.AssertNotCalled(t, "Submit", mock.Anything, mock.Anything, mock.Anything, mock.Anything)
}

func TestSetSharedPoolComponentQuota_GuardBlocksOnUnconvergedTiKVRemoteWorkerTopology(t *testing.T) {
	sharedPoolRepo := repomocks.NewSharedPoolRepo(t)
	lockRepo := repomocks.NewLockRepo(t)
	dagExecutor := executormocks.NewRemoteExecutor(t)
	scaleManager := bractionmocks.NewScaleManager(t)

	sharedPool := newTKRWSharedPool(tkrwEnabledBoth(t), 3)
	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,
		transactioner:    &mockQuotaTransactioner{},
		scaleManager:     scaleManager,
		dataplaneManager: newTKRWDataplaneManager(t),
	}

	targetComponent := &entity.Component{
		Id:            testTKRWComponentID,
		ComponentType: apiv1.ComponentType_TIKV,
		NodeQuota: &apiv1.NodeQuota{
			MaxNodeCount: 9,
			MinNodeCount: 3,
		},
	}

	result, err := server.SetSharedPoolComponentQuota(context.Background(), testTKRWSharedPoolID, targetComponent)
	require.Error(t, err)
	assert.Equal(t, codes.FailedPrecondition, status.Code(err))
	assert.Nil(t, result)
	dagExecutor.AssertNotCalled(t, "Submit", mock.Anything, mock.Anything, mock.Anything, mock.Anything)
}