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/config"
	"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"
	clustergatewaymocks "github.com/tidbcloud/cluster-service-ng/pkg/clustergateway/mocks"
	dataplanemocks "github.com/tidbcloud/cluster-service-ng/pkg/dataplane/mocks"
	dataplanemgrmocks "github.com/tidbcloud/cluster-service-ng/pkg/dataplanemgr/mocks"
	"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"
)

func TestScaleSharedPoolBatch(t *testing.T) {
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	lockRepo := repomnocks.NewLockRepo(t)
	dataplaneClient := dataplanemgrmocks.NewClient(t)
	dagExecutor := executormocks.NewRemoteExecutor(t)
	scaleManager := bractionmocks.NewScaleManager(t)
	dataplaneManager := dataplanemocks.NewManager(t)
	ruConfigRepo := repomnocks.NewRuConfigRepo(t)
	tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t)
	clusterGatewayMock := clustergatewaymocks.NewClient(t)

	sharedPool := createTestSharedPool()
	sharedPool.ComponentGroups = append(sharedPool.ComponentGroups, &entity.Component{
		Id:            "cluster-a",
		ComponentType: apiv1.ComponentType_TIPROXY,
		Phase:         apiv1.Component_AVAILABLE,
		NodeQuota: &entity.NodeQuota{
			NodeCount:    2,
			MinNodeCount: 1,
			MaxNodeCount: 4,
			NodeSpec: &entity.NodeSpec{
				Cpu:      8,
				MemoryGi: 16,
			},
		},
	})

	sharedPoolRepo.On("Load", mock.Anything, "testsharedpoolid", false).Return(sharedPool, nil)
	lockRepo.On("Lock", mock.Anything, mock.MatchedBy(func(lock *entity.Lock) bool {
		return lock.Key == "testsharedpoolid"
	})).Return(func() error { return nil }, nil)
	dataplaneClient.On("ListNodePools", mock.Anything, mock.Anything).Return(&dplanemgrv1.ListNodePoolsResponse{
		NodePools: []*dplanemgrv1.NodePool{{
			CpuCore:               8,
			MemoryGi:              16,
			AvailableCpuMilliCore: 7800,
			AvailableMemoryMi:     15000,
			NodeType:              dplanemgrv1.NodeType_Dedicated,
			Version:               "v1.0.0",
		}},
	}, nil)
	dagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil)
	scaleManager.On("AnyFullfillingDecision", mock.Anything, mock.Anything).Return(false, nil)
	dagExecutor.On("Submit", mock.Anything, dag.PlanUpdateComponent, mock.MatchedBy(func(args actions.ActionArgs) bool {
		componentIDs, _ := args.GetString(workflow.ComponentIdsKey)
		componentTypes, _ := args.GetString(workflow.ComponentTypesKey)
		tidbOps, _ := args.GetString(entity.GetUpdateComponentKey(&entity.Component{Id: "default", ComponentType: apiv1.ComponentType_TIDB}))
		tiproxyOps, _ := args.GetString(entity.GetUpdateComponentKey(&entity.Component{Id: "cluster-a", ComponentType: apiv1.ComponentType_TIPROXY}))

		return componentIDs == "default,cluster-a" &&
			componentTypes == "tidb,tiproxy" &&
			tidbOps == "update-node-count" &&
			tiproxyOps == "update-node-count"
	}), mock.Anything).Return(nil)

	server := &serverImpl{
		cloudProvider:    "aws",
		region:           "us-west-2",
		sharedPoolRepo:   sharedPoolRepo,
		lockRepo:         lockRepo,
		dataplaneManager: dataplaneManager,
		dataplaneClient:  dataplaneClient,
		dagExecutor:      dagExecutor,
		transactioner:    &mockTransactioner{},
		bizCfg: &config.BizConfig{
			DefaultDataplaneAccount: "test-account",
		},
		scaleManager:      scaleManager,
		ruConfigRepo:      ruConfigRepo,
		tidbClusterRepo:   tidbClusterRepo,
		clustergatewayCli: clusterGatewayMock,
	}

	result, err := server.ScaleSharedPoolBatch(context.Background(), "testsharedpoolid", []*entity.Component{
		{
			Id:            "cluster-a",
			ComponentType: apiv1.ComponentType_TIPROXY,
			NodeQuota: &entity.NodeQuota{
				NodeCount: 3,
			},
		},
		{
			Id:            "default",
			ComponentType: apiv1.ComponentType_TIDB,
			NodeQuota: &entity.NodeQuota{
				NodeCount: 3,
			},
		},
	}, false)
	require.NoError(t, err)
	require.NotNil(t, result)
	assert.Equal(t, apiv1.SharedPool_MODIFYING, result.GetPhase())
}

