package premium_dm

import (
	"context"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/AlekSi/pointer"
	"github.com/avast/retry-go/v4"
	pinerrors "github.com/pingcap/errors"
	"github.com/samber/lo"
	"github.com/tidbcloud/dataflow-service-ng/app/event_bus/publisher/center"
	"github.com/tidbcloud/dataflow-service-ng/app/internal/console_auditlog"
	"github.com/tidbcloud/dataflow-service-ng/app/models/domain"
	"github.com/tidbcloud/dataflow-service-ng/app/models/premium_dm"
	"github.com/tidbcloud/dataflow-service-ng/app/models/sink_endpoint"
	"github.com/tidbcloud/dataflow-service-ng/app/repo"
	"github.com/tidbcloud/dataflow-service-ng/app/repo/entity"
	apputils "github.com/tidbcloud/dataflow-service-ng/app/utils"
	"github.com/tidbcloud/dataflow-service-ng/idl/pbgen/proto/premium_dm/v1beta2"
	"github.com/tidbcloud/dataflow-service-ng/internal/config"
	"github.com/tidbcloud/dataflow-service-ng/pkg/agent"
	"github.com/tidbcloud/dataflow-service-ng/pkg/crypto"
	encryptor_wrapper "github.com/tidbcloud/dataflow-service-ng/pkg/encryptor-wrapper"
	"github.com/tidbcloud/dataflow-service-ng/pkg/grpc/grpcli"
	"github.com/tidbcloud/dataflow-service-ng/pkg/log"
	agentpb "github.com/tidbcloud/dataflow-service/dataflow-agent-ng/idl/pbgen/proto/ng_agent"
	"github.com/tidbcloud/eventbus-datastruct/event"
	"github.com/tidbcloud/pkgv2/bizerr"
	"github.com/tidbcloud/pkgv2/db"
	commonv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/common/v1"
	mgmt "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
	"github.com/tidbcloud/tiinf/sdk/auditlog"
	"go.uber.org/zap"
	"google.golang.org/protobuf/types/known/emptypb"
)

const maxRetryTimesOnExecSQL uint = 3

// agentSQLLogicalError is returned when the regional agent answers HTTP 200 but sets
// error_msg on ExecSQL / ExecSQLOnDMSource (MySQL/TiDB logical errors).
type agentSQLLogicalError struct {
	msg string
}

func (e *agentSQLLogicalError) Error() string {
	return e.msg
}

type PremiumDMService struct {
	v1beta2.UnimplementedPremiumMigrationServiceServer

	dao                 *premium_dm.PremiumDMTaskDAO
	taskWorker          *premium_dm.PremiumDMTaskWorker
	mgmtClient          grpcli.MgmtSvcGrpcClient
	encryptor           encryptor_wrapper.Encryptor
	agentConfig         agent.ClientConfig
	dbProvider          db.TransProvider
	sinkEndpointModel   sink_endpoint.SinkEndpointModel
	bizSinkEndpointRepo repo.BizSinkEndpointRepo
	eventPublisher      center.EventCenterPublisher
}

func NewPremiumDMService(
	dao *premium_dm.PremiumDMTaskDAO,
	taskWorker *premium_dm.PremiumDMTaskWorker,
	mgmtClient grpcli.MgmtSvcGrpcClient,
	encryptor encryptor_wrapper.Encryptor,
	agentConfig agent.ClientConfig,
	dbProvider db.TransProvider,
	sinkEndpointModel sink_endpoint.SinkEndpointModel,
	bizSinkEndpointRepo repo.BizSinkEndpointRepo,
	eventPublisher center.EventCenterPublisher,
) *PremiumDMService {
	return &PremiumDMService{
		dao:                 dao,
		taskWorker:          taskWorker,
		mgmtClient:          mgmtClient,
		encryptor:           encryptor,
		agentConfig:         agentConfig,
		dbProvider:          dbProvider,
		sinkEndpointModel:   sinkEndpointModel,
		bizSinkEndpointRepo: bizSinkEndpointRepo,
		eventPublisher:      eventPublisher,
	}
}

func validateCreateTaskReq(req *v1beta2.CreatePremiumMigrationReq, cluster *mgmt.Cluster) error {
	allowedConnTypesBySourceType := make(map[v1beta2.PremiumDMSourceType]map[v1beta2.PremiumDMConnType]struct{})
	for _, supportedSource := range premiumSupportedSourcesByCloudProvider(cluster.GetCloudProvider()) {
		allowedConnTypes := make(map[v1beta2.PremiumDMConnType]struct{}, len(supportedSource.GetConnMethods()))
		for _, connType := range supportedSource.GetConnMethods() {
			allowedConnTypes[connType] = struct{}{}
		}
		allowedConnTypesBySourceType[supportedSource.GetSourceType()] = allowedConnTypes
	}
	for i, source := range req.GetSources() {
		if source == nil {
			continue
		}
		allowedConnTypes, ok := allowedConnTypesBySourceType[source.GetSourceType()]
		if !ok {
			return bizerr.New(bizerr.ErrCodeParams,
				"source type %s is not supported for cloud provider %s",
				source.GetSourceType().String(),
				cluster.GetCloudProvider(),
			)
		}
		if source.GetConnProfile() == nil {
			return bizerr.New(bizerr.ErrCodeParams, "source %d conn profile is required", i)
		}
		if _, ok := allowedConnTypes[source.GetConnProfile().GetConnType()]; !ok {
			return bizerr.New(bizerr.ErrCodeParams,
				"connection type %s is not supported for source type %s and cloud provider %s",
				source.GetConnProfile().GetConnType().String(),
				source.GetSourceType().String(),
				cluster.GetCloudProvider(),
			)
		}
	}
	return nil
}

