package dm_svc

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/AlekSi/pointer"
	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/hashicorp/go-multierror"
	set "github.com/hashicorp/go-set/v3"
	"github.com/tidbcloud/pkgv2/bizerr"
	"github.com/tidbcloud/pkgv2/clusterproxy/serverless/dm"
	"github.com/tidbcloud/pkgv2/clusterproxy/tidbcloud"
	"github.com/tidbcloud/pkgv2/import/cloud_storage"
	"github.com/tidbcloud/serverless-service/internal/config"
	"github.com/tidbcloud/serverless-service/internal/domain/constant"
	"github.com/tidbcloud/serverless-service/internal/domain/entity"
	"github.com/tidbcloud/serverless-service/internal/domain/facade"
	"github.com/tidbcloud/serverless-service/internal/domain/repo"
	"github.com/tidbcloud/serverless-service/internal/domain/service"
	"github.com/tidbcloud/serverless-service/internal/errorcode"
	"github.com/tidbcloud/serverless-service/internal/utils"
	tidbmgmtservicev1 "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
	"go.uber.org/zap"
)

const MySQLAccessDeniedErrorString = "Error 1045 (28000)"

// DMTaskWorker is responsible for managing the lifecycle of DM (Data Migration) tasks.
// It handles different stages of task execution including source creation, binding, task creation, and deletion.
type DMTaskWorker struct {
	dmTaskRepo                 repo.DMTaskRepo
	clusterProxyProviderFacade facade.ClusterProxyProviderFacade
	encryptorFacade            facade.EncryptorFacade
	mgmtSvc                    facade.MgmtSvcFacade
	privateLinkConnectionSvc   service.PrivateLinkConnectionService
}

// dmTaskRuntime carries ephemeral runtime-only data derived at worker run time
// that should not be persisted into the database model (e.g. endpoints which
// may change). It is passed down the call chain instead of adding transient
// fields onto entity.DMTask.
type dmTaskRuntime struct {
	clusterProxyURL string
}

func NewDMTaskWorker(
	dmTaskRepo repo.DMTaskRepo,
	clusterProxyProviderFacade facade.ClusterProxyProviderFacade,
	encryptorFacade facade.EncryptorFacade,
	mgmtSvc facade.MgmtSvcFacade,
	privateLinkConnectionSvc service.PrivateLinkConnectionService,
) *DMTaskWorker {
	return &DMTaskWorker{
		dmTaskRepo:                 dmTaskRepo,
		clusterProxyProviderFacade: clusterProxyProviderFacade,
		encryptorFacade:            encryptorFacade,
		mgmtSvc:                    mgmtSvc,
		privateLinkConnectionSvc:   privateLinkConnectionSvc,
	}
}

func (w *DMTaskWorker) Run(ctx context.Context) error {
	logger := utils.GetWorkerLoggerFromContext(ctx)
	startTime := time.Now()
	logger.Info("start to run dm task worker")

	tasks, clusterMap, err := w.fetchRunnableTasks(ctx, logger)
	if err != nil {
		return err
	}
	logger.Info("collected tasks", zap.Int("taskCount", len(tasks)))

	var errs error
	for _, task := range tasks {
		runtime, ok := w.buildRuntimeForTask(logger, task, clusterMap)
		if !ok {
			continue
		}
		if err := w.handleTask(ctx, task, runtime); err != nil {
			errs = multierror.Append(errs, err)
		}
	}

	logger.Info("finished to run dm task worker", zap.Int("taskCount", len(tasks)),
		zap.Any("costSeconds", time.Since(startTime).Seconds()))
	return errs
}