func TestScaleSharedPoolBatchForceBypassesTiKVRemoteWorkerGuard(t *testing.T) {
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	lockRepo := repomnocks.NewLockRepo(t)
	dataplaneClient := dataplanemgrmocks.NewClient(t)
	dagExecutor := executormocks.NewRemoteExecutor(t)
	scaleManager := bractionmocks.NewScaleManager(t)
	dataplaneManager := newTKRWDataplaneManager(t)
	ruConfigRepo := repomnocks.NewRuConfigRepo(t)
	tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t)
	clusterGatewayMock := clustergatewaymocks.NewClient(t)

	sharedPool := newTKRWSharedPool(tkrwEnabledBoth(t), 3)
	setupTKRWScaleSuccessMocks(sharedPool, dataplaneClient, dagExecutor)

	sharedPoolRepo.On("Load", mock.Anything, testTKRWSharedPoolID, false).Return(sharedPool, nil)
	lockRepo.On("Lock", mock.Anything, mock.MatchedBy(func(lock *entity.Lock) bool {
		return lock.Key == testTKRWSharedPoolID
	})).Return(func() error { return nil }, nil)
	scaleManager.On("AnyFullfillingDecision", mock.Anything, mock.Anything).Return(false, nil)
	setupAllDAGMocksWithConflict(dagExecutor, dag.PlanEnableCoprocessorWorker, runningDagStatus(tkrwKey(t, tkrwEnabledBoth(t))))

	server := &serverImpl{
		cloudProvider:    "aws",
		region:           "us-west-2",
		sharedPoolRepo:   sharedPoolRepo,
		lockRepo:         lockRepo,
		dataplaneManager: dataplaneManager,
		dataplaneClient:  dataplaneClient,
		dagExecutor:      dagExecutor,
		transactioner:    &mockTransactioner{},
		bizCfg: &config.BizConfig{
			DefaultDataplaneAccount: "test-account",
		},
		scaleManager:      scaleManager,
		ruConfigRepo:      ruConfigRepo,
		tidbClusterRepo:   tidbClusterRepo,
		clustergatewayCli: clusterGatewayMock,
	}

	scaleComponents := []*entity.Component{
		{
			Id:            testTKRWComponentID,
			ComponentType: apiv1.ComponentType_TIKV,
			NodeQuota: &entity.NodeQuota{
				NodeCount: 6,
			},
		},
	}

	result, err := server.ScaleSharedPoolBatch(context.Background(), testTKRWSharedPoolID, scaleComponents, true)
	require.NoError(t, err)
	require.NotNil(t, result)
	assert.Equal(t, apiv1.SharedPool_MODIFYING, result.GetPhase())
}

func TestScaleSharedPoolBatchTiKVRemoteWorkerTerminalDriftGuard(t *testing.T) {
	tests := []struct {
		name        string
		force       bool
		wantSuccess bool
	}{
		{name: "non-force request is blocked", force: false, wantSuccess: false},
		{name: "force request bypasses guard", force: true, wantSuccess: true},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
			lockRepo := repomnocks.NewLockRepo(t)
			dataplaneClient := dataplanemgrmocks.NewClient(t)
			dagExecutor := executormocks.NewRemoteExecutor(t)
			scaleManager := bractionmocks.NewScaleManager(t)
			dataplaneManager := newTKRWDataplaneManager(t)
			ruConfigRepo := repomnocks.NewRuConfigRepo(t)
			tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t)
			clusterGatewayMock := clustergatewaymocks.NewClient(t)

			sharedPool := newTKRWSharedPool(tkrwEnabledBoth(t), 3)
			if tt.wantSuccess {
				setupTKRWScaleSuccessMocks(sharedPool, dataplaneClient, dagExecutor)
			}

			sharedPoolRepo.On("Load", mock.Anything, testTKRWSharedPoolID, false).Return(sharedPool, nil)
			lockRepo.On("Lock", mock.Anything, mock.MatchedBy(func(lock *entity.Lock) bool {
				return lock.Key == testTKRWSharedPoolID
			})).Return(func() error { return nil }, nil)
			scaleManager.On("AnyFullfillingDecision", mock.Anything, mock.Anything).Return(false, nil)
			setupAllDAGMocks(dagExecutor)

			server := &serverImpl{
				cloudProvider:    "aws",
				region:           "us-west-2",
				sharedPoolRepo:   sharedPoolRepo,
				lockRepo:         lockRepo,
				dataplaneManager: dataplaneManager,
				dataplaneClient:  dataplaneClient,
				dagExecutor:      dagExecutor,
				transactioner:    &mockTransactioner{},
				bizCfg: &config.BizConfig{
					DefaultDataplaneAccount: "test-account",
				},
				scaleManager:      scaleManager,
				ruConfigRepo:      ruConfigRepo,
				tidbClusterRepo:   tidbClusterRepo,
				clustergatewayCli: clusterGatewayMock,
			}

			scaleComponents := []*entity.Component{
				{
					Id:            testTKRWComponentID,
					ComponentType: apiv1.ComponentType_TIKV,
					NodeQuota: &entity.NodeQuota{
						NodeCount: 6,
					},
				},
			}

			result, err := server.ScaleSharedPoolBatch(context.Background(), testTKRWSharedPoolID, scaleComponents, tt.force)
			if tt.wantSuccess {
				require.NoError(t, err)
				require.NotNil(t, result)
				assert.Equal(t, apiv1.SharedPool_MODIFYING, result.GetPhase())
				return
			}

			require.Error(t, err)
			assert.Equal(t, codes.FailedPrecondition, status.Code(err))
			require.ErrorContains(t, err, "tikv remote worker topology has not converged")
			assert.Nil(t, result)
			dagExecutor.AssertNotCalled(t, "Submit", mock.Anything, mock.Anything, mock.Anything, mock.Anything)
		})
	}
}

func TestCanonicalizeComponentRefs(t *testing.T) {
	componentIDs, componentTypes := canonicalizeComponentRefs(
		[]string{"cluster-b", "default", "cluster-a"},
		[]string{"tiproxy", "tidb", "tiproxy"},
	)

	assert.Equal(t, []string{"default", "cluster-a", "cluster-b"}, componentIDs)
	assert.Equal(t, []string{"tidb", "tiproxy", "tiproxy"}, componentTypes)
}