func validatePremiumDMPublicSourceHosts(
	ctx context.Context, servicePlan commonv1.ServicePlan, sources []*v1beta2.PremiumDMSource,
) error {
	for i, source := range sources {
		if source == nil {
			continue
		}
		if err := validatePremiumDMPublicConnProfile(ctx, servicePlan, source.GetConnProfile()); err != nil {
			return bizerr.New(bizerr.ErrCodeParams, "source %d public host is invalid: %s", i, err.Error())
		}
	}
	return nil
}

func validatePremiumDMPublicConnProfile(
	ctx context.Context, servicePlan commonv1.ServicePlan, connProfile *v1beta2.PremiumDMConnProfile,
) error {
	// Structural validation is plan-independent: a nil conn profile is always
	// invalid and would otherwise be dereferenced downstream and panic.
	if connProfile == nil {
		return bizerr.New(bizerr.ErrCodeParams, "conn profile is required")
	}
	// Public endpoint host validation only applies to Essential V2 clusters.
	if servicePlan != commonv1.ServicePlan_Essential_V2 {
		return nil
	}
	if connProfile.GetConnType() == v1beta2.PremiumDMConnType_PRIVATELINK {
		return nil
	}
	if err := apputils.ValidatePublicEndpointHost(ctx, connProfile.GetHost()); err != nil {
		return bizerr.New(bizerr.ErrCodeParams, "public host %s is invalid: %s", connProfile.GetHost(), err.Error())
	}
	return nil
}

func premiumSupportedSourcesByCloudProvider(cloudProvider string) []*v1beta2.PremiumSupportedSource {
	sources := []*v1beta2.PremiumSupportedSource{
		{
			SourceType:  v1beta2.PremiumDMSourceType_MYSQL,
			DisplayName: "MySQL",
			ConnMethods: []v1beta2.PremiumDMConnType{
				v1beta2.PremiumDMConnType_PUBLIC,
				v1beta2.PremiumDMConnType_PRIVATELINK,
			},
		},
	}
	if strings.EqualFold(cloudProvider, string(entity.CloudProviderTypeAzure)) {
		sources = append(sources, &v1beta2.PremiumSupportedSource{
			SourceType:  v1beta2.PremiumDMSourceType_AZURE_FLEXIBLE_MYSQL,
			DisplayName: "Azure Flexible MySQL",
			ConnMethods: []v1beta2.PremiumDMConnType{v1beta2.PremiumDMConnType_PUBLIC},
		})
	}
	aliCloudRDSConnMethods := []v1beta2.PremiumDMConnType{v1beta2.PremiumDMConnType_PUBLIC}
	if strings.EqualFold(cloudProvider, string(entity.CloudProviderTypeAliCloud)) {
		aliCloudRDSConnMethods = append(aliCloudRDSConnMethods, v1beta2.PremiumDMConnType_PRIVATELINK)
	}
	sources = append(sources, &v1beta2.PremiumSupportedSource{
		SourceType:  v1beta2.PremiumDMSourceType_ALICLOUD_RDS_MYSQL,
		DisplayName: "AliCloud RDS MySQL",
		ConnMethods: aliCloudRDSConnMethods,
	})
	return sources
}

func (s *PremiumDMService) validateCluster(
	ctx context.Context, tiDBID string, logger *zap.Logger,
) (*mgmt.Cluster, error) {
	clusterID, err := strconv.ParseUint(tiDBID, 10, 64)
	if err != nil {
		return nil, err
	}
	listResp, err := s.mgmtClient.ListClusters(ctx, &mgmt.ListClustersRequest{
		Filter: &mgmt.ClusterFilter{ClusterIds: []uint64{clusterID}},
	})
	if err != nil {
		logger.Error("list mgmt clusters error", zap.Error(err), zap.Uint64("clusterID", clusterID))
		return nil, err
	}
	if len(listResp.Clusters) == 0 {
		bizErr := bizerr.New(bizerr.ErrCodeRecordNotFound, "cluster %d is not found", clusterID)
		logger.Error("cluster not found", zap.Error(bizErr), zap.Uint64("clusterID", clusterID))
		return nil, bizErr
	}
	return listResp.Clusters[0], nil
}

func isWhitelisted(orgId uint64, cluster *mgmt.Cluster) error {
	return nil
}

func (s *PremiumDMService) encryptPasswords(logger *zap.Logger, req *v1beta2.CreatePremiumMigrationReq) error {
	encrypted, err := s.encryptor.EncryptBase64(req.GetTarget().GetPassword())
	if err != nil {
		logger.Error("failed to encrypt target database password", zap.Error(err))
		return err
	}
	req.Target.Password = encrypted
	for i := range req.Sources {
		if req.Sources[i] == nil || req.Sources[i].GetConnProfile() == nil {
			continue
		}
		var enc string
		enc, err = s.encryptor.EncryptBase64(req.Sources[i].GetConnProfile().GetPassword())
		if err != nil {
			logger.Error("failed to encrypt source database password", zap.Error(err))
			return err
		}
		req.Sources[i].ConnProfile.Password = enc
	}
	return nil
}

func (s *PremiumDMService) registerPremiumDMClusterStatus(
	ctx context.Context,
	logger *zap.Logger,
	clusterID uint64,
	task *entity.PremiumDMTask,
) (func(), error) {
	clusterStatus, bizErr := s.mgmtClient.SetClusterCustomStatus(ctx, &mgmt.SetClusterCustomStatusRequest{
		CustomStatus: &mgmt.ClusterCustomStatus{
			ClusterId:  clusterID,
			StatusName: task.ID,
			StatusType: mgmt.CustomStatusType_CUSTOM_STATUS_TYPE_DATA_MIGRATION,
			Mutex: &mgmt.ClusterCustomStatus_CustomMutex{
				ClusterAction: []mgmt.ClusterAction{
					mgmt.ClusterAction_CLUSTER_ACTION_DELETE_CLUSTER,
				},
			},
		},
	})
	if bizErr != nil {
		logger.Error("failed to register cluster custom status for premium dm task", zap.Error(bizErr))
		return nil, bizErr
	}
	if clusterStatus == nil {
		err := bizerr.New(bizerr.ErrCodeDB, "cluster custom status response is nil")
		logger.Error("failed to register cluster custom status for premium dm task", zap.Error(err))
		return nil, err
	}
	if clusterStatus.GetStatusId() == 0 {
		err := bizerr.New(bizerr.ErrCodeDB, "cluster custom status response has invalid zero status id")
		logger.Error("failed to register cluster custom status for premium dm task", zap.Error(err))
		return nil, err
	}
	task.ClusterStatusID = clusterStatus.GetStatusId()
	return func() {
		rctx := context.WithoutCancel(ctx)
		delErr := s.mgmtClient.DeleteClusterCustomStatus(rctx, &mgmt.DeleteClusterCustomStatusRequest{
			ClusterId: clusterID,
			StatusId:  task.ClusterStatusID,
		})
		if delErr != nil {
			logger.Error("failed to compensate cluster custom status for premium dm task",
				zap.Uint64("cluster_status_id", task.ClusterStatusID),
				zap.Error(delErr))
			return
		}
		logger.Info("deleted cluster custom status after create premium dm task failed",
			zap.Uint64("cluster_status_id", task.ClusterStatusID))
	}, nil
}