// fetchRunnableTasks returns tasks in runnable stages and a map of clusterId->Cluster
func (w *DMTaskWorker) fetchRunnableTasks(
	ctx context.Context,
	logger *zap.Logger,
) ([]*entity.DMTask, map[string]*tidbmgmtservicev1.Cluster, error) {
	tasks, bizErr := w.dmTaskRepo.ListTasks(ctx, &entity.DMTaskFilter{
		Stages: []constant.DMTaskStage{
			constant.DMTaskStageSourceCreating,
			constant.DMTaskStageDeleting,
			constant.DMTaskStageTaskCreating,
			constant.DMTaskStageSourceBounding,
			constant.DMTaskStageRunning,
			constant.DMTaskStageStopped,
			constant.DMTaskStageFailed,
		},
	})
	if bizErr != nil {
		logger.Error("failed to list dm tasks", zap.Error(bizErr))
		return nil, nil, bizErr
	}
	if len(tasks) == 0 {
		return tasks, map[string]*tidbmgmtservicev1.Cluster{}, nil
	}

	clusterIDs := set.New[uint64](len(tasks))
	for _, t := range tasks {
		cID, err := strconv.ParseUint(t.ClusterID, 10, 64)
		if err != nil {
			return nil, nil, err
		}
		clusterIDs.Insert(cID)
	}
	clusters, err := w.mgmtSvc.ListAllClusters(ctx, &tidbmgmtservicev1.ListClustersRequest{
		Filter: &tidbmgmtservicev1.ClusterFilter{
			ClusterIds:    clusterIDs.Slice(),
			ExcludeStates: facade.DeletedClusterStatus,
			View:          tidbmgmtservicev1.ClusterFilter_CLUSTER_VIEW_BASIC,
			ClusterType:   tidbmgmtservicev1.DeployType_DEPLOY_TYPE_SERVERLESS.Enum(),
		},
	})
	if err != nil {
		logger.Error("failed to list clusters from mgmt svc", zap.Error(err))
		return nil, nil, err
	}
	clusterMap := make(map[string]*tidbmgmtservicev1.Cluster, len(clusters))
	for i := range clusters {
		clusterMap[strconv.FormatUint(clusters[i].ClusterId, 10)] = clusters[i]
	}
	return tasks, clusterMap, nil
}

// buildRuntimeForTask derives runtime info for a task. Returns ok=false when task should be skipped.
func (w *DMTaskWorker) buildRuntimeForTask(
	logger *zap.Logger,
	task *entity.DMTask,
	clusterMap map[string]*tidbmgmtservicev1.Cluster,
) (*dmTaskRuntime, bool) {
	cluster, exist := clusterMap[task.ClusterID]
	if !exist {
		logger.Error("cluster not exist", zap.String("cluster_id", task.ClusterID))
		return nil, false
	}
	serverlessCluster := cluster.GetServerlessCluster()
	if serverlessCluster == nil || serverlessCluster.GetCluster() == nil {
		logger.Error("cluster missing serverless cluster info", zap.String("cluster_id", task.ClusterID))
		return nil, false
	}
	endpoint := serverlessCluster.GetCluster().GetProxyEndpoint()
	if endpoint == "" {
		logger.Error(
			"invalid cluster proxy endpoint",
			zap.String("cluster_id", task.ClusterID),
			zap.String("endpoint", endpoint),
		)
		return nil, false
	}
	return &dmTaskRuntime{clusterProxyURL: endpoint}, true
}

func (w *DMTaskWorker) handleTask(ctx context.Context, task *entity.DMTask, rt *dmTaskRuntime) error {
	logger := utils.GetWorkerLoggerFromContext(ctx).With(
		zap.String("task_id", task.ID),
		zap.String("task_name", task.Name),
	)

	switch task.Stage {
	case constant.DMTaskStageSourceCreating:
		err := w.OnTaskSourceCreating(ctx, task, rt)
		if err != nil {
			return err
		}
	case constant.DMTaskStageSourceBounding:
		err := w.OnTaskSourceBounding(ctx, task, rt)
		if err != nil {
			return err
		}
	case constant.DMTaskStageTaskCreating:
		err := w.OnTaskCreating(ctx, task, rt)
		if err != nil {
			return err
		}
	case constant.DMTaskStageDeleting:
		err := w.OnTaskDeleting(ctx, task, rt)
		if err != nil {
			return err
		}
	case constant.DMTaskStageRunning, constant.DMTaskStageStopped, constant.DMTaskStageFailed:
		err := w.OnTaskPolling(ctx, task, rt)
		if err != nil {
			return err
		}
	default:
		logger.Warn("unknown dm task stage", zap.String("stage", string(task.Stage)))
		return errors.New("unknown dm task stage")
	}
	return nil
}

