package premium_dm

import (
	"context"
	"fmt"
	"testing"
	"time"

	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	center_mock "github.com/tidbcloud/dataflow-service-ng/app/event_bus/publisher/center/mock"
	"github.com/tidbcloud/dataflow-service-ng/app/repo"
	"github.com/tidbcloud/dataflow-service-ng/app/repo/entity"
	"github.com/tidbcloud/dataflow-service-ng/pkg/agent"
	grpclimock "github.com/tidbcloud/dataflow-service-ng/pkg/grpc/grpcli/mock"
	"github.com/tidbcloud/eventbus-datastruct/event"
	"github.com/tidbcloud/pkgv2/db"
	tidbmgmtservicev1 "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
	"gorm.io/gorm"
)

type mockDMClusterRegionalClient struct {
	mock.Mock
}

func (m *mockDMClusterRegionalClient) ListDMClusters(
	ctx context.Context, sharedPoolID uint64,
) (*agent.ListDMClustersResponse, error) {
	args := m.Called(ctx, sharedPoolID)
	if resp := args.Get(0); resp != nil {
		return resp.(*agent.ListDMClustersResponse), args.Error(1)
	}
	return nil, args.Error(1)
}

func (m *mockDMClusterRegionalClient) GetDMClusterByID(
	ctx context.Context, sharedPoolID uint64, dmClusterID string,
) (*agent.DMCluster, error) {
	args := m.Called(ctx, sharedPoolID, dmClusterID)
	if dmCluster := args.Get(0); dmCluster != nil {
		return dmCluster.(*agent.DMCluster), args.Error(1)
	}
	return nil, args.Error(1)
}

func (m *mockDMClusterRegionalClient) CreateDMCluster(
	ctx context.Context, sharedPoolID uint64, req *agent.CreateDMClusterRequest,
) (*agent.DMCluster, error) {
	args := m.Called(ctx, sharedPoolID, req)
	if dmCluster := args.Get(0); dmCluster != nil {
		return dmCluster.(*agent.DMCluster), args.Error(1)
	}
	return nil, args.Error(1)
}

func (m *mockDMClusterRegionalClient) ScaleDMCluster(
	ctx context.Context, sharedPoolID uint64, dmClusterID string, req *agent.ScaleDMClusterRequest,
) error {
	return m.Called(ctx, sharedPoolID, dmClusterID, req).Error(0)
}

func (m *mockDMClusterRegionalClient) DeleteDMCluster(ctx context.Context, sharedPoolID uint64, dmClusterID string) error {
	return m.Called(ctx, sharedPoolID, dmClusterID).Error(0)
}

func TestOnTaskDeletingDeletesClusterCustomStatusWhenPresent(t *testing.T) {
	ctx := context.Background()
	testDB, err := repo.GetFakeDB()
	require.NoError(t, err)

	dao := NewPremiumDMTaskDAO(db.NewProvider(testDB))
	mgmtCli := grpclimock.NewMgmtSvcGrpcClient(t)
	publisher := center_mock.NewEventCenterPublisher(t)
	worker := NewPremiumDMTaskWorker(dao, mgmtCli, nil, agent.ClientConfig{}, nil, nil, nil, nil, publisher)
	worker.newProxyFunc = func(ctx context.Context, tiDBID string, dmCluster *agent.DMCluster) (*PremiumDMProxy, error) {
		return nil, nil
	}

	task := buildDeletingPremiumDMTask(6789)
	require.NoError(t, db.NewProvider(testDB).WriteDB(ctx).Create(task).Error)
	defer cleanupDeletingPremiumDMTask(t, testDB, task.ID)
	publisher.On("PublishAsync", mock.MatchedBy(func(ev *event.DisplayEvent) bool {
		return ev.ContextID == task.ID && ev.DisplayStatus == event.EventStatusSucceeded
	})).Once()

	var deleteReq *tidbmgmtservicev1.DeleteClusterCustomStatusRequest
	deleteCall := mgmtCli.On("DeleteClusterCustomStatus", mock.Anything, mock.Anything).
		Run(func(args mock.Arguments) {
			deleteReq = args.Get(1).(*tidbmgmtservicev1.DeleteClusterCustomStatusRequest)
		}).
		Return(nil)
	defer deleteCall.Unset()

	err = worker.OnTaskDeleting(ctx, task, nil)
	require.NoError(t, err)
	require.NotNil(t, deleteReq)
	require.Equal(t, uint64(10001), deleteReq.GetClusterId())
	require.Equal(t, uint64(6789), deleteReq.GetStatusId())

	stored, bizErr := dao.GetTask(ctx, task.TiDBID, task.ID, nil)
	require.NoError(t, bizErr)
	require.Equal(t, entity.PremiumDMTaskStageDeleted, stored.Stage)
	require.Empty(t, stored.TargetEncryptedPassword)
}

func TestOnTaskDeletingRetryableFailureDoesNotPublishEvent(t *testing.T) {
	ctx := context.Background()
	publisher := center_mock.NewEventCenterPublisher(t)
	worker := NewPremiumDMTaskWorker(
		nil, nil, nil, agent.ClientConfig{}, nil, nil, nil, nil, publisher,
	)
	worker.newProxyFunc = func(
		context.Context, string, *agent.DMCluster,
	) (*PremiumDMProxy, error) {
		return nil, fmt.Errorf("temporary proxy failure")
	}

	err := worker.OnTaskDeleting(ctx, buildDeletingPremiumDMTask(0), nil)
	require.EqualError(t, err, "temporary proxy failure")
	publisher.AssertNotCalled(t, "PublishAsync", mock.Anything)
}

