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"
	dplanemgrv1 "github.com/tidbcloud/cluster-service-ng/api/dataplane-manager/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"
	dataplanemocks "github.com/tidbcloud/cluster-service-ng/pkg/dataplanemgr/mocks"
	"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 mockNodeVersionTransactioner struct{}

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

func TestSetSharedPoolComponentNodeVersion_BasicValidation(t *testing.T) {
	testCases := []struct {
		name          string
		sharedPoolId  string
		compId        string
		compType      apiv1.ComponentType
		nodeVersion   string
		expectedError string
		expectedCode  codes.Code
	}{
		{
			name:          "empty shared pool ID",
			sharedPoolId:  "",
			compId:        "test-comp",
			compType:      apiv1.ComponentType_TIKV,
			nodeVersion:   "1.28",
			expectedError: "sharedPool id is required",
			expectedCode:  codes.InvalidArgument,
		},
		{
			name:          "unspecified component type",
			sharedPoolId:  "testpool",
			compId:        "test-comp",
			compType:      apiv1.ComponentType_COMPONENT_TYPE_UNSPECIFIED,
			nodeVersion:   "1.28",
			expectedError: "component type is required",
			expectedCode:  codes.InvalidArgument,
		},
		{
			name:          "empty node version",
			sharedPoolId:  "testpool",
			compId:        "test-comp",
			compType:      apiv1.ComponentType_TIKV,
			nodeVersion:   "",
			expectedError: "node_version is required",
			expectedCode:  codes.InvalidArgument,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			server := &serverImpl{
				transactioner: &mockNodeVersionTransactioner{},
			}

			result, err := server.SetSharedPoolComponentNodeVersion(context.Background(), tc.sharedPoolId, tc.compId, tc.compType, tc.nodeVersion)

			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 TestSetSharedPoolComponentNodeVersion_BlocksTiKVRemoteWorkerRoutingSwitch(t *testing.T) {
	for _, runningPlan := range []string{
		dag.PlanEnableCoprocessorWorker,
		dag.PlanDisableCoprocessorWorker,
	} {
		t.Run(runningPlan, func(t *testing.T) {
			sharedPool := newTKRWSharedPool(tkrwEnabledBoth(t), 3)
			sharedPoolRepo, lockRepo, dagExecutor, scaleManager := newTKRWHandlerMocksWithRunningDag(
				t, sharedPool, runningPlan, runningDagStatus(tkrwKey(t, sharedPool.GetEnabledTikvRemoteWorkers())))

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

			component, err := server.SetSharedPoolComponentNodeVersion(
				context.Background(), testTKRWSharedPoolID, testTKRWComponentID, apiv1.ComponentType_TIKV, "1.28")

			require.Error(t, err)
			assert.Nil(t, component)
			assert.Equal(t, codes.FailedPrecondition, status.Code(err))
			assert.Equal(t, apiv1.Component_AVAILABLE,
				sharedPoolDefaultComponentMap(sharedPool)[apiv1.ComponentType_TIKV].GetPhase())
			dagExecutor.AssertNotCalled(t, "Submit", mock.Anything, mock.Anything, mock.Anything, mock.Anything)
		})
	}
}

func TestSetSharedPoolComponentNodeVersion_BlocksTiKVRemoteWorkerTerminalDrift(t *testing.T) {
	sharedPool := newTKRWSharedPool(tkrwEnabledBoth(t), 3)
	sharedPoolRepo, lockRepo, dagExecutor, scaleManager := newTKRWHandlerMocks(t, sharedPool)

	server := &serverImpl{
		sharedPoolRepo:   sharedPoolRepo,
		lockRepo:         lockRepo,
		dagExecutor:      dagExecutor,
		transactioner:    &mockNodeVersionTransactioner{},
		scaleManager:     scaleManager,
		dataplaneManager: newTKRWDataplaneManager(t),
	}

	component, err := server.SetSharedPoolComponentNodeVersion(
		context.Background(), testTKRWSharedPoolID, testTKRWComponentID, apiv1.ComponentType_TIKV, "1.28")

	require.Error(t, err)
	assert.Nil(t, component)
	assert.Equal(t, codes.FailedPrecondition, status.Code(err))
	assert.Equal(t, apiv1.Component_AVAILABLE,
		sharedPoolDefaultComponentMap(sharedPool)[apiv1.ComponentType_TIKV].GetPhase())
	dagExecutor.AssertNotCalled(t, "Submit", mock.Anything, mock.Anything, mock.Anything, mock.Anything)
}

func TestSetSharedPoolComponentNodeVersion_SuccessfulUpdate(t *testing.T) {
	mockSharedPoolRepo := repomocks.NewSharedPoolRepo(t)
	mockLockRepo := repomocks.NewLockRepo(t)
	mockDagExecutor := executormocks.NewRemoteExecutor(t)
	mockScaleManager := bractionmocks.NewScaleManager(t)
	mockDataplaneClient := dataplanemocks.NewClient(t)

	sharedPool := &entity.SharedPool{
		SharedPool: apiv1.SharedPool{
			Id:            "testpool",
			DataplaneName: "test-dataplane",
			Phase:         apiv1.SharedPool_AVAILABLE,
			ComponentGroups: []*entity.Component{
				{
					Id:            "test-comp",
					ComponentType: apiv1.ComponentType_TIKV,
					Phase:         apiv1.Component_AVAILABLE,
					NodeQuota: &apiv1.NodeQuota{
						NodeSpec: &apiv1.NodeSpec{
							Cpu:      4,
							MemoryGi: 16,
						},
					},
				},
			},
		},
	}

	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)
	mockDataplaneClient.On("ListNodePools", mock.Anything, mock.Anything).Return(&dplanemgrv1.ListNodePoolsResponse{
		NodePools: []*dplanemgrv1.NodePool{
			{
				CpuCore:    4,
				MemoryGi:   16,
				K8SVersion: "1.28",
				NodeType:   dplanemgrv1.NodeType_Dedicated,
			},
		},
	}, nil)
	mockDagExecutor.On("Submit", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)

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

	ctx := context.Background()
	result, err := server.SetSharedPoolComponentNodeVersion(ctx, "testpool", "test-comp", apiv1.ComponentType_TIKV, "1.28")

	require.NoError(t, err)
	assert.NotNil(t, result)
	assert.Equal(t, "test-comp", result.Id)
	assert.Equal(t, apiv1.ComponentType_TIKV, result.ComponentType)

	mockSharedPoolRepo.AssertExpectations(t)
	mockLockRepo.AssertExpectations(t)
	mockDagExecutor.AssertExpectations(t)
}

func TestSetSharedPoolComponentNodeVersion_ComponentNotFound(t *testing.T) {
	mockSharedPoolRepo := repomocks.NewSharedPoolRepo(t)
	mockLockRepo := repomocks.NewLockRepo(t)
	mockDagExecutor := executormocks.NewRemoteExecutor(t)
	mockScaleManager := bractionmocks.NewScaleManager(t)

	sharedPool := &entity.SharedPool{
		SharedPool: apiv1.SharedPool{
			Id:    "testpool",
			Phase: apiv1.SharedPool_AVAILABLE,
			ComponentGroups: []*entity.Component{
				{
					Id:            "other-comp",
					ComponentType: apiv1.ComponentType_TIDB,
					Phase:         apiv1.Component_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, mock.Anything).Return(false, nil)

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

	ctx := context.Background()
	result, err := server.SetSharedPoolComponentNodeVersion(ctx, "testpool", "non-existent-comp", apiv1.ComponentType_TIKV, "1.28")

	require.Error(t, err)
	assert.Contains(t, err.Error(), "not found")
	statusErr, ok := status.FromError(err)
	require.True(t, ok)
	assert.Equal(t, codes.NotFound, statusErr.Code())
	assert.Nil(t, result)
}

func TestSetSharedPoolComponentNodeVersion_SharedPoolNotFound(t *testing.T) {
	mockSharedPoolRepo := repomocks.NewSharedPoolRepo(t)
	mockLockRepo := repomocks.NewLockRepo(t)

	mockSharedPoolRepo.On("Load", mock.Anything, "non-existent-pool", false).Return(nil, status.Errorf(codes.NotFound, "sharedPool not found"))
	mockLockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil)

	server := &serverImpl{
		sharedPoolRepo: mockSharedPoolRepo,
		lockRepo:       mockLockRepo,
		transactioner:  &mockNodeVersionTransactioner{},
	}

	ctx := context.Background()
	result, err := server.SetSharedPoolComponentNodeVersion(ctx, "non-existent-pool", "test-comp", apiv1.ComponentType_TIKV, "1.28")

	require.Error(t, err)
	assert.Contains(t, err.Error(), "not found")
	statusErr, ok := status.FromError(err)
	require.True(t, ok)
	assert.Equal(t, codes.NotFound, statusErr.Code())
	assert.Nil(t, result)
}

func TestSetSharedPoolComponentNodeVersion_SharedPoolCreating(t *testing.T) {
	mockSharedPoolRepo := repomocks.NewSharedPoolRepo(t)
	mockLockRepo := repomocks.NewLockRepo(t)
	mockDagExecutor := executormocks.NewRemoteExecutor(t)
	mockScaleManager := bractionmocks.NewScaleManager(t)

	sharedPool := &entity.SharedPool{
		SharedPool: apiv1.SharedPool{
			Id:    "testpool",
			Phase: apiv1.SharedPool_CREATING,
			ComponentGroups: []*entity.Component{
				{
					Id:            "test-comp",
					ComponentType: apiv1.ComponentType_TIKV,
					Phase:         apiv1.Component_CREATING,
				},
			},
		},
	}

	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:  &mockNodeVersionTransactioner{},
		scaleManager:   mockScaleManager,
	}

	ctx := context.Background()
	result, err := server.SetSharedPoolComponentNodeVersion(ctx, "testpool", "test-comp", apiv1.ComponentType_TIKV, "1.28")

	require.Error(t, err)
	assert.Contains(t, err.Error(), "creating")
	statusErr, ok := status.FromError(err)
	require.True(t, ok)
	assert.Equal(t, codes.FailedPrecondition, statusErr.Code())
	assert.Nil(t, result)
}

func TestSetSharedPoolComponentNodeVersion_UseDefaultComponentId(t *testing.T) {
	mockSharedPoolRepo := repomocks.NewSharedPoolRepo(t)
	mockLockRepo := repomocks.NewLockRepo(t)
	mockDagExecutor := executormocks.NewRemoteExecutor(t)
	mockScaleManager := bractionmocks.NewScaleManager(t)
	mockDataplaneClient := dataplanemocks.NewClient(t)

	sharedPool := &entity.SharedPool{
		SharedPool: apiv1.SharedPool{
			Id:            "testpool",
			DataplaneName: "test-dataplane",
			Phase:         apiv1.SharedPool_AVAILABLE,
			ComponentGroups: []*entity.Component{
				{
					Id:            "default-comp",
					ComponentType: apiv1.ComponentType_PD, // PD component is used to get default ID
					Phase:         apiv1.Component_AVAILABLE,
					NodeQuota: &apiv1.NodeQuota{
						NodeSpec: &apiv1.NodeSpec{
							Cpu:      4,
							MemoryGi: 16,
						},
					},
				},
				{
					Id:            "default-comp",
					ComponentType: apiv1.ComponentType_TIKV,
					Phase:         apiv1.Component_AVAILABLE,
					NodeQuota: &apiv1.NodeQuota{
						NodeSpec: &apiv1.NodeSpec{
							Cpu:      4,
							MemoryGi: 16,
						},
					},
				},
			},
		},
	}

	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)
	mockDataplaneClient.On("ListNodePools", mock.Anything, mock.Anything).Return(&dplanemgrv1.ListNodePoolsResponse{
		NodePools: []*dplanemgrv1.NodePool{
			{
				CpuCore:    4,
				MemoryGi:   16,
				K8SVersion: "1.28",
				NodeType:   dplanemgrv1.NodeType_Dedicated,
			},
		},
	}, nil)
	mockDagExecutor.On("Submit", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)

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

	ctx := context.Background()
	// Pass empty compId to use default
	result, err := server.SetSharedPoolComponentNodeVersion(ctx, "testpool", "", apiv1.ComponentType_TIKV, "1.28")

	require.NoError(t, err)
	assert.NotNil(t, result)
	assert.Equal(t, "default-comp", result.Id)

	mockSharedPoolRepo.AssertExpectations(t)
	mockLockRepo.AssertExpectations(t)
	mockDagExecutor.AssertExpectations(t)
}

func TestSetSharedPoolComponentNodeVersion_DifferentComponentTypes(t *testing.T) {
	componentTypes := []apiv1.ComponentType{
		apiv1.ComponentType_TIDB,
		apiv1.ComponentType_TIKV,
		apiv1.ComponentType_PD,
		apiv1.ComponentType_TIFLASH_WRITE,
		apiv1.ComponentType_TIFLASH_COMPUTE,
		apiv1.ComponentType_TIPROXY,
	}

	for _, compType := range componentTypes {
		t.Run(compType.String(), func(t *testing.T) {
			mockSharedPoolRepo := repomocks.NewSharedPoolRepo(t)
			mockLockRepo := repomocks.NewLockRepo(t)
			mockDagExecutor := executormocks.NewRemoteExecutor(t)
			mockScaleManager := bractionmocks.NewScaleManager(t)
			mockDataplaneClient := dataplanemocks.NewClient(t)

			sharedPool := &entity.SharedPool{
				SharedPool: apiv1.SharedPool{
					Id:            "testpool",
					DataplaneName: "test-dataplane",
					Phase:         apiv1.SharedPool_AVAILABLE,
					ComponentGroups: []*entity.Component{
						{
							Id:            "test-comp",
							ComponentType: compType,
							Phase:         apiv1.Component_AVAILABLE,
							NodeQuota: &apiv1.NodeQuota{
								NodeSpec: &apiv1.NodeSpec{
									Cpu:      4,
									MemoryGi: 16,
								},
							},
						},
					},
				},
			}

			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)
			mockDataplaneClient.On("ListNodePools", mock.Anything, mock.Anything).Return(&dplanemgrv1.ListNodePoolsResponse{
				NodePools: []*dplanemgrv1.NodePool{
					{
						CpuCore:    4,
						MemoryGi:   16,
						K8SVersion: "1.28",
						NodeType:   dplanemgrv1.NodeType_Dedicated,
					},
				},
			}, nil)
			mockDagExecutor.On("Submit", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)

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

			ctx := context.Background()
			result, err := server.SetSharedPoolComponentNodeVersion(ctx, "testpool", "test-comp", compType, "1.28")

			require.NoError(t, err)
			assert.NotNil(t, result)
			assert.Equal(t, compType, result.ComponentType)
		})
	}
}