//nolint:funlen,gocyclo
func (s *PremiumDMService) CreatePremiumMigration(
	ctx context.Context, req *v1beta2.CreatePremiumMigrationReq,
) (resp *v1beta2.PremiumMigration, err error) {
	logger := log.GetLogger(ctx).With(
		zap.String("cluster_id", req.GetTidbId()),
		zap.String("name", strings.TrimSpace(req.GetDisplayName())),
	)
	var cluster *mgmt.Cluster
	var migrationID string
	defer func() {
		if cluster == nil {
			return
		}
		console_auditlog.SubmitAsync(ctx, auditlog.EventCreateMigrationJob, cluster,
			buildPremiumMigrationCreateDetails(req, migrationID, cluster, err), err == nil)
	}()

	cluster, err = s.validateCluster(ctx, req.GetTidbId(), logger)
	if err != nil {
		return nil, err
	}
	if bizErr := validateCreateTaskReq(req, cluster); bizErr != nil {
		return nil, bizErr
	}
	if bizErr := validatePremiumDMPublicSourceHosts(
		ctx, apputils.NewCluster(cluster).GetServicePlan(), req.GetSources()); bizErr != nil {
		return nil, bizErr
	}
	if bizErr := isWhitelisted(cluster.GetOrganizationId(), cluster); bizErr != nil {
		return nil, bizErr
	}
	if err := s.encryptPasswords(logger, req); err != nil {
		return nil, err
	}

	// Batch-load all PRIVATELINK endpoints for this cluster's shared pool.
	clusterWrapper := apputils.NewCluster(cluster)
	id2endpoint, err := premium_dm.BuildEndpointMapForCluster(ctx, s.bizSinkEndpointRepo, clusterWrapper)
	if err != nil {
		return nil, err
	}
	task, bizErr := CreateV1beta1TaskReqToDMTask(ctx, req, cluster, id2endpoint)
	if bizErr != nil {
		return nil, bizErr
	}

	compensateClusterStatus, err := s.registerPremiumDMClusterStatus(ctx, logger, cluster.GetClusterId(), task)
	if err != nil {
		return nil, err
	}
	shouldCompensateClusterStatus := true
	defer func() {
		if !shouldCompensateClusterStatus || task.ClusterStatusID == 0 {
			return
		}
		compensateClusterStatus()
	}()

	if err := s.dbProvider.Transaction(ctx, func(ctx context.Context) error {
		// TODO(ideascf): move to model layer and add repo layer for premium dm
		if err := s.dbProvider.WriteDB(ctx).Create(task).Error; err != nil {
			return err
		}
		// Create sink endpoint usage records for PRIVATELINK sources
		err := s.createSinkEndpointUsagesForReqSources(
			ctx, task.ID, req.GetSources(),
			entity.BizSinkEndpointUsageTypePremiumDMTask,
		)
		if err != nil {
			return err
		}
		return nil
	}); err != nil {
		logger.Error("failed to create premium dm task", zap.Error(err))
		return nil, err
	}
	shouldCompensateClusterStatus = false

	premium_dm.PublishPremiumDMEvent(
		ctx, s.eventPublisher, task, premium_dm.PremiumDMOperationCreate, event.EventStatusStarted,
	)

	resp, err = DMTaskAndStatusToV1beta1PbTask(ctx, task, id2endpoint)
	if resp != nil {
		migrationID = resp.GetMigrationId()
	}
	return resp, err
}

func (s *PremiumDMService) CreatePremiumMigrationPrecheck(
	ctx context.Context, req *v1beta2.CreatePremiumMigrationReq,
) (*v1beta2.CreatePremiumMigrationPrecheckResp, error) {
	logger := log.GetLogger(ctx).With(
		zap.String("cluster_id", req.GetTidbId()),
		zap.String("name", strings.TrimSpace(req.GetDisplayName())),
	)
	cluster, bizErr := s.validateCluster(ctx, req.GetTidbId(), logger)
	if bizErr != nil {
		return nil, bizErr
	}
	if bizErr := validateCreateTaskReq(req, cluster); bizErr != nil {
		return nil, bizErr
	}
	if bizErr := validatePremiumDMPublicSourceHosts(
		ctx, apputils.NewCluster(cluster).GetServicePlan(), req.GetSources()); bizErr != nil {
		return nil, bizErr
	}
	if bizErr := isWhitelisted(cluster.GetOrganizationId(), cluster); bizErr != nil {
		return nil, bizErr
	}
	if err := s.encryptPasswords(logger, req); err != nil {
		return nil, err
	}
	precheck, bizErr := CreateV1beta2TaskReqToPremiumDMPrecheck(
		ctx, req, cluster, config.Global().PremiumDMConfig.PrecheckSkipConfigs,
	)
	if bizErr != nil {
		return nil, bizErr
	}
	worker := premium_dm.NewPremiumDMPrecheckWorker(s.dao, s.taskWorker)
	if err := s.dbProvider.Transaction(ctx, func(ctx context.Context) error {
		// TODO(ideascf): move to model layer and add repo layer for premium dm
		if bizErr := worker.CreatePrecheck(ctx, precheck); bizErr != nil {
			logger.Error("failed to create precheck", zap.Error(bizErr))
			return bizErr
		}
		// Create sink endpoint usage records for PRIVATELINK sources
		err := s.createSinkEndpointUsagesForReqSources(
			ctx, precheck.ID, req.GetSources(),
			entity.BizSinkEndpointUsageTypePremiumDMPrecheck,
		)
		if err != nil {
			return err
		}
		return nil
	}); err != nil {
		logger.Error("failed to create premium dm precheck", zap.Error(err))
		return nil, err
	}
	return &v1beta2.CreatePremiumMigrationPrecheckResp{
		PrecheckId: precheck.ID,
	}, nil
}