func TestOnTaskDeletingSkipsClusterCustomStatusDeletionWhenMissing(t *testing.T) {
	ctx := context.Background()
	testDB, err := repo.GetFakeDB()
	require.NoError(t, err)

	dao := NewPremiumDMTaskDAO(db.NewProvider(testDB))
	mgmtCli := grpclimock.NewMgmtSvcGrpcClient(t)
	worker := NewPremiumDMTaskWorker(dao, mgmtCli, nil, agent.ClientConfig{}, nil, nil, nil, nil, nil)
	worker.newProxyFunc = func(ctx context.Context, tiDBID string, dmCluster *agent.DMCluster) (*PremiumDMProxy, error) {
		return nil, nil
	}

	task := buildDeletingPremiumDMTask(0)
	require.NoError(t, db.NewProvider(testDB).WriteDB(ctx).Create(task).Error)
	defer cleanupDeletingPremiumDMTask(t, testDB, task.ID)

	err = worker.OnTaskDeleting(ctx, task, nil)
	require.NoError(t, err)

	stored, bizErr := dao.GetTask(ctx, task.TiDBID, task.ID, nil)
	require.NoError(t, bizErr)
	require.Equal(t, entity.PremiumDMTaskStageDeleted, stored.Stage)
	require.Empty(t, stored.TargetEncryptedPassword)
}

func TestDeleteDMClusterIfNoWorkers(t *testing.T) {
	tests := map[string]struct {
		dmCluster        *agent.DMCluster
		getErr           error
		deleteErr        error
		expectedErr      string
		expectDeleteCall bool
	}{
		"success": {
			dmCluster: &agent.DMCluster{
				WorkerNodeQuota: &agent.NodeQuota{NodeCount: 0},
			},
			expectDeleteCall: true,
		},
		"reject non-zero worker quota": {
			dmCluster: &agent.DMCluster{
				WorkerNodeQuota: &agent.NodeQuota{NodeCount: 2},
			},
			expectedErr: "delete premium dm cluster is not allowed when worker node count is 2",
		},
		"get cluster failure": {
			getErr:      fmt.Errorf("get failed"),
			expectedErr: "get failed",
		},
		"delete failure": {
			dmCluster: &agent.DMCluster{
				WorkerNodeQuota: &agent.NodeQuota{NodeCount: 0},
			},
			deleteErr:        fmt.Errorf("delete failed"),
			expectedErr:      "delete failed",
			expectDeleteCall: true,
		},
	}

	for name, tc := range tests {
		t.Run(name, func(t *testing.T) {
			const (
				sharedPoolID = uint64(101)
				dmClusterID  = "202"
			)
			ctx := context.Background()
			mgmtCli := grpclimock.NewMgmtSvcGrpcClient(t)
			regionalClient := &mockDMClusterRegionalClient{}
			worker := NewPremiumDMTaskWorker(nil, mgmtCli, nil, agent.ClientConfig{}, nil, nil, nil, nil, nil)

			var constructedResourcePoolID uint64
			worker.newDMClusterRegionalClientByResourcePoolIDFunc = func(
				ctx context.Context, resourcePoolID uint64,
			) (dmClusterRegionalClient, error) {
				constructedResourcePoolID = resourcePoolID
				return regionalClient, nil
			}

			getCall := regionalClient.On("GetDMClusterByID", ctx, sharedPoolID, "202").
				Return(tc.dmCluster, tc.getErr)
			if tc.expectDeleteCall {
				regionalClient.On("DeleteDMCluster", ctx, sharedPoolID, "202").Return(tc.deleteErr).Once()
			}

			err := worker.DeleteDMClusterIfNoWorkers(ctx, sharedPoolID, dmClusterID)
			if tc.expectedErr != "" {
				require.EqualError(t, err, tc.expectedErr)
			} else {
				require.NoError(t, err)
			}
			require.Equal(t, sharedPoolID, constructedResourcePoolID)

			getCall.Unset()
			regionalClient.AssertExpectations(t)
		})
	}
}

func buildDeletingPremiumDMTask(clusterStatusID uint64) *entity.PremiumDMTask {
	now := time.Now().UnixNano()
	return &entity.PremiumDMTask{
		ID:                      fmt.Sprintf("dmtsk-del-%d", now),
		Name:                    fmt.Sprintf("del-%d", now),
		Stage:                   entity.PremiumDMTaskStageDeleting,
		TiDBID:                  "10001",
		TenantID:                "20002",
		CloudProvider:           string(entity.CloudProviderTypeAWS),
		Region:                  "us-west-2",
		Mode:                    entity.PremiumDMTaskModeFull,
		ImportMode:              entity.PremiumDMTaskImportModeLogical,
		TargetUser:              "root",
		TargetEncryptedPassword: "encrypted-password",
		ClusterStatusID:         clusterStatusID,
	}
}

func cleanupDeletingPremiumDMTask(t *testing.T, testDB *gorm.DB, taskID string) {
	t.Helper()
	require.NoError(t, testDB.Exec("DELETE FROM premium_dm_sources WHERE dm_task_id = ?", taskID).Error)
	require.NoError(t, testDB.Exec("DELETE FROM premium_dm_tasks WHERE id = ?", taskID).Error)
}