// OnTaskSourceCreating handles the source creation stage of a DM task.
// It creates DM sources in the cluster and updates their status to bounding stage.
func (w *DMTaskWorker) OnTaskSourceCreating(ctx context.Context, task *entity.DMTask, rt *dmTaskRuntime) error {
	logger := utils.GetWorkerLoggerFromContext(ctx).With(
		zap.String("task_id", task.ID),
		zap.String("task_name", task.Name),
		zap.String("stage", string(task.Stage)),
	)
	logger.Info("OnTaskSourceCreating", zap.Int("total_sources", len(task.DMSources)))
	namespace, service := constant.DMProxyTarget()
	proxy, bizErr := w.clusterProxyProviderFacade.GetServerlessHTTPProxy(ctx, rt.clusterProxyURL, namespace, service)
	if bizErr != nil {
		logger.Error("failed to get dm proxy from cluster", zap.Error(bizErr))
		return bizErr
	}

	for i := range task.DMSources {
		s := &task.DMSources[i]
		if s.Stage != constant.DMSourceStageCreating {
			continue
		}

		req := &dm.CreateSourceRequest{
			Source: dm.Source{
				Enable: true,
				// For mode "all", users do not need to manually provide the GTID; it will be handled automatically in the code.
				EnableGtid: s.BinlogGTID != "" && task.Mode == constant.DMTaskModeIncr,
				Port:       s.ConnProfile.Port,
				SourceName: s.ID,
				User:       s.ConnProfile.User,
			},
		}

		switch s.ConnProfile.Type {
		case constant.ConnTypePublic:
			req.Source.Host = s.ConnProfile.Target
		case constant.ConnTypePrivateLink:
			req.Source.Host = s.EndpointDomain
		}

		password, err := w.encryptorFacade.DecryptBase64(s.ConnProfile.Password)
		if err != nil {
			logger.Error("failed to decrypt source password", zap.Error(err))
			return err
		}
		passwordStr := dm.NewNullableString(aws.String(password))
		req.Source.Password = *passwordStr
		if s.ConnProfile.Security != nil {
			security := dm.NewNullableSecurity(&dm.Security{
				CertAllowedCn:  s.ConnProfile.Security.CertAllowedCn,
				SslCaContent:   string(s.ConnProfile.Security.SslCaContent),
				SslCertContent: string(s.ConnProfile.Security.SslCertContent),
				SslKeyContent:  string(s.ConnProfile.Security.SslKeyContent),
			})
			req.Source.Security = *security
		}
		_, err = proxy.CreateServerlessDMSource(ctx, req)
		if err != nil {
			logger.Error("failed to create dm source", zap.Error(err))
			return err
		}
		err = w.dmTaskRepo.UpdateSource(ctx, s.ID, map[string]any{
			"stage":           constant.DMSourceStageBounding,
			"endpoint_id":     s.EndpointID,
			"endpoint_domain": s.EndpointDomain,
		})
		if err != nil {
			logger.Error("failed to update dm source stage to bounding", zap.Error(err))
			return err
		}
	}

	// Update task stage to source bounding when all sources are created
	err := w.dmTaskRepo.UpdateTask(
		ctx,
		task.ID,
		[]constant.DMTaskStage{constant.DMTaskStageSourceCreating},
		map[string]any{
			"stage": constant.DMTaskStageSourceBounding,
		},
	)
	if err != nil {
		logger.Error("failed to update dm task stage to sourceBounding", zap.Error(err))
		return err
	}

	return nil
}

// OnTaskSourceBounding waits for all DM sources to be bound to workers in the DM cluster.
// Once all sources are bound, it transitions the task to the task creation stage.
func (w *DMTaskWorker) OnTaskSourceBounding(ctx context.Context, task *entity.DMTask, rt *dmTaskRuntime) error {
	logger := utils.GetWorkerLoggerFromContext(ctx).With(
		zap.String("task_id", task.ID),
		zap.String("task_name", task.Name),
		zap.String("stage", string(task.Stage)))
	logger.Info("OnTaskSourceBounding", zap.Int("total_sources", len(task.DMSources)))
	allBounded := true
	namespace, service := constant.DMProxyTarget()
	proxy, bizErr := w.clusterProxyProviderFacade.GetServerlessHTTPProxy(ctx, rt.clusterProxyURL, namespace, service)
	if bizErr != nil {
		logger.Error("failed to get dm proxy from cluster", zap.Error(bizErr))
		return bizErr
	}

	for i := range task.DMSources {
		s := &task.DMSources[i]
		if s.Stage == constant.DMSourceStageBounded {
			continue
		}

		source, err := proxy.GetServerlessDMSource(ctx, s.ID)
		if err != nil {
			logger.Error("failed to get dm source from proxy", zap.Error(err))
			return err
		}
		if len(source.StatusList) > 0 && source.StatusList[0].WorkerName != "" {
			err = w.dmTaskRepo.UpdateSource(ctx, s.ID, map[string]any{
				"stage":        constant.DMSourceStageBounded,
				"conn_profile": s.ConnProfile,
			})
			if err != nil {
				logger.Error("failed to update dm source stage to bounded", zap.Error(err))
				return err
			}
			continue
		}
		allBounded = false
	}

	if allBounded {
		err := w.dmTaskRepo.UpdateTask(
			ctx,
			task.ID,
			[]constant.DMTaskStage{constant.DMTaskStageSourceBounding},
			map[string]any{
				"stage": constant.DMTaskStageTaskCreating,
			},
		)
		if err != nil {
			logger.Error("failed to update dm task stage to taskCreating", zap.Error(err))
			return err
		}
	}

	return nil
}