func (s *PremiumDMService) GetPremiumMigration(
	ctx context.Context, req *v1beta2.GetPremiumMigrationReq,
) (*v1beta2.PremiumMigration, error) {
	logger := log.GetLogger(ctx).With(
		zap.String("cluster_id", req.GetTidbId()),
	)
	task, bizErr := s.dao.GetTask(
		ctx,
		req.GetTidbId(),
		req.GetMigrationId(),
		entity.NonDeletedPremiumDMTaskStages,
	)
	if bizErr != nil {
		logger.Error("get dm task failed", zap.Error(bizErr))
		return nil, bizErr
	}
	endpointMap, err := s.buildId2EndpointForTasks(ctx, task)
	if err != nil {
		return nil, err
	}
	return DMTaskAndStatusToV1beta1PbTask(ctx, task, endpointMap)
}

func (s *PremiumDMService) ListPremiumMigrations(
	ctx context.Context, req *v1beta2.ListPremiumMigrationsReq,
) (*v1beta2.ListPremiumMigrationsResp, error) {
	logger := log.GetLogger(ctx).With(
		zap.String("cluster_id", req.GetTidbId()),
	)
	var err error
	filter := &entity.PremiumDMTaskFilter{
		ClusterID: req.GetTidbId(),
	}
	pager := apputils.NewPagerNumberPager()
	err = pager.Normalize(req)
	if err != nil {
		return nil, err
	}
	filter.Offset = pointer.To(int(pager.Offset()))
	filter.Limit = pointer.To(int(pager.PageSize()))
	filter.OrderBy = "created_at desc"
	filter.Stages = entity.NonDeletedPremiumDMTaskStages

	tasks, err := s.dao.ListTasks(ctx, filter)
	if err != nil {
		return nil, err
	}
	total, err := s.dao.CountTasks(ctx, filter)
	if err != nil {
		return nil, err
	}
	logger.Info(fmt.Sprintf("total: %d", total))
	nextPageToken := pager.NextPageToken(uint32(total))
	resp := &v1beta2.ListPremiumMigrationsResp{
		PremiumMigrations: make([]*v1beta2.PremiumMigration, 0, len(tasks)),
		TotalSize:         uint32(total),
		NextPageToken:     nextPageToken,
	}
	endpointMap, err := s.buildId2EndpointForTasks(ctx, tasks...)
	if err != nil {
		return nil, err
	}
	for _, t := range tasks {
		task, err := DMTaskAndStatusToV1beta1PbTask(ctx, t, endpointMap)
		if err != nil {
			return nil, err
		}
		resp.PremiumMigrations = append(resp.PremiumMigrations, task)
	}
	return resp, nil
}

func (s *PremiumDMService) DeletePremiumMigration(
	ctx context.Context, req *v1beta2.DeletePremiumMigrationReq,
) (resp *v1beta2.PremiumMigration, err error) {
	logger := log.GetLogger(ctx).With(
		zap.String("cluster_id", req.GetTidbId()),
	)
	var cluster *mgmt.Cluster
	var taskName string
	defer func() {
		if cluster == nil {
			return
		}
		console_auditlog.SubmitAsync(ctx, auditlog.EventDeleteMigrationJob, cluster,
			buildPremiumMigrationOperateDetails(req.GetMigrationId(), taskName, cluster, err), err == nil)
	}()

	cluster, err = s.validateCluster(ctx, req.TidbId, logger)
	if err != nil {
		return nil, err
	}
	if bizErr := isWhitelisted(cluster.GetOrganizationId(), cluster); bizErr != nil {
		return nil, bizErr
	}
	task, bizErr := s.dao.GetTask(ctx, req.GetTidbId(), req.GetMigrationId(), nil)
	if bizErr != nil {
		logger.Error("get dm task failed", zap.Error(bizErr))
		return nil, bizErr
	}
	taskName = task.Name
	endpointMap, err := s.buildId2EndpointForTasks(ctx, task)
	if err != nil {
		return nil, err
	}
	if task.Stage == entity.PremiumDMTaskStageDeleted || task.Stage == entity.PremiumDMTaskStageDeleting {
		return DMTaskAndStatusToV1beta1PbTask(ctx, task, endpointMap)
	}
	var supersededOperation premium_dm.PremiumDMOperation
	supersededStage := task.Stage
	switch task.Stage {
	case entity.PremiumDMTaskStagePausing:
		supersededOperation = premium_dm.PremiumDMOperationPause
	case entity.PremiumDMTaskStageResuming:
		supersededOperation = premium_dm.PremiumDMOperationResume
	}
	bizErr = s.dao.UpdateTask(ctx, task.ID, entity.ActivePremiumDMTaskStages, map[string]any{
		"stage":        entity.PremiumDMTaskStageDeleting,
		"deleted_time": time.Now(),
	})
	if bizErr != nil {
		logger.Error("update dm task failed", zap.Error(bizErr))
		return nil, bizErr
	}
	task.Stage = entity.PremiumDMTaskStageDeleting
	if supersededOperation != "" {
		supersededTask := *task
		supersededTask.Stage = supersededStage
		premium_dm.PublishPremiumDMEvent(
			ctx, s.eventPublisher, &supersededTask, supersededOperation, event.EventStatusCanceled,
		)
	}
	premium_dm.PublishPremiumDMEvent(
		ctx, s.eventPublisher, task, premium_dm.PremiumDMOperationDelete, event.EventStatusStarted,
	)
	task, bizErr = s.dao.GetTask(ctx, req.GetTidbId(), req.GetMigrationId(), nil)
	if bizErr != nil {
		logger.Error("get dm task failed", zap.Error(bizErr))
		return nil, bizErr
	}
	return DMTaskAndStatusToV1beta1PbTask(ctx, task, endpointMap)
}

