package cluster_job

import (
	"context"
	"testing"

	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	appcontainer "github.com/tidbcloud/dataflow-service/app/container"
	grpclimock "github.com/tidbcloud/dataflow-service/pkg/grpc/grpcli/mock"
	"github.com/tidbcloud/pkgv2/bizerr"
	mgmt "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
)

func TestRegisterPauseLock(t *testing.T) {
	ctx := context.Background()
	clusterID := uint64(9527)
	taskID := uint64(20260703)

	mockMgmt := grpclimock.NewMgmtSvcGrpcClient(t)
	oldMgmt := appcontainer.Global().MgmtSvcGrpcClient
	appcontainer.Global().MgmtSvcGrpcClient = mockMgmt
	defer func() {
		appcontainer.Global().MgmtSvcGrpcClient = oldMgmt
	}()

	tests := map[string]struct {
		registerFn         func(context.Context, uint64, uint64) bizerr.BizError
		expectedStatusType mgmt.CustomStatusType
		expectedStatusName string
	}{

		"dm pause lock": {
			registerFn:         RegisterDMPauseLock,
			expectedStatusType: mgmt.CustomStatusType_CUSTOM_STATUS_TYPE_DATA_MIGRATION,
			expectedStatusName: dmPauseLockStatusName(taskID),
		},
	}

	for name, tc := range tests {
		t.Run(name, func(t *testing.T) {
			call := mockMgmt.On("SetClusterCustomStatus", mock.Anything, mock.Anything).
				Run(func(args mock.Arguments) {
					req, ok := args.Get(1).(*mgmt.SetClusterCustomStatusRequest)
					require.True(t, ok)
					require.NotNil(t, req.CustomStatus)
					require.Equal(t, clusterID, req.CustomStatus.ClusterId)
					require.Equal(t, tc.expectedStatusType, req.CustomStatus.StatusType)
					require.Equal(t, tc.expectedStatusName, req.CustomStatus.StatusName)
					require.NotNil(t, req.CustomStatus.Mutex)
					require.Equal(t, []mgmt.ClusterAction{mgmt.ClusterAction_CLUSTER_ACTION_PAUSE_CLUSTER},
						req.CustomStatus.Mutex.GetClusterAction())
				}).
				Return(&mgmt.ClusterCustomStatus{StatusId: 1}, nil)
			defer call.Unset()

			err := tc.registerFn(ctx, clusterID, taskID)
			require.NoError(t, err)
		})
	}
}

func TestIsDMPauseLockRegistered(t *testing.T) {
	ctx := context.Background()
	clusterID := uint64(9527)
	dmJobID := uint64(20260703)

	mockMgmt := grpclimock.NewMgmtSvcGrpcClient(t)
	oldMgmt := appcontainer.Global().MgmtSvcGrpcClient
	appcontainer.Global().MgmtSvcGrpcClient = mockMgmt
	defer func() {
		appcontainer.Global().MgmtSvcGrpcClient = oldMgmt
	}()

	tests := map[string]struct {
		response *mgmt.IsClusterCustomStatusExistedResponse
		expected bool
	}{
		"lock exists": {
			response: &mgmt.IsClusterCustomStatusExistedResponse{
				IsExisted: true,
				ExistedCustomStatuses: []*mgmt.ClusterCustomStatus{
					{StatusName: dmPauseLockStatusName(dmJobID)},
				},
			},
			expected: true,
		},
		"lock not found": {
			response: &mgmt.IsClusterCustomStatusExistedResponse{
				IsExisted: true,
				ExistedCustomStatuses: []*mgmt.ClusterCustomStatus{
					{StatusName: "other"},
				},
			},
			expected: false,
		},
		"status type absent": {
			response: &mgmt.IsClusterCustomStatusExistedResponse{
				IsExisted: false,
			},
			expected: false,
		},
	}

	for name, tc := range tests {
		t.Run(name, func(t *testing.T) {
			queryCall := mockMgmt.On("IsClusterCustomStatusExisted", mock.Anything, mock.Anything).
				Run(func(args mock.Arguments) {
					req, ok := args.Get(1).(*mgmt.IsClusterCustomStatusExistedRequest)
					require.True(t, ok)
					require.Equal(t, clusterID, req.ClusterId)
					require.Equal(t, mgmt.CustomStatusType_CUSTOM_STATUS_TYPE_DATA_MIGRATION, req.StatusType)
				}).
				Return(tc.response, nil)
			defer queryCall.Unset()

			existed, err := IsDMPauseLockRegistered(ctx, clusterID, dmJobID)
			require.NoError(t, err)
			require.Equal(t, tc.expected, existed)
		})
	}
}