// OnTaskCreating creates the actual DM task in the cluster after all sources are bound.
// It configures the migration rules, source configurations, and target database settings.
func (w *DMTaskWorker) OnTaskCreating(ctx context.Context, task *entity.DMTask, rt *dmTaskRuntime) error {
	logger := utils.GetWorkerLoggerFromContext(ctx).With(
		zap.String("task_id", task.ID),
		zap.String("task_name", task.Name),
		zap.String("stage", string(task.Stage)),
	)
	logger.Info("OnTaskCreating", zap.Int("total_sources", len(task.DMSources)))
	namespace, service := constant.DMProxyTarget()
	proxy, bizErr := w.clusterProxyProviderFacade.GetServerlessHTTPProxy(ctx, rt.clusterProxyURL, namespace, service)
	if bizErr != nil {
		logger.Error("failed to get dm proxy from cluster", zap.Error(bizErr))
		return bizErr
	}

	password, err := w.encryptorFacade.DecryptBase64(task.TargetEncryptedPassword)
	if err != nil {
		logger.Error("failed to decrypt target password", zap.Error(err))
		return err
	}

	folderName := fmt.Sprintf("%s/%s/", task.TenantID, task.ID)
	bucket := fmt.Sprintf("%s-%s-dm", config.Stage(), task.Region)
	uri := aws.String(fmt.Sprintf("s3://%s/%s", bucket, folderName))

	// Modify URI for AliCloud OSS if needed
	uri, err = w.modifyURIForAliCloudOSS(task, uri, logger)
	if err != nil {
		logger.Error("failed to modify uri", zap.Error(err))
		return err
	}

	dmConf := config.Global().DMConfig
	req := &dm.CreateTaskRequest{
		Task: dm.Task{
			Name:                      task.ID,
			OnDuplicate:               dm.TASKONDUPLICATE_REPLACE,
			EnhanceOnlineSchemaChange: true,
			// do precheck using separate api before creating task
			IgnoreCheckingItems: []string{"all"},
			SourceConfig: dm.TaskSourceConfig{
				FullMigrateConf: &dm.TaskFullMigrateConf{
					Consistency:        aws.String("auto"),
					DataDir:            uri,
					ImportMode:         pointer.To(dm.TASKFULLMIGRATECONFIMPORTMODE_LOGICAL),
					OnDuplicateLogical: pointer.To(dm.TASKFULLMIGRATECONFONDUPLICATELOGICAL_REPLACE),
					ImportThreads:      pointer.To(dmConf.ImportThreads),
					ExportThreads:      pointer.To(dmConf.ExportThreads),
				},
				IncrMigrateConf: &dm.TaskIncrMigrateConf{
					ReplThreads: pointer.To(dmConf.ReplThreads),
					ReplBatch:   pointer.To(dmConf.ReplBatch),
				},
			},
			TargetConfig: dm.TaskTargetDataBase{
				Host:     constant.DMTaskTargetHost,
				Password: password,
				Port:     constant.DMTaskTargetPort,
				User:     task.TargetUser,
			},
			TaskMode: buildTaskMode(task.Mode),
		},
	}
	if task.ShardMode != nil {
		req.Task.ShardMode = task.ShardMode
	}

	// Configure source configurations and migration rules
	sourceConfs := make([]dm.TaskSourceConf, 0, len(task.DMSources))
	migrateRules := make([]dm.TaskTableMigrateRule, 0, len(task.DMSources))
	for i := range task.DMSources {
		s := &task.DMSources[i]
		sourceConfs = append(sourceConfs, dm.TaskSourceConf{
			BinlogGtid: aws.String(s.BinlogGTID),
			BinlogName: aws.String(s.BinlogName),
			BinlogPos:  aws.Int32(s.BinlogPos),
			SourceName: s.ID,
		})
		if len(s.MigrationRules) > 0 {
			migrateRules = append(migrateRules, buildTaskMigrateRulesFromMigrationRules(s)...)
		} else {
			migrateRules = append(migrateRules, buildTaskMigrateRulesFromBlockAllowRules(s)...)
		}
	}
	req.Task.SourceConfig.SourceConf = sourceConfs
	req.Task.TableMigrateRule = migrateRules

	_, err = proxy.CreateServerlessDMTask(ctx, req)
	var openapiErr *tidbcloud.DMOpenapiError
	if errors.As(err, &openapiErr) && openapiErr.ErrorCode == constant.CodeSchedulerSubTaskExist {
		logger.Warn("dm task already created in serverless cluster", zap.Error(err))
	} else if err != nil {
		logger.Error("failed to create dm task", zap.Error(err))
		return err
	}

	sources := make([]string, 0, len(task.DMSources))
	for i := range task.DMSources {
		sources = append(sources, task.DMSources[i].ID)
	}
	err = proxy.StartServerlessDMTask(ctx, task.ID, &dm.StartTaskRequest{
		SourceNameList: sources,
	})
	if err != nil {
		logger.Error("failed to start dm task", zap.Error(err))
		return err
	}

	err = w.dmTaskRepo.UpdateTask(
		ctx,
		task.ID,
		[]constant.DMTaskStage{constant.DMTaskStageTaskCreating},
		map[string]any{
			"stage": constant.DMTaskStageRunning,
		},
	)
	if err != nil {
		logger.Error("failed to update dm task stage to running", zap.Error(err))
		return err
	}

	return nil
}

