package premium_dm

import (
	"context"
	"testing"

	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	center_mock "github.com/tidbcloud/dataflow-service-ng/app/event_bus/publisher/center/mock"
	modelpremiumdm "github.com/tidbcloud/dataflow-service-ng/app/models/premium_dm"
	"github.com/tidbcloud/dataflow-service-ng/app/repo"
	"github.com/tidbcloud/dataflow-service-ng/app/repo/entity"
	"github.com/tidbcloud/dataflow-service-ng/idl/pbgen/proto/premium_dm/v1beta2"
	"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"
	mgmt "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
)

func TestPausePremiumMigrationAcceptsAsynchronously(t *testing.T) {
	ctx := context.Background()
	testDB, err := repo.GetFakeDB()
	require.NoError(t, err)
	dao := modelpremiumdm.NewPremiumDMTaskDAO(db.NewProvider(testDB))
	mgmtCli := grpclimock.NewMgmtSvcGrpcClient(t)
	publisher := center_mock.NewEventCenterPublisher(t)
	service := NewPremiumDMService(
		dao, nil, mgmtCli, nil, agent.ClientConfig{}, db.NewProvider(testDB), nil, nil, publisher,
	)
	task := newAsyncControllerTask(entity.PremiumDMTaskStageRunning)
	require.NoError(t, db.NewProvider(testDB).WriteDB(ctx).Create(task).Error)
	defer cleanupPremiumDMTask(t, testDB, task.ID)

	mgmtCli.On("ListClusters", mock.Anything, mock.Anything).Return(&mgmt.ListClustersResponse{
		Clusters: []*mgmt.Cluster{buildPremiumDMCluster()},
	}, nil).Once()
	publisher.On("PublishAsync", mock.MatchedBy(func(ev *event.DisplayEvent) bool {
		return ev.ContextID == task.ID &&
			ev.DisplayEventType == string(entity.PauseDataMigration) &&
			ev.DisplayStatus == event.EventStatusStarted
	})).Once()

	resp, err := service.PausePremiumMigration(ctx, &v1beta2.PausePremiumMigrationReq{
		TidbId: task.TiDBID, MigrationId: task.ID,
	})
	require.NoError(t, err)
	require.NotNil(t, resp)
	stored, bizErr := dao.GetTask(ctx, task.TiDBID, task.ID, nil)
	require.NoError(t, bizErr)
	require.Equal(t, entity.PremiumDMTaskStagePausing, stored.Stage)
	require.Nil(t, stored.OperationDispatchedAt)
	publisher.AssertExpectations(t)
}

func TestResumePremiumMigrationAcceptsAsynchronously(t *testing.T) {
	for _, initialStage := range []entity.PremiumDMTaskStage{
		entity.PremiumDMTaskStageStopped,
		entity.PremiumDMTaskStageFailed,
	} {
		t.Run(string(initialStage), func(t *testing.T) {
			ctx := context.Background()
			testDB, err := repo.GetFakeDB()
			require.NoError(t, err)
			dao := modelpremiumdm.NewPremiumDMTaskDAO(db.NewProvider(testDB))
			mgmtCli := grpclimock.NewMgmtSvcGrpcClient(t)
			publisher := center_mock.NewEventCenterPublisher(t)
			service := NewPremiumDMService(
				dao, nil, mgmtCli, nil, agent.ClientConfig{}, db.NewProvider(testDB), nil, nil, publisher,
			)
			task := newAsyncControllerTask(initialStage)
			require.NoError(t, db.NewProvider(testDB).WriteDB(ctx).Create(task).Error)
			defer cleanupPremiumDMTask(t, testDB, task.ID)

			mgmtCli.On("ListClusters", mock.Anything, mock.Anything).Return(&mgmt.ListClustersResponse{
				Clusters: []*mgmt.Cluster{buildPremiumDMCluster()},
			}, nil).Once()
			publisher.On("PublishAsync", mock.MatchedBy(func(ev *event.DisplayEvent) bool {
				return ev.ContextID == task.ID &&
					ev.DisplayEventType == string(entity.ResumeDataMigration) &&
					ev.DisplayStatus == event.EventStatusStarted
			})).Once()

			resp, err := service.ResumePremiumMigration(ctx, &v1beta2.ResumePremiumMigrationReq{
				TidbId: task.TiDBID, MigrationId: task.ID,
			})
			require.NoError(t, err)
			require.NotNil(t, resp)
			stored, bizErr := dao.GetTask(ctx, task.TiDBID, task.ID, nil)
			require.NoError(t, bizErr)
			require.Equal(t, entity.PremiumDMTaskStageResuming, stored.Stage)
			require.Nil(t, stored.OperationDispatchedAt)
			publisher.AssertExpectations(t)
		})
	}
}