func (s *PremiumDMService) PausePremiumMigration(
	ctx context.Context, req *v1beta2.PausePremiumMigrationReq,
) (_ *emptypb.Empty, err error) {
	logger := log.GetLogger(ctx).With(
		zap.String("cluster_id", req.GetTidbId()),
		zap.String("migration_id", req.GetMigrationId()),
	)
	var cluster *mgmt.Cluster
	var taskName string
	defer func() {
		if cluster == nil {
			return
		}
		console_auditlog.SubmitAsync(ctx, auditlog.EventSuspendMigrationJob, cluster,
			buildPremiumMigrationOperateDetails(req.GetMigrationId(), taskName, cluster, err), err == nil)
	}()

	cluster, err = s.validateCluster(ctx, req.GetTidbId(), logger)
	if err != nil {
		return nil, err
	}

	var acceptedTask *entity.PremiumDMTask
	err = s.dbProvider.Transaction(ctx, func(ctx context.Context) error {
		task, bizErr := s.dao.GetTask(ctx, req.GetTidbId(), req.GetMigrationId(), nil)
		if bizErr != nil {
			logger.Error("get dm task failed", zap.Error(bizErr))
			return bizErr
		}
		taskName = task.Name
		if task.Stage != entity.PremiumDMTaskStageRunning {
			return bizerr.New(bizerr.ErrCodeParams, "dm task %s is not running", task.ID)
		}

		rowsAffected, updErr := s.dao.UpdateTaskWithRowsAffected(
			ctx,
			task.ID,
			[]entity.PremiumDMTaskStage{entity.PremiumDMTaskStageRunning},
			map[string]any{
				"stage":                   entity.PremiumDMTaskStagePausing,
				"operation_dispatched_at": nil,
			},
		)
		if updErr != nil {
			logger.Error("update dm task stage failed", zap.Error(updErr))
			return updErr
		}
		if rowsAffected == 0 {
			return bizerr.New(
				bizerr.ErrCodeOptimisticLockConflict,
				"premium dm task %s was changed by another operation",
				task.ID,
			)
		}
		task.Stage = entity.PremiumDMTaskStagePausing
		task.OperationDispatchedAt = nil
		acceptedTask = task
		return nil
	})
	if err != nil {
		logger.Error("failed to accept pause for dm task", zap.Error(err))
		return nil, err
	}
	premium_dm.PublishPremiumDMEvent(
		ctx, s.eventPublisher, acceptedTask, premium_dm.PremiumDMOperationPause, event.EventStatusStarted,
	)
	return &emptypb.Empty{}, nil
}

func (s *PremiumDMService) ResumePremiumMigration(
	ctx context.Context, req *v1beta2.ResumePremiumMigrationReq,
) (_ *emptypb.Empty, err error) {
	logger := log.GetLogger(ctx).With(
		zap.String("cluster_id", req.GetTidbId()),
		zap.String("migration_id", req.GetMigrationId()),
	)
	var cluster *mgmt.Cluster
	var taskName string
	defer func() {
		if cluster == nil {
			return
		}
		console_auditlog.SubmitAsync(ctx, auditlog.EventResumeMigrationJob, cluster,
			buildPremiumMigrationOperateDetails(req.GetMigrationId(), taskName, cluster, err), err == nil)
	}()

	cluster, err = s.validateCluster(ctx, req.GetTidbId(), logger)
	if err != nil {
		return nil, err
	}

	var acceptedTask *entity.PremiumDMTask
	err = s.dbProvider.Transaction(ctx, func(ctx context.Context) error {
		task, bizErr := s.dao.GetTask(ctx, req.GetTidbId(), req.GetMigrationId(), nil)
		if bizErr != nil {
			logger.Error("get dm task failed", zap.Error(bizErr))
			return bizErr
		}
		taskName = task.Name
		if task.Stage != entity.PremiumDMTaskStageStopped && task.Stage != entity.PremiumDMTaskStageFailed {
			return bizerr.New(
				bizerr.ErrCodeParams,
				"premium dm task %s cannot be resumed, current stage: %s",
				task.ID,
				task.Stage,
			)
		}
		rowsAffected, updErr := s.dao.UpdateTaskWithRowsAffected(
			ctx,
			task.ID,
			[]entity.PremiumDMTaskStage{
				entity.PremiumDMTaskStageStopped,
				entity.PremiumDMTaskStageFailed,
			},
			map[string]any{
				"stage":                   entity.PremiumDMTaskStageResuming,
				"operation_dispatched_at": nil,
			},
		)
		if updErr != nil {
			logger.Error("update premium dm task stage failed", zap.Error(updErr))
			return updErr
		}
		if rowsAffected == 0 {
			return bizerr.New(
				bizerr.ErrCodeOptimisticLockConflict,
				"premium dm task %s was changed by another operation",
				task.ID,
			)
		}
		task.Stage = entity.PremiumDMTaskStageResuming
		task.OperationDispatchedAt = nil
		acceptedTask = task
		return nil
	})
	if err != nil {
		logger.Error("failed to accept resume for premium migration", zap.Error(err))
		return nil, err
	}
	premium_dm.PublishPremiumDMEvent(
		ctx, s.eventPublisher, acceptedTask, premium_dm.PremiumDMOperationResume, event.EventStatusStarted,
	)
	return &emptypb.Empty{}, nil
}