// modifyURIForAliCloudOSS modifies the S3 URI for AliCloud OSS compatibility.
// Since DM doesn't support OSS ARN authentication yet, we need to pass access keys
// through query parameters for AliCloud OSS access.
func (w *DMTaskWorker) modifyURIForAliCloudOSS(task *entity.DMTask, uri *string, logger *zap.Logger) (*string, error) {
	if task.CloudProvider != config.ProviderAliCloud {
		return uri, nil
	}

	sourceUrl, err := url.Parse(*uri)
	if err != nil {
		logger.Error("failed to parse source url", zap.Error(err), zap.String("uri", *uri))
		return nil, err
	}

	query := sourceUrl.Query()
	query.Set("provider", "alibaba")
	query.Set("region", task.Region)
	query.Set("endpoint", cloud_storage.GenerateOssInternalEndpoint(task.Region))

	// TODO: Replace with OSS ARN authentication when DM supports it
	query.Set("secret-access-key", config.Global().OssConfig.SecretAccessKey)
	query.Set("access-key", config.Global().OssConfig.AccessKeyID)
	sourceUrl.RawQuery = query.Encode()

	return aws.String(sourceUrl.String()), nil
}

// OnTaskDeleting handles the deletion of DM tasks and their associated sources.
// It gracefully handles cases where resources may have already been deleted.
func (w *DMTaskWorker) OnTaskDeleting(ctx context.Context, task *entity.DMTask, rt *dmTaskRuntime) error {
	logger := utils.GetWorkerLoggerFromContext(ctx).With(
		zap.String("task_id", task.ID),
		zap.String("task_name", task.Name),
		zap.String("stage", string(task.Stage)),
	)
	logger.Info("OnTaskDeleting")
	namespace, service := constant.DMProxyTarget()
	proxy, bizErr := w.clusterProxyProviderFacade.GetServerlessHTTPProxy(ctx, rt.clusterProxyURL, namespace, service)
	if bizErr != nil {
		logger.Error("get serverless proxy failed", zap.Error(bizErr))
		return bizErr
	}

	// build source list
	sourceNames := make([]string, 0, len(task.DMSources))
	for i := range task.DMSources {
		sourceNames = append(sourceNames, task.DMSources[i].ID)
	}

	err := proxy.StopServerlessDMTask(ctx, task.ID, &dm.StopTaskRequest{
		SourceNameList: sourceNames,
	})
	var openapiErr *tidbcloud.DMOpenapiError
	if errors.As(err, &openapiErr) && openapiErr.ErrorCode == constant.CodeSchedulerSubTaskOpTaskNotExist {
		logger.Warn("dm task not found in serverless cluster, may be already delete or not started", zap.Error(err))
	} else if err != nil {
		logger.Error("delete dm task failed", zap.Error(err))
		return err
	}

	// The delete task removes metadata information from the downstream database.
	// If the database connection fails (for example, due to a changed user password), an error will be reported.
	// Workaround: When deleting the DM task, if a MySQL access denied error occurs (e.g., due to changed credentials),
	// we ignore the error. This is because the metadata cleanup cannot be completed, but the task can still be considered deleted.
	// This allows force deletion to proceed even if the downstream database is inaccessible.
	err = proxy.DeleteServerlessDMTask(ctx, task.ID)
	if errors.As(err, &openapiErr) && openapiErr.ErrorCode == constant.CodeSchedulerTaskNotExist {
		logger.Warn("dm task not found in serverless cluster, may be already deleted or not started", zap.Error(err))
	} else if err != nil {
		// ignore mysql access denied error
		if !strings.Contains(err.Error(), MySQLAccessDeniedErrorString) {
			logger.Error("delete dm task failed", zap.Error(err))
			return err
		}
	}

	for i := range task.DMSources {
		logger := logger.With(zap.String("dm_source_id", task.DMSources[i].ID))
		if task.DMSources[i].Stage == constant.DMSourceStageDeleted {
			continue
		}
		err = proxy.DeleteServerlessDMSource(ctx, task.DMSources[i].ID)
		if errors.As(err, &openapiErr) && openapiErr.ErrorCode == constant.CodeSchedulerSourceCfgNotExist {
			logger.Warn(
				"dm source not found in serverless cluster, may be already deleted or not created",
				zap.Error(err),
			)
		} else if err != nil {
			logger.Error("delete dm source failed", zap.Error(err), zap.String("dm_source_id", task.DMSources[i].ID))
			return err
		}

		if strings.HasPrefix(task.DMSources[i].EndpointID, constant.PrivateLinkConnectionIDPrefix) {
			err := w.privateLinkConnectionSvc.DeregisterPrivateLinkConnection(
				ctx,
				task.DMSources[i].ClusterID,
				task.DMSources[i].EndpointID,
				task.DMSources[i].ID)
			if err != nil {
				logger.Error("failed to deregister private link connection", zap.Error(err))
				return err
			}
		}

		task.DMSources[i].ConnProfile.Password = ""
		task.DMSources[i].ConnProfile.Security = nil
		err = w.dmTaskRepo.UpdateSource(ctx, task.DMSources[i].ID, map[string]any{
			"stage":        constant.DMSourceStageDeleted,
			"deleted_time": time.Now(),
			"conn_profile": task.DMSources[i].ConnProfile,
		})
		if err != nil {
			logger.Error(
				"failed to update dm source stage to deleted",
				zap.Error(err),
				zap.String("dm_source_id", task.DMSources[i].ID),
			)
			return err
		}
		logger.Info("successfully deleted dm source", zap.String("source_id", task.DMSources[i].ID))
	}

	if task.ClusterStatusID != 0 {
		uint64ClusterId, err := strconv.ParseUint(task.ClusterID, 10, 64)
		if err != nil {
			logger.Error("failed to parse cluster id", zap.Error(err))
			return err
		}

		err = w.mgmtSvc.DeleteClusterCustomStatus(ctx, uint64ClusterId, task.ClusterStatusID)
		if err != nil {
			logger.Error("failed to delete cluster custom status for dm task", zap.Error(err))
			return err
		}
	}

	err = w.dmTaskRepo.UpdateTask(ctx, task.ID, constant.NonDeletedDMTaskStages, map[string]any{
		"stage":                     constant.DMTaskStageDeleted,
		"deleted_time":              time.Now(),
		"target_encrypted_password": "",
	})
	if err != nil {
		logger.Error("failed to update dm task stage to deleted", zap.Error(err))
		return err
	}

	logger.Info("successfully deleted dm task and all sources")
	return nil
}