func TestPausePremiumMigrationRejectsDuplicateTransition(t *testing.T) {
	ctx := context.Background()
	testDB, err := repo.GetFakeDB()
	require.NoError(t, err)
	dao := modelpremiumdm.NewPremiumDMTaskDAO(db.NewProvider(testDB))
	mgmtCli := grpclimock.NewMgmtSvcGrpcClient(t)
	service := NewPremiumDMService(
		dao, nil, mgmtCli, nil, agent.ClientConfig{}, db.NewProvider(testDB), nil, nil, nil,
	)
	task := newAsyncControllerTask(entity.PremiumDMTaskStagePausing)
	require.NoError(t, db.NewProvider(testDB).WriteDB(ctx).Create(task).Error)
	defer cleanupPremiumDMTask(t, testDB, task.ID)
	mgmtCli.On("ListClusters", mock.Anything, mock.Anything).Return(&mgmt.ListClustersResponse{
		Clusters: []*mgmt.Cluster{buildPremiumDMCluster()},
	}, nil).Once()

	resp, err := service.PausePremiumMigration(ctx, &v1beta2.PausePremiumMigrationReq{
		TidbId: task.TiDBID, MigrationId: task.ID,
	})
	require.Nil(t, resp)
	require.ErrorContains(t, err, "is not running")
}

func TestDeletePremiumMigrationCancelsSupersededTransition(t *testing.T) {
	for _, initialStage := range []entity.PremiumDMTaskStage{
		entity.PremiumDMTaskStagePausing,
		entity.PremiumDMTaskStageResuming,
	} {
		t.Run(string(initialStage), func(t *testing.T) {
			ctx := context.Background()
			testDB, err := repo.GetFakeDB()
			require.NoError(t, err)
			dao := modelpremiumdm.NewPremiumDMTaskDAO(db.NewProvider(testDB))
			mgmtCli := grpclimock.NewMgmtSvcGrpcClient(t)
			publisher := center_mock.NewEventCenterPublisher(t)
			service := NewPremiumDMService(
				dao, nil, mgmtCli, nil, agent.ClientConfig{}, db.NewProvider(testDB), nil, nil, publisher,
			)
			task := newAsyncControllerTask(initialStage)
			require.NoError(t, db.NewProvider(testDB).WriteDB(ctx).Create(task).Error)
			defer cleanupPremiumDMTask(t, testDB, task.ID)

			mgmtCli.On("ListClusters", mock.Anything, mock.Anything).Return(&mgmt.ListClustersResponse{
				Clusters: []*mgmt.Cluster{buildPremiumDMCluster()},
			}, nil).Once()
			operation := entity.PauseDataMigration
			if initialStage == entity.PremiumDMTaskStageResuming {
				operation = entity.ResumeDataMigration
			}
			publisher.On("PublishAsync", mock.MatchedBy(func(ev *event.DisplayEvent) bool {
				return ev.ContextID == task.ID &&
					ev.DisplayEventType == string(operation) &&
					ev.DisplayStatus == event.EventStatusCanceled
			})).Once()
			publisher.On("PublishAsync", mock.MatchedBy(func(ev *event.DisplayEvent) bool {
				return ev.ContextID == task.ID &&
					ev.DisplayEventType == string(entity.DeleteDataMigration) &&
					ev.DisplayStatus == event.EventStatusStarted
			})).Once()

			resp, err := service.DeletePremiumMigration(ctx, &v1beta2.DeletePremiumMigrationReq{
				TidbId: task.TiDBID, MigrationId: task.ID,
			})
			require.NoError(t, err)
			require.Equal(t, v1beta2.PremiumMigration_DELETING, resp.State)
			publisher.AssertExpectations(t)
		})
	}
}

func newAsyncControllerTask(stage entity.PremiumDMTaskStage) *entity.PremiumDMTask {
	return &entity.PremiumDMTask{
		ID: "async-task-1", Name: "async-migration", Stage: stage,
		TiDBID: "10001", TenantID: "20002", ProjectID: "30003",
		Mode: entity.PremiumDMTaskModeAll, ImportMode: entity.PremiumDMTaskImportModeLogical,
		CloudProvider: string(entity.CloudProviderTypeAWS), Region: "us-west-2",
	}
}