func (s *PremiumDMService) PremiumMigrationCancelPrecheck(
	ctx context.Context, req *v1beta2.CancelPremiumMigrationPrecheckReq,
) (*emptypb.Empty, error) {
	logger := log.GetLogger(ctx).With(
		zap.String("cluster_id", req.GetTidbId()),
		zap.String("precheck_id", req.GetPrecheckId()),
	)

	cluster, bizErr := s.validateCluster(ctx, req.GetTidbId(), logger)
	if bizErr != nil {
		return nil, bizErr
	}
	if bizErr := isWhitelisted(cluster.GetOrganizationId(), cluster); bizErr != nil {
		return nil, bizErr
	}

	precheck, bizErr := s.dao.GetPrecheck(ctx, req.GetTidbId(), req.GetPrecheckId())
	if bizErr != nil {
		logger.Error("failed to get precheck", zap.Error(bizErr))
		return nil, bizErr
	}

	if precheck.Stage == string(entity.PremiumDMPrecheckStageFinished) ||
		precheck.Stage == string(entity.PremiumDMPrecheckStageFailed) ||
		precheck.Stage == string(entity.PremiumDMPrecheckStageCanceled) {
		return &emptypb.Empty{}, nil
	}

	bizErr = s.dao.UpdatePrecheck(ctx, precheck.ID, map[string]any{
		"stage":         entity.PremiumDMPrecheckStageCanceled,
		"canceled_time": time.Now(),
	})
	if bizErr != nil {
		logger.Error("failed to cancel precheck", zap.Error(bizErr))
		return nil, bizErr
	}

	return &emptypb.Empty{}, nil
}

func (s *PremiumDMService) PremiumMigrationGetPrecheck(
	ctx context.Context, req *v1beta2.GetPremiumMigrationPrecheckReq,
) (*v1beta2.PremiumMigrationPrecheck, error) {
	logger := log.GetLogger(ctx).With(
		zap.String("cluster_id", req.GetTidbId()),
		zap.String("precheck_id", req.GetPrecheckId()),
	)

	cluster, bizErr := s.validateCluster(ctx, req.GetTidbId(), logger)
	if bizErr != nil {
		return nil, bizErr
	}
	if bizErr := isWhitelisted(cluster.GetOrganizationId(), cluster); bizErr != nil {
		return nil, bizErr
	}

	precheck, bizErr := s.dao.GetPrecheck(ctx, req.GetTidbId(), req.GetPrecheckId())
	if bizErr != nil {
		logger.Error("failed to get precheck", zap.Error(bizErr))
		return nil, bizErr
	}

	return PremiumDMPrecheckToV1beta2PbPrecheck(precheck), nil
}

// ListPremiumMigrationPrechecks lists precheck records of a tidb cluster in
// reverse-chronological order. Per-item details are not included; callers
// fetch them via PremiumMigrationGetPrecheck.
func (s *PremiumDMService) ListPremiumMigrationPrechecks(
	ctx context.Context, req *v1beta2.ListPremiumMigrationPrechecksReq,
) (*v1beta2.ListPremiumMigrationPrechecksResp, error) {
	logger := log.GetLogger(ctx).With(
		zap.String("cluster_id", req.GetTidbId()),
	)

	cluster, bizErr := s.validateCluster(ctx, req.GetTidbId(), logger)
	if bizErr != nil {
		return nil, bizErr
	}
	if bizErr := isWhitelisted(cluster.GetOrganizationId(), cluster); bizErr != nil {
		return nil, bizErr
	}

	pager := apputils.NewPagerNumberPager()
	if err := pager.Normalize(req); err != nil {
		return nil, err
	}
	filter := &entity.PremiumDMPrecheckFilter{
		ClusterID: req.GetTidbId(),
		OrderBy:   "created_at desc",
		Offset:    pointer.To(int(pager.Offset())),
		Limit:     pointer.To(int(pager.PageSize())),
	}
	prechecks, bizErr := s.dao.ListPrechecks(ctx, filter)
	if bizErr != nil {
		logger.Error("failed to list prechecks", zap.Error(bizErr))
		return nil, bizErr
	}
	total, bizErr := s.dao.CountPrechecks(ctx, filter)
	if bizErr != nil {
		logger.Error("failed to count prechecks", zap.Error(bizErr))
		return nil, bizErr
	}

	resp := &v1beta2.ListPremiumMigrationPrechecksResp{
		Prechecks:     make([]*v1beta2.PremiumMigrationPrecheckSummary, 0, len(prechecks)),
		TotalSize:     total,
		NextPageToken: pager.NextPageToken(total),
	}
	for _, p := range prechecks {
		resp.Prechecks = append(resp.Prechecks, PremiumDMPrecheckToV1beta2PbSummary(p))
	}
	return resp, nil
}

func (s *PremiumDMService) GetPremiumMigrationSourceTables(
	ctx context.Context, req *v1beta2.PremiumMigrationSourceDBConnProfile,
) (*v1beta2.GetSourceTablesResp, error) {
	sql := `SELECT table_schema, table_name
FROM information_schema.tables
WHERE table_schema NOT IN ('information_schema', 'mysql', 'performance_schema', 'sys')
  AND table_type = 'BASE TABLE'
ORDER BY table_schema, table_name`
	execRes, err := s.execSQLOnSource(ctx, req, sql)
	if err != nil {
		return nil, err
	}
	res := &v1beta2.GetSourceTablesResp{
		Tables: make([]*v1beta2.GetSourceTablesResp_Table, len(execRes.Rows)),
	}
	schemaIndex := -1
	tableIndex := -1
	for j, t := range execRes.ColumnNames {
		if strings.EqualFold("table_schema", t) {
			schemaIndex = j
		} else if strings.EqualFold("table_name", t) {
			tableIndex = j
		}
	}
	if schemaIndex == -1 || tableIndex == -1 {
		return nil, errors.New("invalid table schema or table name")
	}
	for i, row := range execRes.GetRows() {
		res.Tables[i] = &v1beta2.GetSourceTablesResp_Table{
			Schema: row.TypedColumnValues[schemaIndex].GetStringValue(),
			Table:  row.TypedColumnValues[tableIndex].GetStringValue(),
		}
	}
	return res, nil
}