func (w *DMTaskWorker) OnTaskPolling(ctx context.Context, task *entity.DMTask, rt *dmTaskRuntime) error {
	logger := utils.GetWorkerLoggerFromContext(ctx).With(
		zap.String("task_id", task.ID),
		zap.String("task_name", task.Name),
		zap.String("stage", string(task.Stage)),
	)
	logger.Debug("OnTaskPolling: querying dm task status")

	namespace, service := constant.DMProxyTarget()
	proxy, bizErr := w.clusterProxyProviderFacade.GetServerlessHTTPProxy(ctx, rt.clusterProxyURL, namespace, service)
	if bizErr != nil {
		logger.Error("get serverless proxy failed", zap.Error(bizErr))
		return bizErr
	}

	status, err := proxy.QueryServerlessDMTaskStatus(ctx, task.ID, []string{})
	var openapiErr *tidbcloud.DMOpenapiError
	if errors.As(err, &openapiErr) && openapiErr.ErrorCode == constant.CodeSchedulerTaskNotExist {
		logger.Warn("dm task not found in serverless cluster, may be already deleted or not started", zap.Error(err))
	} else if err != nil {
		logger.Error("query dm task status failed", zap.Error(err))
		return bizerr.New(errorcode.ErrCodeInternalError, "")
	}

	// Build map sourceName -> *dm.SubTaskStatus from API
	currentStatus := make(map[string]*dm.SubTaskStatus, len(status.Data))
	for _, st := range status.Data {
		// we will fetch/merge existing stored status per source below
		currentStatus[st.SourceName] = &st
	}

	stage := constant.DMTaskStageRunning
	var errMsgs []string
	for i := range task.DMSources {
		s := &task.DMSources[i]
		st, ok := currentStatus[s.ID]
		if !ok {
			continue
		}

		// Determine current unit and only update corresponding sub-structure
		unit := strings.ToLower(st.Unit)
		old := s.SubTaskStatus // existing persisted (zero-value if empty)

		// Preserve previously stored sections if current unit is different and API didn't provide that section
		if old != nil {
			if unit != "dump" && old.DumpStatus != nil {
				st.DumpStatus = old.DumpStatus
			}
			if unit != "load" && old.LoadStatus != nil {
				st.LoadStatus = old.LoadStatus
			}
			if unit != "sync" && old.SyncStatus != nil {
				st.SyncStatus = old.SyncStatus
			}
		}

		// Only persist if something actually changed (simple heuristic: compare unit or stage or presence of section)
		str, err := json.Marshal(st)
		if err != nil {
			return fmt.Errorf("json marshal sub_task_status for source %s failed: %w", s.ID, err)
		}
		updates := map[string]any{"sub_task_status": str}
		if err := w.dmTaskRepo.UpdateSource(ctx, s.ID, updates); err != nil {
			logger.Error("update sub_task_status failed", zap.String("source_id", s.ID), zap.Error(err))
			return err
		}
		if st.Stage == dm.TASKSTAGE_STOPPED || st.Stage == dm.TASKSTAGE_PAUSED {
			isFailed := st.ErrorMsg != nil && *st.ErrorMsg != ""
			if isFailed {
				stage = constant.DMTaskStageFailed
				errMsgs = append(errMsgs, *st.ErrorMsg)
				continue
			}

			if stage != constant.DMTaskStageFailed {
				stage = constant.DMTaskStageStopped
			}
		}
	}

	updates := map[string]any{
		"stage": stage,
	}

	if len(errMsgs) > 0 {
		errMsg := strings.Join(errMsgs, "\n")
		updates["err_msg"] = errMsg
		if task.Stage != constant.DMTaskStageFailed || task.ErrMsg != errMsg {
			updates["failed_time"] = time.Now()
		}
	} else {
		updates["err_msg"] = ""
		updates["failed_time"] = nil
	}

	err = w.dmTaskRepo.UpdateTask(ctx, task.ID, constant.ActiveDMTaskStages, updates)
	if err != nil {
		logger.Error("failed to update dm task stage during polling", zap.Error(err))
		return err
	}

	return nil
}

// buildTaskMigrateRulesFromBlockAllowRules generates DM TaskTableMigrateRule(s) from BlockAllowRules for a given DMSource.
// For DM, BlockAllowRules must be set first. If needed, MigrateRule can be configured additionally.
// The DM OpenAPI does not expose BlockAllowRules directly; instead, BlockAllowRules are generated from MigrateRule.
// If you only configure table_migrate_rule like:
//
//	[
//	  {
//	    "source": {"schema": "db", "source_name": "mysql-01", "table": "tb"},
//	    "target": {"schema": "db", "table": "tb"}
//	  }
//	]
//
// The generated BlockAllowRules will be: do-dbs: [], do-tables: ["db.tb"].
// This means all databases are included by default, but only the specified tables (db.tb) are allowed.
// However, this does not meet the requirement, because all databases will be synchronized except for the tables listed in do-tables.
// To ensure only the intended tables are synchronized, you must also add a rule like:
//
//	{
//	  "source": {"schema": "db", "source_name": "mysql-01"},
//	  "target": {"schema": "db"}
//	}
//
// This ensures the correct BlockAllowRules are generated, so only the specified tables in the specified database are synchronized.
func buildTaskMigrateRulesFromBlockAllowRules(s *entity.DMSource) []dm.TaskTableMigrateRule {
	taskMigrateRules := make([]dm.TaskTableMigrateRule, 0)
	schemaExists := make(map[string]struct{})
	for _, schema := range s.BlockAllowRules.DoDBs {
		migrateRule := dm.TaskTableMigrateRule{}
		migrateRule.Source.Schema = schema
		migrateRule.Source.SourceName = s.ID
		migrateRule.Target = &dm.TaskTableMigrateRuleTarget{
			Schema: pointer.To(schema),
		}
		taskMigrateRules = append(taskMigrateRules, migrateRule)
		schemaExists[schema] = struct{}{}
	}
	for _, table := range s.BlockAllowRules.DoTables {
		migrateRule := dm.TaskTableMigrateRule{}
		migrateRule.Source.SourceName = s.ID
		migrateRule.Source.Schema = table.Schema
		migrateRule.Source.Table = table.Table
		migrateRule.Target = &dm.TaskTableMigrateRuleTarget{
			Schema: pointer.To(table.Schema),
		}
		if migrateRule.Source.Table != "*" && migrateRule.Source.Table != "~.*" {
			migrateRule.Target.Table = pointer.To(table.Table)
		}
		taskMigrateRules = append(taskMigrateRules, migrateRule)
		if _, ok := schemaExists[table.Schema]; !ok {
			migrateRule := dm.TaskTableMigrateRule{}
			migrateRule.Source.Schema = table.Schema
			migrateRule.Source.SourceName = s.ID
			migrateRule.Target = &dm.TaskTableMigrateRuleTarget{
				Schema: pointer.To(table.Schema),
			}
			taskMigrateRules = append(taskMigrateRules, migrateRule)
			schemaExists[table.Schema] = struct{}{}
		}
	}
	return taskMigrateRules
}