func (s *PremiumDMService) ValidatePremiumMigrationSourceConn(
	ctx context.Context, req *v1beta2.PremiumMigrationSourceDBConnProfile,
) (*emptypb.Empty, error) {
	_, err := s.execSQLOnSource(ctx, req, "SELECT 1")
	if err != nil {
		return nil, err
	}
	return &emptypb.Empty{}, nil
}

//nolint:funlen // SQL execution setup has many sequential steps
func (s *PremiumDMService) execSQLOnSource(
	ctx context.Context,
	req *v1beta2.PremiumMigrationSourceDBConnProfile,
	sql string,
) (*agentpb.ExecSQLOnDMSourceResponse, error) {
	logger := log.GetLogger(ctx).With(
		zap.String("cluster_id", req.GetTidbId()),
	)
	logger.With(zap.String("sql", sql)).Info("exec sql on dm source")
	cluster, err := s.validateCluster(ctx, req.GetTidbId(), logger)
	if err != nil {
		return nil, err
	}
	if err := validatePremiumDMPublicConnProfile(
		ctx, apputils.NewCluster(cluster).GetServicePlan(), req.GetConnProfile()); err != nil {
		return nil, err
	}
	clusterID, err := strconv.ParseUint(req.GetTidbId(), 10, 64)
	if err != nil {
		return nil, err
	}
	agentCli, err := agent.NewNGRegionalClientByClusterID(ctx, clusterID, s.mgmtClient, &s.agentConfig)
	if err != nil {
		return nil, pinerrors.Annotate(err, "failed to create agent client")
	}
	var resp *agentpb.ExecSQLOnDMSourceResponse
	err = retry.Do(func() error {
		pemResp, err := agentCli.GetPublicKey(ctx, &agentpb.GetPublicKeyRequest{})
		if err != nil {
			return err
		}
		pem := pemResp.GetPem()
		cipherPasswd, err := crypto.EncryptWithPEM(pem, []byte(req.GetConnProfile().GetPassword()))
		if err != nil {
			return retry.Unrecoverable(err)
		}
		execReq := &agentpb.ExecSQLOnDMSourceRequest{
			Sql: sql,
			ConnProfile: &agentpb.ConnProfile{
				User: req.GetConnProfile().GetUser(),
				Port: req.GetConnProfile().GetPort(),
				EncryptedPassword: &agentpb.EncryptedData{
					Pem:  pem,
					Data: cipherPasswd,
				},
			},
		}
		// Always use ConnProfile_Host: the agent's ExecSQLOnDMSource only reads the
		// host field to build the MySQL DSN. For PrivateLink sources we resolve the
		// sink-endpoint address and pass it as host.
		switch req.GetConnProfile().GetConnType() {
		case v1beta2.PremiumDMConnType_PRIVATELINK:
			endpointAddr, resolveErr := s.resolveSinkEndpointAddress(
				ctx, clusterID, req.GetConnProfile().GetSinkEndpointId())
			if resolveErr != nil {
				return retry.Unrecoverable(resolveErr)
			}
			execReq.ConnProfile.ConnType = agentpb.ConnType_PRIVATE_LINK
			execReq.ConnProfile.Endpoint = &agentpb.ConnProfile_Host{Host: endpointAddr}
		default:
			execReq.ConnProfile.Endpoint = &agentpb.ConnProfile_Host{Host: req.GetConnProfile().GetHost()}
		}
		if req.GetConnProfile().Security != nil {
			execReq.ConnProfile.Security = &agentpb.Security{
				CertAllowedCn:  req.GetConnProfile().Security.CertAllowedCn,
				SslCaContent:   req.GetConnProfile().Security.SslCaContent,
				SslCertContent: req.GetConnProfile().Security.SslCertContent,
				SslKeyContent:  req.GetConnProfile().Security.SslKeyContent,
			}
		}
		resp, err = agentCli.ExecSQLOnDMSource(ctx, execReq)
		if err != nil {
			return retry.Unrecoverable(err)
		}
		if resp.GetPemChanged() {
			return agent.PEMChangedError
		}
		if resp.GetErrorMsg() != "" {
			return &agentSQLLogicalError{msg: resp.GetErrorMsg()}
		}
		return nil
	},
		retry.Attempts(maxRetryTimesOnExecSQL),
		retry.LastErrorOnly(true),
		retry.RetryIf(func(err error) bool {
			return errors.Is(err, agent.PEMChangedError)
		}),
	)

	if err != nil {
		var sqlLogical *agentSQLLogicalError
		if errors.As(err, &sqlLogical) {
			return nil, err
		}
		logger.Info("failed to call execute sql in check database connectivity",
			zap.Error(err), zap.String("url", req.GetTidbId()))
		return nil, errors.New("failed to connect database")
	}
	logger.With(zap.String("sql", sql), zap.Any("resp", resp)).
		Info("success to exec sql on dm source")
	return resp, nil
}