// buildTaskMigrateRulesFromMigrationRules translates MigrationRules into DM TaskTableMigrateRule(s) for a given DMSource.
// Each MigrationRule already specifies how a source schema/table should be rewritten, so we convert it into a table-level migrate rule.
// Similar to BlockAllowRules handling, we also add a schema-level migrate rule per unique source schema to ensure DM generates the correct
// block-allow lists and only migrates the intended schemas/tables.
func buildTaskMigrateRulesFromMigrationRules(s *entity.DMSource) []dm.TaskTableMigrateRule {
	taskMigrateRules := make([]dm.TaskTableMigrateRule, 0)
	schemaExists := make(map[string]struct{})
	for _, rule := range s.MigrationRules {
		migrateRule := dm.TaskTableMigrateRule{}
		migrateRule.Source.SourceName = s.ID
		migrateRule.Source.Schema = rule.Source.SchemaPattern
		migrateRule.Source.Table = rule.Source.TablePattern
		migrateRule.Target = &dm.TaskTableMigrateRuleTarget{
			Schema: pointer.To(rule.Target.Schema),
			Table:  pointer.To(rule.Target.Table),
		}
		taskMigrateRules = append(taskMigrateRules, migrateRule)

		if _, ok := schemaExists[rule.Source.SchemaPattern]; !ok {
			migrateRule := dm.TaskTableMigrateRule{}
			migrateRule.Source.Schema = rule.Source.SchemaPattern
			migrateRule.Source.SourceName = s.ID
			migrateRule.Target = &dm.TaskTableMigrateRuleTarget{
				Schema: pointer.To(rule.Target.Schema),
			}
			taskMigrateRules = append(taskMigrateRules, migrateRule)
			schemaExists[rule.Source.SchemaPattern] = struct{}{}
		}
	}

	return taskMigrateRules
}

// buildTaskMode converts the entity.DMTaskMode to the corresponding dm.TaskMode for the DM API.
func buildTaskMode(mode constant.DMTaskMode) dm.TaskTaskMode {
	switch mode {
	case constant.DMTaskModeAll:
		return dm.TASKTASKMODE_ALL
	case constant.DMTaskModeIncr:
		return dm.TASKTASKMODE_INCREMENTAL
	default:
		return ""
	}
}