//nolint:funlen // connection validation has many sequential steps
func (s *PremiumDMService) ValidatePremiumMigrationMigrationTargetConn(
	ctx context.Context, req *v1beta2.PremiumMigrationTargetDBConnProfile,
) (*emptypb.Empty, error) {
	clusterID, err := strconv.ParseUint(req.GetTidbId(), 10, 64)
	if err != nil {
		return nil, err
	}
	agentCli, err := agent.NewNGRegionalClientByClusterID(ctx, clusterID, s.mgmtClient, &s.agentConfig)
	if err != nil {
		return nil, pinerrors.Annotate(err, "failed to create agent client")
	}
	var execClusterResp *agentpb.ExecSQLResp
	err = retry.Do(func() error {
		pemResp, err := agentCli.GetPublicKey(ctx, &agentpb.GetPublicKeyRequest{})
		if err != nil {
			return err
		}
		pem := pemResp.GetPem()
		cipherPasswd, err := crypto.EncryptWithPEM(pem, []byte(req.GetPassword()))
		if err != nil {
			return retry.Unrecoverable(err)
		}
		execReq := &agentpb.ExecSQLReq{
			Sql: "SELECT 1",
			Cluster: &agentpb.Cluster{
				ClusterId:         clusterID,
				User:              req.User,
				PasswordEncrypted: true,
				EncryptedPassword: &agentpb.EncryptedData{
					Pem:  pem,
					Data: cipherPasswd,
				},
			},
			TypedColumns: true,
			ColumnNum:    1,
		}
		var sqlErr error
		execClusterResp, sqlErr = agentCli.ExecuteSQL(ctx, execReq)
		if sqlErr != nil {
			return retry.Unrecoverable(sqlErr)
		}
		if execClusterResp.GetPemChanged() {
			return agent.PEMChangedError
		}
		if execClusterResp.GetErrorMsg() != "" {
			return &agentSQLLogicalError{msg: execClusterResp.GetErrorMsg()}
		}
		return nil
	},
		retry.Attempts(maxRetryTimesOnExecSQL),
		retry.LastErrorOnly(true),
		retry.RetryIf(func(err error) bool {
			return errors.Is(err, agent.PEMChangedError)
		}),
	)

	if err != nil {
		var sqlLogical *agentSQLLogicalError
		if errors.As(err, &sqlLogical) {
			return nil, err
		}
		log.GetLogger(ctx).Info("failed to call execute sql in check database connectivity",
			zap.Error(err), zap.String("url", req.GetTidbId()))
		return nil, errors.New("failed to connect database")
	}
	return &emptypb.Empty{}, nil
}

// resolveSinkEndpointAddress queries the BizSinkEndpoint by cluster-scoped ID and returns its resolved endpoint address.
func (s *PremiumDMService) resolveSinkEndpointAddress(
	ctx context.Context, clusterID uint64, sinkEndpointID uint64,
) (string, error) {
	resp, err := s.sinkEndpointModel.QuerySinkEndpoint(ctx, &domain.QuerySinkEndpointReq{
		ID:        sinkEndpointID,
		ClusterID: clusterID,
	})
	if err != nil {
		return "", fmt.Errorf("query sink endpoint %d: %w", sinkEndpointID, err)
	}
	ep := resp.SinkEndpoint
	if ep.Status != entity.BizPLEndpointStatusAvailable {
		return "", fmt.Errorf("sink endpoint %d is not available (status=%s)", sinkEndpointID, ep.Status)
	}
	if ep.EndpointAddress == "" {
		return "", fmt.Errorf("sink endpoint %d has no resolved address", sinkEndpointID)
	}
	return ep.EndpointAddress, nil
}

// buildId2EndpointForTasks collects all unique PRIVATELINK sink endpoint IDs from
// the given tasks and batch-loads them into a map keyed by endpoint ID.
func (s *PremiumDMService) buildId2EndpointForTasks(
	ctx context.Context,
	tasks ...*entity.PremiumDMTask,
) (map[uint64]*entity.BizSinkEndpoint, error) {
	idSet := make(map[entity.ID]struct{})
	for _, t := range tasks {
		for i := range t.DMSources {
			if t.DMSources[i].ConnProfile.Type == entity.PremiumConnTypePrivateLink {
				idSet[entity.ID(t.DMSources[i].ConnProfile.SinkEndpointID)] = struct{}{}
			}
		}
	}
	if len(idSet) == 0 {
		return nil, nil
	}
	endpoints, bizErr := s.bizSinkEndpointRepo.List(ctx,
		func(qs entity.BizSinkEndpointQuerySet) entity.BizSinkEndpointQuerySet {
			return qs.IDIn(lo.Keys(idSet)...)
		})
	if bizErr != nil {
		return nil, bizErr
	}
	m := lo.SliceToMap(endpoints,
		func(e entity.BizSinkEndpoint) (uint64, *entity.BizSinkEndpoint) {
			return uint64(e.ID), &e
		},
	)
	return m, nil
}

// createSinkEndpointUsagesForReqSources creates BizSinkEndpointUsage records for
// each source with ConnType=PRIVATELINK in the request.
func (s *PremiumDMService) createSinkEndpointUsagesForReqSources(
	ctx context.Context,
	resourceID string,
	sources []*v1beta2.PremiumDMSource,
	usageType entity.BizSinkEndpointUsageType,
) error {
	for _, src := range sources {
		cp := src.GetConnProfile()
		if cp == nil || cp.GetConnType() != v1beta2.PremiumDMConnType_PRIVATELINK {
			continue
		}
		sinkEndpointID := cp.GetSinkEndpointId()
		if _, err := s.sinkEndpointModel.CreateSinkEndpointUsage(ctx, &domain.CreateSinkEndpointUsageReq{
			SinkEndpointID: uint64(sinkEndpointID),
			OwnerID:        resourceID,
			UsageType:      usageType,
		}); err != nil {
			return fmt.Errorf("create sink endpoint usage for resource %s: %w", resourceID, err)
		}
	}
	return nil
}

func (s *PremiumDMService) GetSupportedSources(
	ctx context.Context, req *v1beta2.PremiumMigrationGetSupportedSourcesReq,
) (*v1beta2.PremiumMigrationGetSupportedSourcesResp, error) {
	logger := log.GetLogger(ctx).With(zap.String("cluster_id", req.GetTidbId()))
	cluster, err := s.validateCluster(ctx, req.GetTidbId(), logger)
	if err != nil {
		return nil, err
	}
	return &v1beta2.PremiumMigrationGetSupportedSourcesResp{
		Sources: premiumSupportedSourcesByCloudProvider(cluster.GetCloudProvider()),
	}, nil
}