package pimport

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

	"github.com/avast/retry-go/v4"
	pingErr "github.com/pingcap/errors"
	"github.com/pingcap/log"
	"github.com/tidbcloud/dataflow-service-ng/app/event_bus/publisher/center"
	"github.com/tidbcloud/dataflow-service-ng/app/internal/cluster_job"
	"github.com/tidbcloud/dataflow-service-ng/app/internal/feature_gate"
	"github.com/tidbcloud/dataflow-service-ng/app/repo/entity"
	"github.com/tidbcloud/dataflow-service-ng/app/utils"
	importadminpb "github.com/tidbcloud/dataflow-service-ng/idl/pbgen/proto/import_into/v1beta2"
	"github.com/tidbcloud/dataflow-service-ng/internal/config"
	serviceconfig "github.com/tidbcloud/dataflow-service-ng/internal/config"
	"github.com/tidbcloud/dataflow-service-ng/internal/db"
	"github.com/tidbcloud/dataflow-service-ng/pkg/agent"
	pkgcrypto "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"
	agentpb "github.com/tidbcloud/dataflow-service/dataflow-agent-ng/idl/pbgen/proto/ng_agent"
	nextgenv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/nextgen/v1"
	mgmt "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"google.golang.org/grpc/codes"
	"google.golang.org/protobuf/encoding/protojson"
	"gorm.io/gorm"
)

const (
	JobIDColName                = "Job_ID"
	GroupKeyColName             = "Group_Key"
	TargetTableColName          = "Target_Table"
	PhaseColName                = "Phase"
	StatusColName               = "Status"
	ImportRowsColName           = "Imported_Rows"
	ResultMessageColName        = "Result_Message"
	CreateTime                  = "Create_Time"
	StartTime                   = "Start_Time"
	EndTimeColName              = "End_Time"
	CurStepColName              = "Cur_Step"
	CurStepProcessedSizeColName = "Cur_Step_Processed_Size"
	CurStepTotalSizeColName     = "Cur_Step_Total_Size"
	CurStepProgressPctColName   = "Cur_Step_Progress_Pct"
	CurStepSpeedColName         = "Cur_Step_Speed"
	CurStepETAColName           = "Cur_Step_ETA"
)

const (
	DetachedOptionName = "DETACHED"
	GroupKeyOptionName = "GROUP_KEY"

	SeparatorOptionName       = "FIELDS_TERMINATED_BY"
	DelimiterOptionName       = "FIELDS_ENCLOSED_BY"
	EscapeOptionName          = "FIELDS_ESCAPED_BY"
	NullOptionName            = "FIELDS_DEFINED_NULL_BY"
	LinesTerminatorOptionName = "LINES_TERMINATED_BY"
	SkipRowSOptionName        = "SKIP_ROWS"
	SplitFileOptionName       = "SPLIT_FILE"
)

const (
	importIntoEstimateTaskIDPrefix     = "import-into-estimate"
	importIntoGetFilesTaskIDPrefix     = "import-into-get-files"
	importIntoCreateSchemaTaskIDPrefix = "import-into-create-schema"
	importIntoSubmitTaskIDPrefix       = "import-into-submit"
	importIntoSubmitTaskPendingPhase   = "submit-task-pending"
)

const (
	importIntoTaskSuccessMessage     = "All table(s) completed to importing successfully"
	importClusterJobCleanupTimeout   = 10 * time.Second
	importClusterJobReconcileTimeout = 10 * time.Second
	importClusterJobRetryBackoff     = time.Minute
)

type ImportIntoModel struct {
	log                      *zap.Logger
	MgmtSvcGrpcClient        grpcli.MgmtSvcGrpcClient
	encryptor                encryptor_wrapper.Encryptor
	ngRegionalClientProvider agent.NGRegionalClientProvider
	featureGate              feature_gate.FeatureGate
	eventPublisher           center.EventCenterPublisher
	clusterJob               cluster_job.Model
}

func NewImportIntoModel(
	mgmtCli grpcli.MgmtSvcGrpcClient, encryptor encryptor_wrapper.Encryptor,
	ngRegionalClientProvider agent.NGRegionalClientProvider,
	featureGate feature_gate.FeatureGate,
	eventPublisher center.EventCenterPublisher,
	clusterJob cluster_job.Model,
) *ImportIntoModel {
	return &ImportIntoModel{
		log:                      zap.NewNop(),
		MgmtSvcGrpcClient:        mgmtCli,
		encryptor:                encryptor,
		ngRegionalClientProvider: ngRegionalClientProvider,
		featureGate:              featureGate,
		eventPublisher:           eventPublisher,
		clusterJob:               clusterJob,
	}
}

func (c *ImportIntoModel) traceIDLogger(ctx context.Context, fields ...zap.Field) *zap.Logger {
	return c.log.With(tilog.TraceIDFromContext(ctx)...).With(fields...)
}

func (c *ImportIntoModel) sugaredTraceIDLogger(ctx context.Context, fields ...zap.Field) *zap.SugaredLogger {
	return c.traceIDLogger(ctx, fields...).Sugar()
}

func (c *ImportIntoModel) newNGRegionalClientByClusterID(ctx context.Context, clusterID uint64) (*agent.NGRegionalClient, error) {
	return c.ngRegionalClientProvider.NewByClusterID(ctx, clusterID)
}

func (c *ImportIntoModel) shouldProcessImportIntoTaskInWorkerScope(task *entity.ImportIntoTask) bool {
	if task == nil {
		return false
	}
	if _, ok := serviceconfig.Global().ImportConfig.WorkerEnabled.Enabled(
		config.FeatureGateQueryParams{
			Provider:  "",
			Region:    "",
			TenantID:  task.TenantID.Uint64(),
			ProjectID: task.ProjectID.Uint64(),
			ClusterID: task.ClusterID.Uint64(),
		},
	); ok {
		return true
	}
	return false
}

func FormatStringOptionValue(s string) string {
	switch s {
	case "\\'", "'":
		return "'\\''"
	case "\r\n", "\\r\\n":
		return "'\\r\\n'"
	case "\\r", "\r":
		return "'\\r'"
	case "\\n", "\n":
		return "'\\n'"
	case "\\t", "\t":
		return "'\\t'"
	case "\\", "\\\\":
		return "'\\\\'"
	case "\"":
		return "'\"'"
	case "\\N":
		return "'\\\\N'"
	default:
		var builder strings.Builder
		for _, r := range s {
			switch r {
			case '\'':
				builder.WriteString(`\'`)
			case '\\':
				builder.WriteString(`\\`)
			case '\n':
				builder.WriteString(`\n`)
			case '\r':
				builder.WriteString(`\r`)
			case '\x00': // NUL byte
				builder.WriteString(`\0`)
			case '\x1a': // Ctrl+Z
				builder.WriteString(`\Z`)
			case '`':
				continue
			default:
				builder.WriteRune(r)
			}
		}
		return fmt.Sprintf("'%s'", builder.String())
	}
}

func (c *ImportIntoModel) Run(ctx context.Context) error {
	logger := c.sugaredTraceIDLogger(ctx)
	reconcileCtx, reconcileCancel := context.WithTimeout(ctx, importClusterJobReconcileTimeout)
	if err := c.ensureActiveImportClusterJobs(reconcileCtx); err != nil {
		logger.With(zap.Error(err)).Error("register active Import Into custom statuses error")
	}
	if reconcileCtx.Err() == nil {
		if err := c.cleanupTerminalImportClusterJobs(reconcileCtx); err != nil {
			logger.With(zap.Error(err)).Error("clean up terminal Import Into custom statuses error")
		}
	}
	reconcileCancel()
	statusHandlerMap := map[entity.ImportIntoStatus]func(context.Context, *entity.ImportIntoTask) error{
		entity.ImportIntoStatusPreparing: c.onImportIntoPreparing,
		entity.ImportIntoStatusScalingUp: c.onImportIntoScalingUp,
		entity.ImportIntoStatusCanceling: c.onImportIntoCanceling,
		entity.ImportIntoStatusImporting: c.onImportIntoImporting,
	}
	for status, handler := range statusHandlerMap {
		if err := c.processImportIntoTasksByStatus(ctx, status, handler); err != nil {
			logger.With(zap.Error(err)).Error("process import into tasks by status error")
		}
	}
	return nil
}

// ensureActiveImportClusterJobs deliberately ignores WorkerEnabled scope. Task processing is scoped,
// but reconciliation repairs cluster-level mutex state and must remain globally available.
func (c *ImportIntoModel) ensureActiveImportClusterJobs(ctx context.Context) error {
	var tasks []entity.ImportIntoTask
	if err := db.TransactionWithContext(ctx, func(tx *gorm.DB) error {
		return tx.Where("custom_status_registered = ? AND status IN ?", false, []entity.ImportIntoStatus{
			entity.ImportIntoStatusPreparing,
			entity.ImportIntoStatusScalingUp,
			entity.ImportIntoStatusImporting,
			entity.ImportIntoStatusCanceling,
		}).Where("custom_status_retry_at IS NULL OR custom_status_retry_at <= ?", time.Now()).
			Order("custom_status_retry_at, id").
			Limit(serviceconfig.Global().ImportConfig.MaxImportTasksRetrievedOnWorker).Find(&tasks).Error
	}); err != nil {
		return err
	}

	var registerErr error
	for i := range tasks {
		if ctxErr := ctx.Err(); ctxErr != nil {
			return errors.Join(registerErr, ctxErr)
		}
		claimed, err := c.claimImportClusterJobReconciliation(ctx, &tasks[i])
		if err != nil {
			registerErr = errors.Join(registerErr, err)
			continue
		}
		if !claimed {
			continue
		}
		logger := c.sugaredTraceIDLogger(ctx,
			zap.Uint64("importID", tasks[i].ID.Uint64()),
			zap.Uint64("clusterID", tasks[i].ClusterID.Uint64()),
		)
		if err := c.ensureImportClusterJob(ctx, logger, &tasks[i]); err != nil {
			logger.With(zap.Error(err)).Error("failed to register Import Into custom status")
			registerErr = errors.Join(registerErr, err)
		}
	}
	return registerErr
}

func (c *ImportIntoModel) ensureImportClusterJob(
	ctx context.Context,
	logger *zap.SugaredLogger,
	imp *entity.ImportIntoTask,
) error {
	if imp.CustomStatusRegistered {
		return nil
	}
	if err := c.clusterJob.EnsureImportClusterJob(ctx, imp.ClusterID.Uint64(), imp.ID.Uint64()); err != nil {
		return errors.Join(err, c.deferImportClusterJobReconciliation(ctx, imp.ID))
	}
	if err := db.TransactionWithContext(context.WithoutCancel(ctx), func(tx *gorm.DB) error {
		return tx.Model(&entity.ImportIntoTask{}).
			Where("id = ?", imp.ID).
			Updates(map[string]any{
				"custom_status_registered": true,
				"custom_status_retry_at":   nil,
			}).Error
	}); err != nil {
		return err
	}
	imp.CustomStatusRegistered = true
	logger.Info("registered Import Into custom status")
	return nil
}

// cleanupTerminalImportClusterJobs deliberately ignores WorkerEnabled scope for the same reason as
// ensureActiveImportClusterJobs: terminal mutex cleanup must survive worker-scope changes.
func (c *ImportIntoModel) cleanupTerminalImportClusterJobs(ctx context.Context) error {
	var tasks []entity.ImportIntoTask
	if err := db.TransactionWithContext(ctx, func(tx *gorm.DB) error {
		return tx.Where("custom_status_registered = ? AND status IN ?", true, []entity.ImportIntoStatus{
			entity.ImportIntoStatusCompleted,
			entity.ImportIntoStatusFailed,
			entity.ImportIntoStatusCanceled,
		}).Where("custom_status_retry_at IS NULL OR custom_status_retry_at <= ?", time.Now()).
			Order("custom_status_retry_at, id").
			Limit(serviceconfig.Global().ImportConfig.MaxImportTasksRetrievedOnWorker).Find(&tasks).Error
	}); err != nil {
		return err
	}

	var cleanupErr error
	for i := range tasks {
		if ctxErr := ctx.Err(); ctxErr != nil {
			return errors.Join(cleanupErr, ctxErr)
		}
		claimed, err := c.claimImportClusterJobReconciliation(ctx, &tasks[i])
		if err != nil {
			cleanupErr = errors.Join(cleanupErr, err)
			continue
		}
		if !claimed {
			continue
		}
		logger := c.sugaredTraceIDLogger(ctx,
			zap.Uint64("importID", tasks[i].ID.Uint64()),
			zap.Uint64("clusterID", tasks[i].ClusterID.Uint64()),
		)
		if err := c.cleanupImportClusterJob(ctx, logger, &tasks[i]); err != nil {
			logger.With(zap.Error(err)).Error("failed to clean up Import Into custom status")
			cleanupErr = errors.Join(cleanupErr, err)
		}
	}
	return cleanupErr
}

func (c *ImportIntoModel) cleanupImportClusterJob(
	ctx context.Context,
	logger *zap.SugaredLogger,
	imp *entity.ImportIntoTask,
) error {
	if !imp.CustomStatusRegistered {
		return nil
	}
	cleanupCtx, cleanupCancel := context.WithTimeout(ctx, importClusterJobCleanupTimeout)
	defer cleanupCancel()
	if err := c.clusterJob.DeregisterImportClusterJob(
		cleanupCtx, imp.ClusterID.Uint64(), imp.ID.Uint64(),
	); err != nil {
		return errors.Join(err, c.deferImportClusterJobReconciliation(ctx, imp.ID))
	}
	if err := db.TransactionWithContext(context.WithoutCancel(ctx), func(tx *gorm.DB) error {
		return tx.Model(&entity.ImportIntoTask{}).
			Where("id = ?", imp.ID).
			Updates(map[string]any{
				"custom_status_registered": false,
				"custom_status_retry_at":   nil,
			}).Error
	}); err != nil {
		return err
	}
	imp.CustomStatusRegistered = false
	logger.Info("cleaned up Import Into custom status")
	return nil
}

func (c *ImportIntoModel) claimImportClusterJobReconciliation(
	ctx context.Context,
	imp *entity.ImportIntoTask,
) (bool, error) {
	now := time.Now()
	claimed := false
	err := db.TransactionWithContext(ctx, func(tx *gorm.DB) error {
		result := tx.Model(&entity.ImportIntoTask{}).
			Where("id = ? AND (custom_status_retry_at IS NULL OR custom_status_retry_at <= ?)", imp.ID, now).
			UpdateColumn("custom_status_retry_at", now.Add(importClusterJobRetryBackoff))
		claimed = result.RowsAffected > 0
		return result.Error
	})
	if err != nil {
		return false, err
	}
	return claimed, nil
}

func (c *ImportIntoModel) deferImportClusterJobReconciliation(ctx context.Context, importID entity.ID) error {
	return db.TransactionWithContext(context.WithoutCancel(ctx), func(tx *gorm.DB) error {
		return tx.Model(&entity.ImportIntoTask{}).
			Where("id = ?", importID).
			UpdateColumn("custom_status_retry_at", time.Now().Add(importClusterJobRetryBackoff)).Error
	})
}

type ImportIntoTask struct {
	entity.ImportIntoTask
	Cluster *mgmt.Cluster
}

func (c *ImportIntoModel) ListIncompletedImportIntoTasks(ctx context.Context) ([]ImportIntoTask, error) {
	statuses := []entity.ImportIntoStatus{
		entity.ImportIntoStatusPreparing,
		entity.ImportIntoStatusScalingUp,
		entity.ImportIntoStatusImporting,
		entity.ImportIntoStatusCanceling,
	}
	return c.ListImportIntoTasksByProcessor(ctx, func(q *gorm.DB) *gorm.DB {
		return q.Where("status in ?", statuses)
	})
}

func (c *ImportIntoModel) ListRecentlyFinishedTerminalImportIntoTasks(ctx context.Context, t time.Time) ([]ImportIntoTask, error) {
	statuses := []entity.ImportIntoStatus{
		entity.ImportIntoStatusCompleted,
		entity.ImportIntoStatusFailed,
	}
	return c.ListImportIntoTasksByProcessor(ctx, func(q *gorm.DB) *gorm.DB {
		return q.Where("status in ?", statuses).Where("finished_at >= ?", t)
	})
}

func (c *ImportIntoModel) ListImportIntoTasksByProcessor(ctx context.Context, proc func(*gorm.DB) *gorm.DB) ([]ImportIntoTask, error) {
	logger := c.sugaredTraceIDLogger(ctx)
	var entities []entity.ImportIntoTask
	err := db.TransactionWithContext(ctx, func(tx *gorm.DB) error {
		query := tx.Model(entity.ImportIntoTask{})
		query = proc(query)
		err := query.Scan(&entities).Error
		if err != nil {
			return err
		}
		return nil
	})
	if err != nil {
		return nil, err
	}
	if len(entities) == 0 {
		return nil, nil
	}
	var clusterIDs []uint64
	for _, ent := range entities {
		clusterIDs = append(clusterIDs, ent.ClusterID.Uint64())
	}
	listResp, err := c.MgmtSvcGrpcClient.ListClusters(ctx, &mgmt.ListClustersRequest{Filter: &mgmt.ClusterFilter{ClusterIds: clusterIDs}, PageSize: uint32(len(clusterIDs))})
	if err != nil {
		logger.With(zap.Error(err), zap.Uint64s("clusterIDs", clusterIDs)).Error("list mgmt clusters error")
		return nil, err
	}
	clusterMap := make(map[uint64]*mgmt.Cluster, len(listResp.Clusters))
	for _, cluster := range listResp.Clusters {
		clusterMap[cluster.ClusterId] = cluster
	}
	var result []ImportIntoTask
	for _, ent := range entities {
		cluster, ok := clusterMap[ent.ClusterID.Uint64()]
		if !ok {
			logger.With(zap.Uint64("clusterID", ent.ClusterID.Uint64()), zap.Uint64("importTaskID", ent.ID.Uint64())).Warn("missing cluster in clusterMap")
			continue
		}
		result = append(result, ImportIntoTask{
			ImportIntoTask: ent,
			Cluster:        cluster,
		})
	}
	return result, nil
}

func (c *ImportIntoModel) processImportIntoTasksByStatus(ctx context.Context, status entity.ImportIntoStatus, handler func(context.Context, *entity.ImportIntoTask) error) error {
	logger := c.sugaredTraceIDLogger(ctx, zap.String("status", string(status)))
	entities, err := c.ListImportIntoTasksByProcessor(ctx, func(q *gorm.DB) *gorm.DB {
		return q.Where("status = ?", status).Limit(serviceconfig.Global().ImportConfig.MaxImportTasksRetrievedOnWorker)
	})
	if err != nil {
		return err
	}
	for _, ent := range entities {
		if !c.shouldProcessImportIntoTaskInWorkerScope(&ent.ImportIntoTask) {
			logger.With(
				zap.Uint64("taskID", ent.ID.Uint64()),
				zap.Uint64("tenantID", ent.TenantID.Uint64()),
				zap.Uint64("projectID", ent.ProjectID.Uint64()),
				zap.Uint64("clusterID", ent.ClusterID.Uint64()),
			).Info("skip import into task because worker scope is disabled")
			continue
		}
		if ent.Cluster.GetNextGenCluster().GetState() == nextgenv1.Cluster_DELETED {
			prevStatus := ent.Status
			if shouldReleaseImportIntoReservationOnFailureStatus(ent.Status) {
				agentClient, releaseErr := c.newNGRegionalClientByClusterID(ctx, ent.ClusterID.Uint64())
				if releaseErr != nil {
					logger.With(zap.Uint64("id", ent.ID.Uint64()), zap.Uint64("clusterId", ent.ClusterID.Uint64()), zap.Error(releaseErr)).Warn("failed to create agent client for tikv reservation cleanup")
				} else {
					c.releaseTiKVCapacityReservation(ctx, logger, agentClient, ent.TenantID.Uint64(), ent.ProjectID.Uint64(), ent.ClusterID.Uint64(), ent.ID.Uint64())
				}
			}
			err := db.TransactionWithContext(ctx, func(tx *gorm.DB) error {
				ent.Status = entity.ImportIntoStatusFailed
				ent.Message = "Cluster has been deleted"
				ent.ImportResults.ControlPlaneFailed = true
				cleanImportIntoStoredURIs(&ent.ImportIntoTask)
				for i := 0; i != len(ent.ImportResults.SubTasks); i++ {
					switch ent.ImportResults.SubTasks[i].ImportIntoJobStatus {
					case entity.ImportIntoJobStatusFailed, entity.ImportIntoJobStatusFinished, entity.ImportIntoJobStatusCanceled:
						continue
					}
					setImportIntoSubTaskStatus(&ent.ImportResults.SubTasks[i], entity.ImportIntoJobStatusFailed)
					ent.ImportResults.SubTasks[i].Message = ent.Message
				}
				return ent.Update(tx,
					entity.ImportIntoTaskDBSchema.Status,
					entity.ImportIntoTaskDBSchema.Message,
					entity.ImportIntoTaskDBSchema.ImportResults,
					entity.ImportIntoTaskDBSchema.SourceURI,
				)
			})
			if err != nil {
				logger.With(zap.Uint64("id", ent.ID.Uint64()), zap.Uint64("clusterId", ent.ClusterID.Uint64()), zap.Error(err)).Error("failed to mark import into task as failed")
			} else {
				publishImportIntoEventIfStatusChanged(
					ctx, c.eventPublisher, &ent.ImportIntoTask, prevStatus,
				)
				if cleanupErr := c.cleanupImportClusterJob(ctx, logger, &ent.ImportIntoTask); cleanupErr != nil {
					logger.With(zap.Error(cleanupErr)).Error("failed to clean up Import Into custom status")
				}
			}
			logger.With(zap.Uint64("id", ent.ID.Uint64()), zap.Uint64("clusterId", ent.ClusterID.Uint64())).Info("cluster hase been deleted, mark task as failed")
			continue
		}
		err = handler(ctx, &ent.ImportIntoTask)
		if err != nil {
			logger.With(zap.Error(err), zap.Uint64("taskId", ent.ID.Uint64())).Error("import into task failed")
			continue
		}
	}
	return nil
}

func appendExternalIDToSourceURIIfNeeded(sourceURI string, clusterID uint64) (string, error) {
	parsedSourceURL, err := url.Parse(sourceURI)
	if err != nil {
		return "", err
	}
	if parsedSourceURL.Scheme != "s3" && parsedSourceURL.Scheme != "oss" {
		return sourceURI, nil
	}

	query := parsedSourceURL.Query()
	query.Set("external-id", fmt.Sprint(clusterID))
	parsedSourceURL.RawQuery = query.Encode()
	return parsedSourceURL.String(), nil
}

func buildImportIntoEstimateTaskID(importID uint64) string {
	return fmt.Sprintf("%s-%d", importIntoEstimateTaskIDPrefix, importID)
}

func buildImportIntoGetFilesTaskID(importID uint64) string {
	return fmt.Sprintf("%s-%d", importIntoGetFilesTaskIDPrefix, importID)
}

func buildImportIntoCreateSchemaTaskID(importID uint64) string {
	return fmt.Sprintf("%s-%d", importIntoCreateSchemaTaskIDPrefix, importID)
}

func buildImportIntoSubmitTaskID(importID uint64, subTaskIndex int) string {
	return fmt.Sprintf("%s-%d-%d", importIntoSubmitTaskIDPrefix, importID, subTaskIndex)
}

func buildImportIntoTiKVReservationID(importID uint64) string {
	return fmt.Sprintf("%d", importID)
}

func isImportIntoSubTaskFileDiscoveryPrepared(task *entity.ImportIntoSubTask) bool {
	return task.ImportIntoTotalSize > 0 ||
		task.ImportIntoSchemaFile != "" ||
		task.ImportIntoFileType != "" ||
		task.ImportIntoFileCompression != ""
}

func isImportIntoFileDiscoveryPrepared(imp *entity.ImportIntoTask) bool {
	if len(imp.ImportResults.SubTasks) == 0 {
		return false
	}
	for i := range imp.ImportResults.SubTasks {
		if !isImportIntoSubTaskFileDiscoveryPrepared(&imp.ImportResults.SubTasks[i]) {
			return false
		}
	}
	return true
}

func isImportIntoSubTaskSchemaPrepared(task *entity.ImportIntoSubTask) bool {
	return task.ImportIntoJobStatus != ""
}

func isImportIntoSchemaOnlySubTask(task *entity.ImportIntoSubTask) bool {
	return task.ImportIntoSchemaFile != "" &&
		task.ImportIntoFileType == agentpb.GetImportFilesResp_UNKNOWN_TYPE.String()
}

func isImportIntoSchemaPrepared(imp *entity.ImportIntoTask) bool {
	if len(imp.ImportResults.SubTasks) == 0 {
		return false
	}
	for i := range imp.ImportResults.SubTasks {
		if !isImportIntoSubTaskSchemaPrepared(&imp.ImportResults.SubTasks[i]) {
			return false
		}
	}
	return true
}

func getImportIntoSharedPoolID(ctx context.Context, agentClient *agent.NGRegionalClient, clusterID uint64) (string, error) {
	tidbGroup, err := agentClient.GetTiDBGroup(ctx, clusterID)
	if err != nil {
		return "", err
	}
	if tidbGroup.SharedPoolId == "" {
		return "", fmt.Errorf("shared pool id is empty for cluster %d", clusterID)
	}
	return tidbGroup.SharedPoolId, nil
}

func (c *ImportIntoModel) isImportIntoTiKVCapacityReservationEnabled(tenantID, projectID, clusterID uint64) bool {
	return c.featureGate.IsFeatureEnabled(
		&feature_gate.QueryFeatureGateReq{
			TenantID:  tenantID,
			ProjectID: projectID,
			ClusterID: clusterID,
		},
		feature_gate.ImportIntoTiKVReservation,
	)
}

func (c *ImportIntoModel) releaseTiKVCapacityReservation(ctx context.Context, logger *zap.SugaredLogger, agentClient *agent.NGRegionalClient, tenantID, projectID, clusterID, importID uint64) {
	if !c.isImportIntoTiKVCapacityReservationEnabled(tenantID, projectID, clusterID) {
		logger.Debugw("skip tikv capacity reservation release because feature is disabled")
		return
	}
	sharedPoolID, err := getImportIntoSharedPoolID(ctx, agentClient, clusterID)
	if err != nil {
		logger.Warnw("failed to get shared pool id when releasing tikv capacity reservation", "error", err)
		return
	}
	reservationID := buildImportIntoTiKVReservationID(importID)
	if err := agentClient.DeleteTiKVCapacityReservation(ctx, sharedPoolID, reservationID); err != nil {
		logger.Warnw("failed to delete tikv capacity reservation",
			"error", err,
			"sharedPoolID", sharedPoolID,
			"reservationID", reservationID,
		)
		return
	}
	logger.Infow("deleted tikv capacity reservation", "sharedPoolID", sharedPoolID, "reservationID", reservationID)
}

func applyImportIntoFileDiscoveryTotalSize(imp *entity.ImportIntoTask) {
	var totalSourceSize uint64
	for i := range imp.ImportResults.SubTasks {
		if imp.ImportResults.SubTasks[i].ImportIntoJobStatus == entity.ImportIntoJobStatusPending {
			totalSourceSize += imp.ImportResults.SubTasks[i].ImportIntoTotalSize
		}
	}
	if totalSourceSize == 0 {
		applyImportIntoBusinessFailure(imp, "import data is empty")
		return
	}
	imp.TotalSize = totalSourceSize
}

const tikvReplicaCount uint64 = 3

func reservedImportIntoTiKVBytes(estimateTiKVSize uint64) uint64 {
	return estimateTiKVSize * tikvReplicaCount
}

func collectImportIntoGetFilesTableMetas(imp *entity.ImportIntoTask) []*agentpb.GetImportFilesReq_TableMeta {
	metas := make([]*agentpb.GetImportFilesReq_TableMeta, 0, len(imp.ImportResults.SubTasks))
	for i := range imp.ImportResults.SubTasks {
		subTask := &imp.ImportResults.SubTasks[i]
		metas = append(metas, &agentpb.GetImportFilesReq_TableMeta{
			Database:     subTask.DataBase,
			Table:        subTask.TableName,
			WildcardPath: subTask.ImportIntoFilesWildcard,
		})
	}
	return metas
}

func buildImportIntoEstimateCsvConfig(imp *entity.ImportIntoTask) (*agentpb.CsvConfig, error) {
	if imp.CustomCSVFormat == "" {
		return nil, nil
	}
	var csvFormat importadminpb.CSVFormat
	if err := protojson.Unmarshal([]byte(imp.CustomCSVFormat), &csvFormat); err != nil {
		return nil, pingErr.Annotate(err, "failed to unmarshal CustomCSVFormat")
	}
	cfg := &agentpb.CsvConfig{}
	if separator, ok, err := restoreImportIntoEstimateCSVLiteral("separator", csvFormat.GetSeparator()); err != nil {
		return nil, err
	} else if ok {
		cfg.Separator = separator
	}
	if delimiter, ok, err := restoreImportIntoEstimateCSVLiteral("delimiter", csvFormat.GetDelimiter()); err != nil {
		return nil, err
	} else if ok {
		cfg.Delimiter = delimiter
	}
	if escapeChar, ok, err := restoreImportIntoEstimateCSVLiteral("escape_char", csvFormat.GetEscapeChar()); err != nil {
		return nil, err
	} else if ok {
		cfg.EscapeChar = &escapeChar
	}
	if lineTerminator, ok, err := restoreImportIntoEstimateCSVLiteral("line_terminator", csvFormat.GetLineTerminator()); err != nil {
		return nil, err
	} else if ok && !isImportIntoEstimateDefaultLineTerminator(lineTerminator) {
		cfg.LinesTerminatedBy = &lineTerminator
	}
	if csvFormat.GetNullValue() != "" {
		nullValue := csvFormat.GetNullValue()
		cfg.NullValue = &nullValue
	}
	return cfg, nil
}

func restoreImportIntoEstimateCSVLiteral(field, raw string) (string, bool, error) {
	if raw == "" {
		return "", false, nil
	}
	if len(raw) == 1 {
		return raw, true, nil
	}
	restored, err := utils.RestoreSpecialChars(raw)
	if err != nil {
		return "", false, pingErr.Annotatef(err, "invalid %s for import estimate csv config", field)
	}
	return restored, true, nil
}

func isImportIntoEstimateDefaultLineTerminator(lineTerminator string) bool {
	switch lineTerminator {
	case "\r\n", "\r", "\n":
		return true
	default:
		return false
	}
}

func collectImportIntoEstimateTableMetas(imp *entity.ImportIntoTask) []*agentpb.CreateImportDataSizeEstimateTaskReq_TableMeta {
	metas := make([]*agentpb.CreateImportDataSizeEstimateTaskReq_TableMeta, 0, len(imp.ImportResults.SubTasks))
	for i := range imp.ImportResults.SubTasks {
		subTask := &imp.ImportResults.SubTasks[i]
		if subTask.ImportIntoJobStatus != entity.ImportIntoJobStatusPending {
			continue
		}
		metas = append(metas, &agentpb.CreateImportDataSizeEstimateTaskReq_TableMeta{
			Database:     subTask.DataBase,
			Table:        subTask.TableName,
			WildcardPath: subTask.ImportIntoFilesWildcard,
		})
	}
	return metas
}

func populateImportIntoSubTaskFromFileDiscovery(task *entity.ImportIntoSubTask, tableMeta *agentpb.GetImportFilesResp_TableMeta) {
	task.ImportIntoTotalSize = tableMeta.GetTotalSize()
	schemaFile := tableMeta.GetSchemaFile()
	if stripped, err := utils.StripURIQuery(schemaFile); err == nil {
		schemaFile = stripped
	} else {
		log.Warn("failed to strip URI query from schema file", zap.Error(err))
	}
	task.ImportIntoSchemaFile = schemaFile
	if len(tableMeta.GetDataFiles()) == 0 {
		task.ImportIntoFileType = agentpb.GetImportFilesResp_UNKNOWN_TYPE.String()
		task.ImportIntoFileCompression = agentpb.GetImportFilesResp_UNKNOWN.String()
		return
	}
	task.ImportIntoFileType = tableMeta.GetDataFiles()[0].GetSourceType().String()
	task.ImportIntoFileCompression = tableMeta.GetDataFiles()[0].GetCompression().String()
}

func populateImportIntoSubTasksFromFileDiscovery(imp *entity.ImportIntoTask, resp *agentpb.GetImportFilesResp) error {
	if len(imp.ImportResults.SubTasks) != 0 {
		if len(resp.GetTableMeta()) != len(imp.ImportResults.SubTasks) {
			return fmt.Errorf("unexpected import files result count %d, expected %d", len(resp.GetTableMeta()), len(imp.ImportResults.SubTasks))
		}
		for i, tableMeta := range resp.GetTableMeta() {
			populateImportIntoSubTaskFromFileDiscovery(&imp.ImportResults.SubTasks[i], tableMeta)
		}
		return nil
	}

	subTasks := make([]entity.ImportIntoSubTask, len(resp.GetTableMeta()))
	for i, tbl := range resp.GetTableMeta() {
		storedWildcard, err := utils.StripURIQuery(tbl.WildcardPath)
		if err != nil {
			return err
		}
		subTasks[i] = entity.ImportIntoSubTask{
			DataBase:                tbl.Database,
			TableName:               tbl.Table,
			ImportIntoFilesWildcard: storedWildcard,
		}
		populateImportIntoSubTaskFromFileDiscovery(&subTasks[i], tbl)
	}
	imp.ImportResults = entity.ImportIntoSubTasks{SubTasks: subTasks}
	return nil
}

func hasPendingImportIntoSubTask(imp *entity.ImportIntoTask) bool {
	for i := range imp.ImportResults.SubTasks {
		if imp.ImportResults.SubTasks[i].ImportIntoJobStatus == entity.ImportIntoJobStatusPending {
			return true
		}
	}
	return false
}

type importIntoEstimatePreparationDecision int

const (
	importIntoEstimatePreparationDecisionWait importIntoEstimatePreparationDecision = iota
	importIntoEstimatePreparationDecisionEstimate
	importIntoEstimatePreparationDecisionFail
	importIntoEstimatePreparationDecisionComplete
)

func resolveImportIntoEstimatePreparationDecision(imp *entity.ImportIntoTask) importIntoEstimatePreparationDecision {
	// Async create-schemas keeps ImportIntoJobStatus empty while a remote schema task is
	// still running, so estimate must not start until every subtask reaches a terminal
	// schema-prepared state.
	if !isImportIntoSchemaPrepared(imp) {
		return importIntoEstimatePreparationDecisionWait
	}
	// Once every subtask has finished schema preparation, at least one pending subtask
	// means there is still import work to do. Estimate should proceed for exactly those
	// pending subtasks, while failed subtasks stay excluded from later phases.
	if hasPendingImportIntoSubTask(imp) {
		return importIntoEstimatePreparationDecisionEstimate
	}
	if hasFailedOrCanceledImportIntoSubTask(imp) {
		return importIntoEstimatePreparationDecisionFail
	}
	return importIntoEstimatePreparationDecisionComplete
}

func applyImportIntoEstimatePreparationDecision(imp *entity.ImportIntoTask) bool {
	switch resolveImportIntoEstimatePreparationDecision(imp) {
	case importIntoEstimatePreparationDecisionEstimate:
		return true
	case importIntoEstimatePreparationDecisionFail:
		applyImportIntoBusinessFailure(imp, "Import failed")
		return false
	case importIntoEstimatePreparationDecisionComplete:
		applyImportIntoCompletion(imp)
		return false
	default:
		return false
	}
}

func hasFailedOrCanceledImportIntoSubTask(imp *entity.ImportIntoTask) bool {
	for i := range imp.ImportResults.SubTasks {
		switch imp.ImportResults.SubTasks[i].ImportIntoJobStatus {
		case entity.ImportIntoJobStatusFailed, entity.ImportIntoJobStatusCanceled:
			return true
		}
	}
	return false
}

func applyImportIntoCompletion(imp *entity.ImportIntoTask) {
	imp.Status = entity.ImportIntoStatusCompleted
	imp.Message = importIntoTaskSuccessMessage
	imp.TargetPassword = ""
	cleanImportIntoStoredURIs(imp)
}

// cleanImportIntoStoredURIs removes credential query parameters from persisted URIs.
// Called when a task is canceled or reaches another terminal state.
func cleanImportIntoStoredURIs(imp *entity.ImportIntoTask) {
	if imp == nil {
		return
	}
	if stripped, err := utils.StripURIQuery(imp.SourceURI); err == nil {
		imp.SourceURI = stripped
	} else if imp.SourceURI != "" {
		log.Warn("failed to strip URI query from source URI", zap.Error(err))
	}
	for i := range imp.ImportResults.SubTasks {
		task := &imp.ImportResults.SubTasks[i]
		if stripped, err := utils.StripURIQuery(task.ImportIntoFilesWildcard); err == nil {
			task.ImportIntoFilesWildcard = stripped
		} else if task.ImportIntoFilesWildcard != "" {
			log.Warn("failed to strip URI query from files wildcard", zap.Error(err))
		}
		if stripped, err := utils.StripURIQuery(task.ImportIntoSchemaFile); err == nil {
			task.ImportIntoSchemaFile = stripped
		} else if task.ImportIntoSchemaFile != "" {
			log.Warn("failed to strip URI query from schema file", zap.Error(err))
		}
	}
}

// isImportTaskNotFoundErr returns true when err is the agent's HTTP/gRPC
// "task not found" response, which indicates the deterministic async task id
// has not been created yet. Callers use this to implement a get-first pattern
// and fall through to CreateXTask only when the task is absent, avoiding an
// unconditional Create payload (with PEM/password material) on every tick.
func isImportTaskNotFoundErr(err error) bool {
	if err == nil {
		return false
	}
	var errResp *agent.ErrorResponse
	if !errors.As(err, &errResp) {
		return false
	}
	return errResp.Code == int64(codes.NotFound)
}

type importIntoAsyncTaskState string

const (
	importIntoAsyncTaskStateSucceeded       importIntoAsyncTaskState = "succeeded"
	importIntoAsyncTaskStateWait            importIntoAsyncTaskState = "wait"
	importIntoAsyncTaskStateTerminalFailure importIntoAsyncTaskState = "terminal_failure"
)

// Shared async task state helpers.
func resolveImportIntoFileDiscoveryTaskState(resp *agentpb.GetImportFilesTaskResp) importIntoAsyncTaskState {
	if resp == nil {
		return importIntoAsyncTaskStateTerminalFailure
	}
	switch resp.GetStatus() {
	case agentpb.ImportFilesTaskStatus_IMPORT_FILES_TASK_STATUS_RUNNING,
		agentpb.ImportFilesTaskStatus_IMPORT_FILES_TASK_STATUS_UNSPECIFIED:
		return importIntoAsyncTaskStateWait
	case agentpb.ImportFilesTaskStatus_IMPORT_FILES_TASK_STATUS_SUCCEEDED:
		return importIntoAsyncTaskStateSucceeded
	case agentpb.ImportFilesTaskStatus_IMPORT_FILES_TASK_STATUS_FAILED:
		return importIntoAsyncTaskStateTerminalFailure
	default:
		return importIntoAsyncTaskStateTerminalFailure
	}
}

func buildImportIntoFileDiscoveryFailureMessage(resp *agentpb.GetImportFilesTaskResp) string {
	if resp == nil {
		return "Internal Server Error"
	}
	if resp.GetStatus() == agentpb.ImportFilesTaskStatus_IMPORT_FILES_TASK_STATUS_FAILED && resp.GetErrorMsg() != "" {
		return resp.GetErrorMsg()
	}
	if resp.GetStatus() == agentpb.ImportFilesTaskStatus_IMPORT_FILES_TASK_STATUS_FAILED {
		return "Import failed"
	}
	return "Internal Server Error"
}

func applyImportIntoFileDiscoveryTaskResult(imp *entity.ImportIntoTask, resp *agentpb.GetImportFilesTaskResp) error {
	return populateImportIntoSubTasksFromFileDiscovery(imp, resp.GetResult())
}

func buildCreateSchemasResultFailureMessage(result *agentpb.CreateSchemasResp) string {
	if result == nil {
		return "Internal Server Error"
	}
	if result.GetPemChanged() {
		return "Failed to create schemas. Please retry the import."
	}
	if result.GetErrorMsg() != "" {
		return result.GetErrorMsg()
	}
	for _, tableResult := range result.GetTableResults() {
		if tableResult.GetStatus() != agentpb.CreateSchemasResp_TableResult_STATUS_FAILED {
			continue
		}
		if tableResult.GetErrorMsg() != "" {
			return tableResult.GetErrorMsg()
		}
		return "Failed to create schemas. Please retry the import."
	}
	return ""
}

func findCreateSchemasTableResult(
	result *agentpb.CreateSchemasResp,
	database string,
	table string,
) *agentpb.CreateSchemasResp_TableResult {
	if result == nil {
		return nil
	}
	for _, tableResult := range result.GetTableResults() {
		if tableResult.GetDatabase() == database && tableResult.GetTable() == table {
			return tableResult
		}
	}
	return nil
}

type createSchemasTableResultKey struct {
	database string
	table    string
}

func indexCreateSchemasTableResults(result *agentpb.CreateSchemasResp) map[createSchemasTableResultKey]*agentpb.CreateSchemasResp_TableResult {
	if result == nil || len(result.GetTableResults()) == 0 {
		return nil
	}
	indexed := make(map[createSchemasTableResultKey]*agentpb.CreateSchemasResp_TableResult, len(result.GetTableResults()))
	for _, tableResult := range result.GetTableResults() {
		indexed[createSchemasTableResultKey{
			database: tableResult.GetDatabase(),
			table:    tableResult.GetTable(),
		}] = tableResult
	}
	return indexed
}

func buildCreateSchemasTableResultFailureMessage(
	tableResult *agentpb.CreateSchemasResp_TableResult,
	fallbackMsg string,
) string {
	if tableResult == nil {
		return ""
	}
	switch tableResult.GetStatus() {
	case agentpb.CreateSchemasResp_TableResult_STATUS_FAILED:
		if tableResult.GetErrorMsg() != "" {
			return tableResult.GetErrorMsg()
		}
		return fallbackMsg
	default:
		return ""
	}
}

func buildCreateSchemasSubTaskFailureMessage(
	task *entity.ImportIntoSubTask,
	result *agentpb.CreateSchemasResp,
	fallbackMsg string,
) string {
	if task == nil {
		return ""
	}
	if result == nil || len(result.GetTableResults()) == 0 {
		return ""
	}
	if tableResult := findCreateSchemasTableResult(result, task.DataBase, task.TableName); tableResult != nil {
		return buildCreateSchemasTableResultFailureMessage(tableResult, fallbackMsg)
	}
	return ""
}

func resolveCreateSchemasTaskState(resp *agentpb.GetSchemasTaskResp) importIntoAsyncTaskState {
	if resp == nil {
		return importIntoAsyncTaskStateTerminalFailure
	}
	switch resp.GetStatus() {
	case agentpb.CreateSchemasTaskStatus_CREATE_SCHEMAS_TASK_STATUS_RUNNING,
		agentpb.CreateSchemasTaskStatus_CREATE_SCHEMAS_TASK_STATUS_UNSPECIFIED:
		return importIntoAsyncTaskStateWait
	case agentpb.CreateSchemasTaskStatus_CREATE_SCHEMAS_TASK_STATUS_SUCCEEDED:
		return importIntoAsyncTaskStateSucceeded
	case agentpb.CreateSchemasTaskStatus_CREATE_SCHEMAS_TASK_STATUS_FAILED:
		return importIntoAsyncTaskStateTerminalFailure
	default:
		return importIntoAsyncTaskStateTerminalFailure
	}
}

func buildCreateSchemasTaskFailureMessage(resp *agentpb.GetSchemasTaskResp) string {
	if resp == nil {
		return "Internal Server Error"
	}
	if resp.GetErrorMsg() != "" {
		return resp.GetErrorMsg()
	}
	if resp.GetStatus() == agentpb.CreateSchemasTaskStatus_CREATE_SCHEMAS_TASK_STATUS_FAILED {
		return "Import failed"
	}
	return buildCreateSchemasResultFailureMessage(resp.GetResult())
}

func applyCreateSchemasResult(task *entity.ImportIntoSubTask, result *agentpb.CreateSchemasResp) {
	if msg := buildCreateSchemasResultFailureMessage(result); msg != "" {
		if taskMsg := buildCreateSchemasSubTaskFailureMessage(task, result, msg); taskMsg != "" {
			setImportIntoSubTaskStatus(task, entity.ImportIntoJobStatusFailed)
			task.Message = taskMsg
		}
		return
	}
	setImportIntoSubTaskStatus(task, entity.ImportIntoJobStatusPending)
	task.Message = ""
}

func applyCreateSchemasTaskResult(task *entity.ImportIntoSubTask, resp *agentpb.GetSchemasTaskResp) {
	if msg := buildCreateSchemasTaskFailureMessage(resp); msg != "" {
		if taskMsg := buildCreateSchemasSubTaskFailureMessage(task, resp.GetResult(), msg); taskMsg != "" {
			setImportIntoSubTaskStatus(task, entity.ImportIntoJobStatusFailed)
			task.Message = taskMsg
		}
		return
	}
	applyCreateSchemasResult(task, resp.GetResult())
}

func collectImportIntoPendingSchemas(imp *entity.ImportIntoTask) []*agentpb.CreateSchemasReq_Schema {
	schemas := make([]*agentpb.CreateSchemasReq_Schema, 0, len(imp.ImportResults.SubTasks))
	for i := range imp.ImportResults.SubTasks {
		task := &imp.ImportResults.SubTasks[i]
		if task.ImportIntoJobStatus != "" {
			continue
		}
		if isImportIntoSchemaOnlySubTask(task) {
			schemas = append(schemas, &agentpb.CreateSchemasReq_Schema{
				Database: task.DataBase,
				Table:    task.TableName,
			})
			continue
		}
		if task.ImportIntoFileType == agentpb.GetImportFilesResp_UNKNOWN_TYPE.String() {
			applyImportIntoSubTaskFailure(task, "Nothing to import")
			continue
		}
		if task.ImportIntoSchemaFile == "" {
			setImportIntoSubTaskStatus(task, entity.ImportIntoJobStatusPending)
			task.Message = ""
			continue
		}
		schemas = append(schemas, &agentpb.CreateSchemasReq_Schema{
			Database: task.DataBase,
			Table:    task.TableName,
		})
	}
	return schemas
}

func applyImportIntoSchemaBatchSuccess(imp *entity.ImportIntoTask) {
	for i := range imp.ImportResults.SubTasks {
		task := &imp.ImportResults.SubTasks[i]
		if task.ImportIntoJobStatus != "" {
			continue
		}
		if isImportIntoSchemaOnlySubTask(task) {
			setImportIntoSchemaOnlySubTaskSuccess(task)
			continue
		}
		if task.ImportIntoSchemaFile == "" || task.ImportIntoFileType == agentpb.GetImportFilesResp_UNKNOWN_TYPE.String() {
			continue
		}
		setImportIntoSubTaskStatus(task, entity.ImportIntoJobStatusPending)
		task.Message = ""
	}
}

func applyImportIntoSchemaBatchFailure(
	imp *entity.ImportIntoTask,
	msg string,
	result *agentpb.CreateSchemasResp,
	controlPlaneFailed bool,
) {
	applyImportIntoFailure(imp, msg, controlPlaneFailed)
	tableResultsByName := indexCreateSchemasTableResults(result)
	if len(tableResultsByName) == 0 {
		return
	}
	for i := range imp.ImportResults.SubTasks {
		task := &imp.ImportResults.SubTasks[i]
		switch task.ImportIntoJobStatus {
		case entity.ImportIntoJobStatusFinished, entity.ImportIntoJobStatusCanceled, entity.ImportIntoJobStatusFailed:
			continue
		default:
			tableResult := tableResultsByName[createSchemasTableResultKey{
				database: task.DataBase,
				table:    task.TableName,
			}]
			if tableResult == nil {
				continue
			}
			switch tableResult.GetStatus() {
			case agentpb.CreateSchemasResp_TableResult_STATUS_SUCCEEDED:
				setImportIntoSubTaskStatus(task, entity.ImportIntoJobStatusPending)
				task.Message = ""
			case agentpb.CreateSchemasResp_TableResult_STATUS_FAILED:
				applyImportIntoSubTaskFailure(task, buildCreateSchemasTableResultFailureMessage(tableResult, msg))
			}
		}
	}
}

func resolveImportIntoJobTaskState(resp *agentpb.GetImportIntoJobTaskResp) importIntoAsyncTaskState {
	if resp == nil {
		return importIntoAsyncTaskStateTerminalFailure
	}
	switch resp.GetStatus() {
	case agentpb.ImportIntoJobTaskStatus_IMPORT_INTO_JOB_TASK_STATUS_RUNNING,
		agentpb.ImportIntoJobTaskStatus_IMPORT_INTO_JOB_TASK_STATUS_UNSPECIFIED:
		return importIntoAsyncTaskStateWait
	case agentpb.ImportIntoJobTaskStatus_IMPORT_INTO_JOB_TASK_STATUS_SUCCEEDED:
		return importIntoAsyncTaskStateSucceeded
	case agentpb.ImportIntoJobTaskStatus_IMPORT_INTO_JOB_TASK_STATUS_FAILED:
		return importIntoAsyncTaskStateTerminalFailure
	default:
		return importIntoAsyncTaskStateTerminalFailure
	}
}

func buildImportIntoJobTaskFailureMessage(resp *agentpb.GetImportIntoJobTaskResp) string {
	if resp == nil {
		return "Internal Server Error"
	}
	if resp.GetStatus() == agentpb.ImportIntoJobTaskStatus_IMPORT_INTO_JOB_TASK_STATUS_FAILED && resp.GetErrorMsg() != "" {
		return resp.GetErrorMsg()
	}
	if resp.GetStatus() == agentpb.ImportIntoJobTaskStatus_IMPORT_INTO_JOB_TASK_STATUS_FAILED {
		return "Import failed"
	}
	return "Internal Server Error"
}

func hasImportIntoPendingSubmitTask(task *entity.ImportIntoSubTask) bool {
	return task.ImportIntoJobID == 0 && task.ImportIntoPhase == importIntoSubmitTaskPendingPhase
}

func markImportIntoPendingSubmitTask(task *entity.ImportIntoSubTask) {
	task.ImportIntoPhase = importIntoSubmitTaskPendingPhase
}

func clearImportIntoPendingSubmitTask(task *entity.ImportIntoSubTask) {
	if task.ImportIntoPhase == importIntoSubmitTaskPendingPhase {
		task.ImportIntoPhase = ""
	}
}

func buildPendingImportIntoJobTaskResp(taskID string) *agentpb.GetImportIntoJobTaskResp {
	return &agentpb.GetImportIntoJobTaskResp{
		TaskId: taskID,
		Status: agentpb.ImportIntoJobTaskStatus_IMPORT_INTO_JOB_TASK_STATUS_RUNNING,
	}
}

func buildWaitImportIntoJobTaskResponses(taskIDs []string) map[string]*agentpb.GetImportIntoJobTaskResp {
	result := make(map[string]*agentpb.GetImportIntoJobTaskResp, len(taskIDs))
	for _, taskID := range taskIDs {
		result[taskID] = buildPendingImportIntoJobTaskResp(taskID)
	}
	return result
}

func applyImportIntoJobTaskResult(task *entity.ImportIntoSubTask, resp *agentpb.GetImportIntoJobTaskResp) {
	clearImportIntoPendingSubmitTask(task)
	task.ImportIntoJobID = uint64(resp.GetResult().GetJobId())
	syncImportIntoSubTaskResult(task)
}

func isRecoverableImportIntoSubmitFailureErr(msg string) bool {
	msg = strings.ToLower(msg)
	return strings.Contains(msg, "there is active job on the target table already") ||
		strings.Contains(msg, "target table is not empty")
}

func buildShowImportJobsByGroupKeySQL(groupKey string) string {
	return fmt.Sprintf("SHOW IMPORT JOBS WHERE group_key = %s", FormatStringOptionValue(groupKey))
}

func importIntoTargetTableName(task *entity.ImportIntoSubTask) string {
	return fmt.Sprintf("`%s`.`%s`", task.DataBase, task.TableName)
}

func importIntoColumnIndexMap(resp *agentpb.ExecSQLResp) (map[string]int, error) {
	if resp == nil {
		return nil, pingErr.New("failed to exec sql")
	}
	indexMap := make(map[string]int, len(resp.ColumnNames))
	for i, cname := range resp.ColumnNames {
		indexMap[cname] = i
	}
	return indexMap, nil
}

func importIntoColumnValuesFromRowWithIndex(resp *agentpb.ExecSQLResp, rowIndex int, indexMap map[string]int, columnNames ...string) (map[string]*agentpb.ColumnValue, error) {
	if resp == nil || rowIndex < 0 || rowIndex >= len(resp.Rows) || resp.Rows[rowIndex] == nil {
		return nil, pingErr.New("failed to exec sql")
	}
	row := resp.Rows[rowIndex]
	if len(resp.ColumnNames) != len(row.TypedColumnValues) {
		return nil, pingErr.New("failed to exec sql")
	}
	values := make(map[string]*agentpb.ColumnValue, len(columnNames))
	for _, col := range columnNames {
		index, ok := indexMap[col]
		if !ok {
			return nil, fmt.Errorf("not found the column %s", col)
		}
		values[col] = row.TypedColumnValues[index]
	}
	return values, nil
}

func importIntoColumnValuesFromRow(resp *agentpb.ExecSQLResp, rowIndex int, columnNames ...string) (map[string]*agentpb.ColumnValue, error) {
	indexMap, err := importIntoColumnIndexMap(resp)
	if err != nil {
		return nil, err
	}
	return importIntoColumnValuesFromRowWithIndex(resp, rowIndex, indexMap, columnNames...)
}

type importIntoJobMatch struct {
	jobID  uint64
	status string
}

type importIntoJobStatusRow struct {
	jobID       uint64
	targetTable string
	taskStatus  *TaskStatus
}

type importIntoJobStatusLookup struct {
	byJobID       map[uint64]*importIntoJobStatusRow
	byTargetTable map[string][]*importIntoJobStatusRow
}

func importIntoNullableTimestampValueAsTime(value *agentpb.ColumnValue) time.Time {
	if value == nil {
		return time.Time{}
	}
	timestamp := value.GetTimestampValue()
	if timestamp == nil {
		return time.Time{}
	}
	return timestamp.AsTime()
}

func parseImportIntoJobStatusFromRow(resp *agentpb.ExecSQLResp, rowIndex int, indexMap map[string]int) (*importIntoJobStatusRow, error) {
	values, err := importIntoColumnValuesFromRowWithIndex(resp, rowIndex, indexMap,
		JobIDColName,
		TargetTableColName,
		PhaseColName,
		StatusColName,
		ImportRowsColName,
		ResultMessageColName,
		CreateTime,
		StartTime,
		EndTimeColName,
		CurStepColName,
		CurStepProcessedSizeColName,
		CurStepTotalSizeColName,
		CurStepProgressPctColName,
		CurStepSpeedColName,
		CurStepETAColName,
	)
	if err != nil {
		return nil, err
	}
	jobID := values[JobIDColName].GetInt64Value()
	if jobID <= 0 {
		return nil, fmt.Errorf("invalid import job id %d", jobID)
	}
	pct := uint64(0)
	if parsedPct, err := parseImportIntoCurrentStepProgressPercent(values[CurStepProgressPctColName].GetStringValue()); err == nil {
		pct = parsedPct
	}
	return &importIntoJobStatusRow{
		jobID:       uint64(jobID),
		targetTable: values[TargetTableColName].GetStringValue(),
		taskStatus: &TaskStatus{
			Phase:                        values[PhaseColName].GetStringValue(),
			Status:                       values[StatusColName].GetStringValue(),
			ImportedRows:                 values[ImportRowsColName].GetInt64Value(),
			ResultMessage:                values[ResultMessageColName].GetStringValue(),
			StartTime:                    importIntoNullableTimestampValueAsTime(values[StartTime]),
			CreateTime:                   importIntoNullableTimestampValueAsTime(values[CreateTime]),
			EndTime:                      importIntoNullableTimestampValueAsTime(values[EndTimeColName]),
			CurrentStep:                  values[CurStepColName].GetStringValue(),
			CurrentStepTotalSize:         values[CurStepTotalSizeColName].GetStringValue(),
			CurrentStepProgressPercent:   int32(pct),
			CurrentStepETA:               values[CurStepETAColName].GetStringValue(),
			CurrentStepProcessedDataSize: values[CurStepProcessedSizeColName].GetStringValue(),
			CurrentStepSpeed:             values[CurStepSpeedColName].GetStringValue(),
		},
	}, nil
}

func findImportIntoJobStatusesFromResp(resp *agentpb.ExecSQLResp) (*importIntoJobStatusLookup, error) {
	indexMap, err := importIntoColumnIndexMap(resp)
	if err != nil {
		return nil, err
	}
	lookup := &importIntoJobStatusLookup{
		byJobID:       make(map[uint64]*importIntoJobStatusRow, len(resp.GetRows())),
		byTargetTable: make(map[string][]*importIntoJobStatusRow, len(resp.GetRows())),
	}
	for i := range resp.GetRows() {
		row, err := parseImportIntoJobStatusFromRow(resp, i, indexMap)
		if err != nil {
			return nil, err
		}
		lookup.byJobID[row.jobID] = row
		lookup.byTargetTable[row.targetTable] = append(lookup.byTargetTable[row.targetTable], row)
	}
	return lookup, nil
}

func (l *importIntoJobStatusLookup) findByJobID(jobID uint64) (*TaskStatus, bool) {
	if l == nil {
		return nil, false
	}
	row, ok := l.byJobID[jobID]
	if !ok {
		return nil, false
	}
	return row.taskStatus, true
}

func (l *importIntoJobStatusLookup) findByTargetTable(targetTable string) ([]*importIntoJobStatusRow, bool) {
	if l == nil {
		return nil, false
	}
	statuses, ok := l.byTargetTable[targetTable]
	return statuses, ok
}

func findRecoverableImportIntoJobIDFromLookup(lookup *importIntoJobStatusLookup, imp *entity.ImportIntoTask, task *entity.ImportIntoSubTask) (uint64, error) {
	targetTable := importIntoTargetTableName(task)
	rows, ok := lookup.findByTargetTable(targetTable)
	if !ok || len(rows) == 0 {
		return 0, fmt.Errorf("no import job matched group key %s and target table %s", imp.GroupKey, targetTable)
	}
	if len(rows) > 1 {
		return 0, fmt.Errorf("multiple import jobs matched group key %s and target table %s", imp.GroupKey, targetTable)
	}
	return rows[0].jobID, nil
}

func findImportIntoMatchingJobs(resp *agentpb.ExecSQLResp, imp *entity.ImportIntoTask, task *entity.ImportIntoSubTask) ([]importIntoJobMatch, error) {
	targetTable := importIntoTargetTableName(task)
	var matchedJobs []importIntoJobMatch
	for i := range resp.GetRows() {
		values, err := importIntoColumnValuesFromRow(resp, i,
			JobIDColName,
			GroupKeyColName,
			TargetTableColName,
			StatusColName,
		)
		if err != nil {
			return nil, err
		}
		if values[GroupKeyColName].GetStringValue() != imp.GroupKey {
			continue
		}
		if values[TargetTableColName].GetStringValue() != targetTable {
			continue
		}
		jobID := values[JobIDColName].GetInt64Value()
		if jobID <= 0 {
			return nil, fmt.Errorf("invalid import job id %d", jobID)
		}
		matchedJobs = append(matchedJobs, importIntoJobMatch{
			jobID:  uint64(jobID),
			status: values[StatusColName].GetStringValue(),
		})
	}
	return matchedJobs, nil
}

func findRecoverableImportIntoJobID(resp *agentpb.ExecSQLResp, imp *entity.ImportIntoTask, task *entity.ImportIntoSubTask) (uint64, error) {
	matchedJobs, err := findImportIntoMatchingJobs(resp, imp, task)
	if err != nil {
		return 0, err
	}

	targetTable := importIntoTargetTableName(task)
	switch len(matchedJobs) {
	case 0:
		return 0, fmt.Errorf("no import job matched group key %s and target table %s", imp.GroupKey, targetTable)
	case 1:
		return matchedJobs[0].jobID, nil
	default:
		return 0, fmt.Errorf("multiple import jobs matched group key %s and target table %s", imp.GroupKey, targetTable)
	}
}

type importIntoExecSQLFunc func(context.Context, *entity.ImportIntoTask, string, []string, int) (*agentpb.ExecSQLResp, error)

type importIntoSQLExecutor struct {
	agentClient       *agent.NGRegionalClient
	clusterID         uint64
	username          string
	decryptedPassword string
	cluster           *agentpb.Cluster
}

func (c *ImportIntoModel) newImportIntoSQLExecutor(
	ctx context.Context,
	imp *entity.ImportIntoTask,
	agentClient *agent.NGRegionalClient,
) (*importIntoSQLExecutor, error) {
	if agentClient == nil {
		var err error
		agentClient, err = c.newNGRegionalClientByClusterID(ctx, imp.ClusterID.Uint64())
		if err != nil {
			return nil, err
		}
	}
	passwd, err := c.encryptor.DecryptBase64(imp.TargetPassword)
	if err != nil {
		return nil, err
	}
	return &importIntoSQLExecutor{
		agentClient:       agentClient,
		clusterID:         imp.ClusterID.Uint64(),
		username:          imp.TargetUsername,
		decryptedPassword: passwd,
	}, nil
}

func (e *importIntoSQLExecutor) ensureCluster(ctx context.Context) error {
	if e.cluster != nil {
		return nil
	}
	pemResp, err := e.agentClient.GetPublicKey(ctx, &agentpb.GetPublicKeyRequest{})
	if err != nil {
		return err
	}
	pem := pemResp.GetPem()
	data, err := pkgcrypto.EncryptWithPEM(pem, []byte(e.decryptedPassword))
	if err != nil {
		return err
	}
	e.cluster = &agentpb.Cluster{
		ClusterId:         e.clusterID,
		User:              e.username,
		PasswordEncrypted: true,
		EncryptedPassword: &agentpb.EncryptedData{
			Data: data,
			Pem:  pem,
		},
	}
	return nil
}

func (e *importIntoSQLExecutor) execute(
	ctx context.Context,
	_ *entity.ImportIntoTask,
	sql string,
	setupSQLs []string,
	columnNum int,
) (*agentpb.ExecSQLResp, error) {
	var resp *agentpb.ExecSQLResp
	err := retry.Do(func() error {
		if err := e.ensureCluster(ctx); err != nil {
			return err
		}
		execSQLReq := &agentpb.ExecSQLReq{
			Sql:          sql,
			SetupSqls:    setupSQLs,
			Cluster:      e.cluster,
			TypedColumns: true,
			ColumnNum:    uint64(columnNum),
		}
		var err error
		resp, err = e.agentClient.ExecuteSQL(ctx, execSQLReq)
		if err != nil {
			return err
		}
		if resp.PemChanged {
			e.cluster = nil
			return agent.PEMChangedError
		}
		if resp.GetErrorMsg() != "" {
			return &execSQLResponseError{msg: resp.GetErrorMsg()}
		}
		return nil
	}, retry.Attempts(serviceconfig.Global().ImportConfig.MaxRetryTimesOnExecImportingQuerySQL), retry.RetryIf(func(err error) bool {
		return errors.Is(err, agent.PEMChangedError)
	}), retry.LastErrorOnly(true))
	return resp, err
}

func recoverImportIntoJobIDWithExec(
	ctx context.Context,
	imp *entity.ImportIntoTask,
	task *entity.ImportIntoSubTask,
	execSQL importIntoExecSQLFunc,
) (uint64, error) {
	showSQL := buildShowImportJobsByGroupKeySQL(imp.GroupKey)
	resp, err := execSQL(ctx, imp, showSQL, nil, 18)
	if err != nil {
		return 0, err
	}
	return findRecoverableImportIntoJobID(resp, imp, task)
}

func (c *ImportIntoModel) recoverImportIntoJobID(ctx context.Context, imp *entity.ImportIntoTask, task *entity.ImportIntoSubTask) (uint64, error) {
	return recoverImportIntoJobIDWithExec(ctx, imp, task, c.executeSQL)
}

func resolveImportIntoEstimateTaskState(resp *agentpb.GetImportDataSizeEstimateTaskResp) importIntoAsyncTaskState {
	if resp == nil {
		return importIntoAsyncTaskStateTerminalFailure
	}
	switch resp.GetStatus() {
	case agentpb.ImportDataSizeEstimateTaskStatus_IMPORT_DATA_SIZE_ESTIMATE_TASK_STATUS_RUNNING,
		agentpb.ImportDataSizeEstimateTaskStatus_IMPORT_DATA_SIZE_ESTIMATE_TASK_STATUS_UNSPECIFIED:
		return importIntoAsyncTaskStateWait
	case agentpb.ImportDataSizeEstimateTaskStatus_IMPORT_DATA_SIZE_ESTIMATE_TASK_STATUS_SUCCEEDED:
		return importIntoAsyncTaskStateSucceeded
	case agentpb.ImportDataSizeEstimateTaskStatus_IMPORT_DATA_SIZE_ESTIMATE_TASK_STATUS_FAILED:
		return importIntoAsyncTaskStateTerminalFailure
	default:
		return importIntoAsyncTaskStateTerminalFailure
	}
}

func buildImportIntoEstimateTaskFailureMessage(resp *agentpb.GetImportDataSizeEstimateTaskResp) string {
	if resp == nil {
		return "Internal Server Error"
	}
	if resp.GetStatus() == agentpb.ImportDataSizeEstimateTaskStatus_IMPORT_DATA_SIZE_ESTIMATE_TASK_STATUS_FAILED {
		msg := resp.GetErrorMsg()
		if msg != "" {
			return msg
		}
		return "failed to estimate import data size"
	}
	return fmt.Sprintf("unexpected import data size estimate task status %s", resp.GetStatus().String())
}

func buildImportIntoEstimateResultFailureMessage(result *agentpb.GetImportDataSizeEstimateTaskResp_Result) string {
	if result == nil {
		return "import data size estimate result is empty"
	}
	if len(result.GetTables()) == 0 {
		return "failed to estimate import data size: all matched tables were skipped, likely because schema files are invalid or unsupported"
	}
	if result.GetTotalSourceSize() == 0 || result.GetTotalTikvSize() == 0 {
		return "import data is empty"
	}
	return ""
}

func applyImportIntoEstimateTaskResult(imp *entity.ImportIntoTask, resp *agentpb.GetImportDataSizeEstimateTaskResp) {
	imp.TotalSize = resp.GetResult().GetTotalSourceSize()
}

func applyImportIntoFailure(imp *entity.ImportIntoTask, msg string, controlPlaneFailed bool) {
	imp.Status = entity.ImportIntoStatusFailed
	imp.Message = msg
	imp.ImportResults.ControlPlaneFailed = imp.ImportResults.ControlPlaneFailed || controlPlaneFailed
	cleanImportIntoStoredURIs(imp)
}

func applyImportIntoSQLFailure(imp *entity.ImportIntoTask, msg string) {
	imp.Status = entity.ImportIntoStatusFailed
	imp.Message = msg
	imp.ImportResults.ControlPlaneFailed = false
	cleanImportIntoStoredURIs(imp)
}

func applyImportIntoControlPlaneFailure(imp *entity.ImportIntoTask, msg string) {
	applyImportIntoFailure(imp, msg, true)
}

func applyImportIntoBusinessFailure(imp *entity.ImportIntoTask, msg string) {
	applyImportIntoFailure(imp, msg, false)
}

func buildImportIntoSourceURIForAgent(imp *entity.ImportIntoTask) (string, error) {
	sourceURIForAgent := imp.SourceURI
	if imp.ServerTier == entity.ServerTierNextGeneration {
		var err error
		sourceURIForAgent, err = appendExternalIDToSourceURIIfNeeded(sourceURIForAgent, imp.ClusterID.Uint64())
		if err != nil {
			return "", err
		}
	}

	parsedSourceURL, err := url.Parse(sourceURIForAgent)
	if err != nil {
		return "", err
	}
	if parsedSourceURL.Scheme == "oss" {
		parsedSourceURL.Scheme = "s3"
	}
	return parsedSourceURL.String(), nil
}

func (c *ImportIntoModel) persistImportIntoTaskState(
	ctx context.Context,
	logger *zap.SugaredLogger,
	imp *entity.ImportIntoTask,
	originalErr error,
	prevStatus entity.ImportIntoStatus,
) error {
	updateFields := []entity.ImportIntoTaskDBSchemaField{
		entity.ImportIntoTaskDBSchema.Message,
		entity.ImportIntoTaskDBSchema.ImportResults,
		entity.ImportIntoTaskDBSchema.TotalSize,
	}
	if isTerminalImportIntoStatus(imp.Status) && imp.FinishedAt == nil {
		now := time.Now()
		imp.FinishedAt = &now
	}
	if isTerminalImportIntoStatus(imp.Status) {
		cleanImportIntoStoredURIs(imp)
	}
	if isTerminalImportIntoStatus(imp.Status) && imp.TargetPassword == "" {
		updateFields = append(updateFields, entity.ImportIntoTaskDBSchema.TargetPassword)
	}
	if isTerminalImportIntoStatus(imp.Status) {
		updateFields = append(updateFields, entity.ImportIntoTaskDBSchema.SourceURI)
	}
	if imp.FinishedAt != nil {
		updateFields = append(updateFields, entity.ImportIntoTaskDBSchema.FinishedAt)
	}
	terminalStatusPersisted := false
	txErr := db.TransactionWithContext(context.Background(), func(tx *gorm.DB) error {
		if updateErr := imp.Update(tx, updateFields...); updateErr != nil {
			return updateErr
		}
		statusUpdate := tx.Model(&entity.ImportIntoTask{}).
			Where("`status` != ? and id = ?", entity.ImportIntoStatusCanceling, imp.ID).
			Update("status", imp.Status)
		if statusUpdate.Error != nil {
			return statusUpdate.Error
		}
		terminalStatusPersisted = isTerminalImportIntoStatus(imp.Status) && statusUpdate.RowsAffected > 0
		return nil
	})
	if txErr == nil {
		publishImportIntoEventIfStatusChanged(ctx, c.eventPublisher, imp, prevStatus)
		if terminalStatusPersisted {
			if cleanupErr := c.cleanupImportClusterJob(ctx, logger, imp); cleanupErr != nil {
				logger.With(zap.Error(cleanupErr)).Error("failed to clean up Import Into custom status")
			}
		}
	}
	if originalErr == nil {
		return txErr
	}
	if txErr != nil {
		logger.Errorw("failed to persist import task state to DB while handling another error",
			"txError", txErr,
			"originalError", originalErr,
		)
	}
	return originalErr
}

func isTerminalImportIntoStatus(status entity.ImportIntoStatus) bool {
	switch status {
	case entity.ImportIntoStatusCompleted, entity.ImportIntoStatusFailed, entity.ImportIntoStatusCanceled:
		return true
	default:
		return false
	}
}

// Preparing-stage helpers.
func (c *ImportIntoModel) prepareImportIntoFileDiscovery(
	ctx context.Context,
	imp *entity.ImportIntoTask,
	agentClient *agent.NGRegionalClient,
	sourceURI string,
	logger *zap.SugaredLogger,
) error {
	if isImportIntoFileDiscoveryPrepared(imp) {
		return nil
	}

	getFilesTaskID := buildImportIntoGetFilesTaskID(imp.ID.Uint64())
	getFilesResp, getFilesErr := agentClient.GetImportFilesTask(ctx, getFilesTaskID)
	if getFilesErr != nil {
		if !isImportTaskNotFoundErr(getFilesErr) {
			logger.Errorw("failed to get import files task",
				"taskID", getFilesTaskID,
				"error", getFilesErr,
			)
			return getFilesErr
		}
		if _, createGetFilesErr := agentClient.CreateImportFilesTask(
			ctx,
			getFilesTaskID,
			imp.ClusterID.Uint64(),
			sourceURI,
			-1,
			collectImportIntoGetFilesTableMetas(imp)...,
		); createGetFilesErr != nil {
			logger.Errorw("failed to create import files task",
				"taskID", getFilesTaskID,
				"error", createGetFilesErr,
			)
			return createGetFilesErr
		}
		logger.Infow("created import files task, waiting for next poll", "taskID", getFilesTaskID)
		return nil
	}

	switch resolveImportIntoFileDiscoveryTaskState(getFilesResp) {
	case importIntoAsyncTaskStateWait:
		logger.Infow("waiting import files task", "taskID", getFilesTaskID)
		return nil
	case importIntoAsyncTaskStateTerminalFailure:
		getFilesTerminalMsg := buildImportIntoFileDiscoveryFailureMessage(getFilesResp)
		logger.Errorw("import files task terminal failure",
			"taskID", getFilesTaskID,
			"message", getFilesTerminalMsg,
			"status", getFilesResp.GetStatus().String(),
		)
		if getFilesResp == nil {
			applyImportIntoControlPlaneFailure(imp, getFilesTerminalMsg)
		} else {
			applyImportIntoBusinessFailure(imp, getFilesTerminalMsg)
		}
		return nil
	case importIntoAsyncTaskStateSucceeded:
	default:
		applyImportIntoControlPlaneFailure(imp, buildImportIntoFileDiscoveryFailureMessage(getFilesResp))
		return nil
	}

	if err := applyImportIntoFileDiscoveryTaskResult(imp, getFilesResp); err != nil {
		logger.Errorw("failed to apply import files task result", "taskID", getFilesTaskID, "error", err)
		applyImportIntoControlPlaneFailure(imp, "Internal Server Error")
		return nil
	}
	if len(imp.ImportResults.SubTasks) == 0 {
		applyImportIntoBusinessFailure(imp, "Import failed")
		return nil
	}
	for i := range imp.ImportResults.SubTasks {
		subTask := &imp.ImportResults.SubTasks[i]
		logger.With(
			zap.String("database", subTask.DataBase),
			zap.String("table", subTask.TableName),
			zap.String("wildcard", utils.RedactSourceURI(subTask.ImportIntoFilesWildcard)),
			zap.String("schema_file", utils.RedactSourceURI(subTask.ImportIntoSchemaFile)),
			zap.String("file_type", subTask.ImportIntoFileType),
			zap.Uint64("total_size", subTask.ImportIntoTotalSize),
			zap.String("compression", subTask.ImportIntoFileCompression),
		).Info("import file discovery completed for sub-task")
	}
	return nil
}

func (c *ImportIntoModel) prepareImportIntoEstimate(
	ctx context.Context,
	imp *entity.ImportIntoTask,
	agentClient *agent.NGRegionalClient,
	sourceURI string,
	logger *zap.SugaredLogger,
) (importIntoEstimateOutcome, uint64, error) {
	estimateTaskID := buildImportIntoEstimateTaskID(imp.ID.Uint64())
	estimateResp, err := agentClient.GetImportDataSizeEstimateTask(ctx, estimateTaskID)
	if err != nil {
		if !isImportTaskNotFoundErr(err) {
			logger.Errorw("failed to get import data size estimate task",
				"taskID", estimateTaskID,
				"error", err,
			)
			return importIntoEstimateOutcomeWait, 0, err
		}
		estimateCsvConfig, estimateCfgErr := buildImportIntoEstimateCsvConfig(imp)
		if estimateCfgErr != nil {
			logger.Errorw("failed to build import estimate csv config",
				"taskID", estimateTaskID,
				"error", estimateCfgErr,
			)
			applyImportIntoBusinessFailure(imp, estimateCfgErr.Error())
			return importIntoEstimateOutcomeFailed, 0, nil
		}
		if _, createEstimateErr := agentClient.CreateImportDataSizeEstimateTask(
			ctx,
			estimateTaskID,
			imp.ClusterID.Uint64(),
			sourceURI,
			-1,
			estimateCsvConfig,
			collectImportIntoEstimateTableMetas(imp)...,
		); createEstimateErr != nil {
			logger.Errorw("failed to create import data size estimate task",
				"taskID", estimateTaskID,
				"error", createEstimateErr,
			)
			return importIntoEstimateOutcomeWait, 0, createEstimateErr
		}
		logger.Infow("created import data size estimate task, waiting for next poll", "taskID", estimateTaskID)
		return importIntoEstimateOutcomeWait, 0, nil
	}
	switch resolveImportIntoEstimateTaskState(estimateResp) {
	case importIntoAsyncTaskStateWait:
		logger.Infow("waiting estimate import data size task", "taskID", estimateTaskID)
		return importIntoEstimateOutcomeWait, 0, nil
	case importIntoAsyncTaskStateTerminalFailure:
		estimateFailureMsg := buildImportIntoEstimateTaskFailureMessage(estimateResp)
		logger.Errorw("import data size estimate task failed",
			"taskID", estimateTaskID,
			"message", estimateFailureMsg,
			"status", estimateResp.GetStatus().String(),
		)
		if estimateResp == nil {
			applyImportIntoControlPlaneFailure(imp, estimateFailureMsg)
		} else {
			applyImportIntoBusinessFailure(imp, estimateFailureMsg)
		}
		return importIntoEstimateOutcomeFailed, 0, nil
	case importIntoAsyncTaskStateSucceeded:
	default:
		applyImportIntoControlPlaneFailure(imp, buildImportIntoEstimateTaskFailureMessage(estimateResp))
		return importIntoEstimateOutcomeFailed, 0, nil
	}

	if estimateErrorMessage := buildImportIntoEstimateResultFailureMessage(estimateResp.GetResult()); estimateErrorMessage != "" {
		logger.Errorw("import data size estimate validation failed",
			"taskID", estimateTaskID,
			"message", estimateErrorMessage,
		)
		applyImportIntoBusinessFailure(imp, estimateErrorMessage)
		return importIntoEstimateOutcomeFailed, 0, nil
	}
	logger.Infow("import data size estimate succeeded",
		"taskID", estimateTaskID,
		"totalSourceSize", estimateResp.GetResult().GetTotalSourceSize(),
		"totalTikvSize", estimateResp.GetResult().GetTotalTikvSize(),
		"tableCount", len(estimateResp.GetResult().GetTables()),
	)
	applyImportIntoEstimateTaskResult(imp, estimateResp)
	return importIntoEstimateOutcomeReady, estimateResp.GetResult().GetTotalTikvSize(), nil
}

// Schema preparation is orchestrated as one deterministic batch task per import.
// A batch failure still fails the whole import, but subtask messages should
// preserve which table actually triggered the schema error.
func (c *ImportIntoModel) prepareImportIntoSchemas(
	ctx context.Context,
	imp *entity.ImportIntoTask,
	agentClient *agent.NGRegionalClient,
	sourceURI string,
	logger *zap.SugaredLogger,
) error {
	// Decrypt the target password lazily: only when at least one sub-task needs
	// a CreateSchemasTask call (i.e., the task is NOT_FOUND on the agent).
	// Once cached, reuse across sub-tasks within the same tick. Avoids repeating
	// the decrypt + agent-side PEM fetch (inside CreateSchemasTask) every tick
	// after all schemas have already been created.
	var (
		cachedPasswd       string
		cachedPasswdLoaded bool
		cachedPasswdErr    error
	)
	loadPassword := func() (string, error) {
		if cachedPasswdLoaded {
			return cachedPasswd, cachedPasswdErr
		}
		cachedPasswdLoaded = true
		passwd, err := c.encryptor.DecryptBase64(imp.TargetPassword)
		if err != nil {
			cachedPasswdErr = err
			logger.Errorw("failed to decrypt target password", "error", err)
			return "", err
		}
		cachedPasswd = passwd
		return cachedPasswd, nil
	}

	schemas := collectImportIntoPendingSchemas(imp)
	if len(schemas) == 0 {
		return nil
	}

	taskID := buildImportIntoCreateSchemaTaskID(imp.ID.Uint64())
	schemaResp, getErr := agentClient.GetSchemasTask(ctx, taskID)
	if getErr != nil {
		if !isImportTaskNotFoundErr(getErr) {
			logger.Errorw("failed to get batch create-schemas task",
				"taskID", taskID,
				"schemaCount", len(schemas),
				"error", getErr,
			)
			return getErr
		}
		passwd, decryptErr := loadPassword()
		if decryptErr != nil {
			return decryptErr
		}
		if _, createErr := agentClient.CreateSchemasTask(ctx, taskID, imp.ClusterID.Uint64(), sourceURI, imp.TargetUsername, passwd, schemas...); createErr != nil {
			logger.Errorw("failed to create batch create-schemas task",
				"taskID", taskID,
				"schemaCount", len(schemas),
				"error", createErr,
			)
			return createErr
		}
		logger.Infow("created batch create-schemas task, waiting for next poll",
			"taskID", taskID,
			"schemaCount", len(schemas),
		)
		return nil
	}

	switch resolveCreateSchemasTaskState(schemaResp) {
	case importIntoAsyncTaskStateWait:
		logger.Infow("waiting batch create-schemas task",
			"taskID", taskID,
			"schemaCount", len(schemas),
		)
		return nil
	case importIntoAsyncTaskStateTerminalFailure:
		msg := buildCreateSchemasTaskFailureMessage(schemaResp)
		var result *agentpb.CreateSchemasResp
		status := ""
		if schemaResp != nil {
			result = schemaResp.GetResult()
			status = schemaResp.GetStatus().String()
		}
		logger.Errorw("batch create-schemas task failed",
			"taskID", taskID,
			"schemaCount", len(schemas),
			"message", msg,
			"status", status,
		)
		applyImportIntoSchemaBatchFailure(imp, msg, result, schemaResp == nil)
		return nil
	case importIntoAsyncTaskStateSucceeded:
		if msg := buildCreateSchemasTaskFailureMessage(schemaResp); msg != "" {
			logger.Errorw("batch create-schemas task returned business error",
				"taskID", taskID,
				"schemaCount", len(schemas),
				"message", msg,
			)
			applyImportIntoSchemaBatchFailure(imp, msg, schemaResp.GetResult(), false)
			return nil
		}
		applyImportIntoSchemaBatchSuccess(imp)
		return nil
	default:
		logger.Errorw("unexpected batch create-schemas task status",
			"taskID", taskID,
			"schemaCount", len(schemas),
			"status", schemaResp.GetStatus().String(),
		)
		applyImportIntoSchemaBatchFailure(imp, "Internal Server Error", schemaResp.GetResult(), true)
		return nil
	}
}

func shouldReleaseImportIntoReservationOnFailureStatus(status entity.ImportIntoStatus) bool {
	switch status {
	case entity.ImportIntoStatusPreparing,
		entity.ImportIntoStatusScalingUp,
		entity.ImportIntoStatusImporting,
		entity.ImportIntoStatusCanceling:
		return true
	default:
		return false
	}
}

// onImportIntoPreparing handles the preparation phase for IMPORT INTO.
// This is where individual IMPORT INTO jobs for each table are initiated and persisted.
func (c *ImportIntoModel) onImportIntoPreparing(ctx context.Context, imp *entity.ImportIntoTask) (err error) {
	logger := c.sugaredTraceIDLogger(ctx, zap.Uint64("importID", imp.ID.Uint64()), zap.Uint64("clusterID", imp.ClusterID.Uint64()))
	logger.Info("begin to prepare IMPORT INTO task")
	agentClient, err := c.newNGRegionalClientByClusterID(ctx, imp.ClusterID.Uint64())
	if err != nil {
		logger.Errorw("failed to create regional client", "error", err)
		return err
	}
	sourceURIWithExternalID, err := buildImportIntoSourceURIForAgent(imp)
	if err != nil {
		logger.Errorw("failed to build source URI for agent", "error", err)
		return err
	}
	reservationEnabled := c.isImportIntoTiKVCapacityReservationEnabled(imp.TenantID.Uint64(), imp.ProjectID.Uint64(), imp.ClusterID.Uint64())
	prevStatus := imp.Status
	defer func() {
		err = c.persistImportIntoTaskState(ctx, logger, imp, err, prevStatus)
	}()
	return runImportIntoPreparingSteps(imp, reservationEnabled, func() {
		applyImportIntoFileDiscoveryTotalSize(imp)
	}, importIntoPreparingHooks{
		prepareFileDiscovery: func() error {
			return c.prepareImportIntoFileDiscovery(ctx, imp, agentClient, sourceURIWithExternalID, logger)
		},
		prepareSchemas: func() error {
			return c.prepareImportIntoSchemas(ctx, imp, agentClient, sourceURIWithExternalID, logger)
		},
		prepareEstimate: func() (importIntoEstimateOutcome, uint64, error) {
			return c.prepareImportIntoEstimate(ctx, imp, agentClient, sourceURIWithExternalID, logger)
		},
		createReservation: func(totalTiKVSize uint64) error {
			sharedPoolID, err := getImportIntoSharedPoolID(ctx, agentClient, imp.ClusterID.Uint64())
			if err != nil {
				logger.Errorw("failed to get shared pool id for tikv capacity reservation", "error", err)
				return err
			}
			reservationID := buildImportIntoTiKVReservationID(imp.ID.Uint64())
			reservedTiKVBytes := reservedImportIntoTiKVBytes(totalTiKVSize)
			if _, createReservationErr := agentClient.CreateTiKVCapacityReservation(ctx, sharedPoolID, reservationID, reservedTiKVBytes); createReservationErr != nil {
				logger.Errorw("failed to create tikv capacity reservation",
					"sharedPoolID", sharedPoolID,
					"reservationID", reservationID,
					"reservedBytes", reservedTiKVBytes,
					"error", createReservationErr,
				)
				return createReservationErr
			}
			return nil
		},
	})
}

// execSQLResponseError is returned when the regional agent returns HTTP 200 but ExecSQLResp.error_msg
// is set (e.g. TiDB IMPORT INTO / SQL errors). It is distinct from *agent.ErrorResponse (HTTP 4xx JSON).
type execSQLResponseError struct {
	msg string
}

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

func importIntoSubtaskMessageFromErr(err error) string {
	if err == nil {
		return "Internal Server Error"
	}
	var errorResponse *agent.ErrorResponse
	if errors.As(err, &errorResponse) && errorResponse.Message != "" {
		return errorResponse.Message
	}
	var sqlRespErr *execSQLResponseError
	if errors.As(err, &sqlRespErr) && sqlRespErr.msg != "" {
		return sqlRespErr.msg
	}
	return "Internal Server Error"
}

// Shared SQL execution helpers for scaling/importing stages.
func (c *ImportIntoModel) executeSQL(ctx context.Context, imp *entity.ImportIntoTask, sql string, setupSQLs []string, columnNum int) (*agentpb.ExecSQLResp, error) {
	logger := c.sugaredTraceIDLogger(ctx, zap.Uint64("importID", imp.ID.Uint64()), zap.Uint64("clusterID", imp.ClusterID.Uint64()))
	executor, err := c.newImportIntoSQLExecutor(ctx, imp, nil)
	if err != nil {
		logger.With(zap.Error(err)).Error("failed to create import into sql executor")
		return nil, err
	}
	return executor.execute(ctx, imp, sql, setupSQLs, columnNum)
}

func (c *ImportIntoModel) buildImportIntoSQL(ctx context.Context, subTask *entity.ImportIntoSubTask, imp *entity.ImportIntoTask, options ...string) (string, []string, error) {
	logger := c.sugaredTraceIDLogger(ctx, zap.Uint64("importID", imp.ID.Uint64()), zap.Uint64("clusterID", imp.ClusterID.Uint64()))
	if imp.GroupKey == "" {
		return "", nil, fmt.Errorf("import into group key is empty")
	}
	withOptions := []string{
		DetachedOptionName,
		fmt.Sprintf("%s=%s", GroupKeyOptionName, FormatStringOptionValue(imp.GroupKey)),
	}
	withOptions = append(withOptions, options...)
	switch subTask.ImportIntoFileType {
	case agentpb.GetImportFilesResp_CSV.String():
		if imp.CustomCSVFormat != "" {
			var customCSVFormat importadminpb.CSVFormat
			err := protojson.Unmarshal([]byte(imp.CustomCSVFormat), &customCSVFormat)
			if err != nil {
				return "", nil, err
			}
			if customCSVFormat.GetSeparator() != "" {
				withOptions = append(withOptions, fmt.Sprintf("%s=%s", SeparatorOptionName, FormatStringOptionValue(customCSVFormat.GetSeparator())))
			}
			if customCSVFormat.GetDelimiter() != "" {
				withOptions = append(withOptions, fmt.Sprintf("%s=%s", DelimiterOptionName, FormatStringOptionValue(customCSVFormat.GetDelimiter())))
			}
			if customCSVFormat.GetEscapeChar() != "" {
				withOptions = append(withOptions, fmt.Sprintf("%s=%s", EscapeOptionName, FormatStringOptionValue(customCSVFormat.GetEscapeChar())))
			}
			if customCSVFormat.GetLineTerminator() != "" {
				val := FormatStringOptionValue(customCSVFormat.GetLineTerminator())
				if customCSVFormat.GetSplitFile() == true && subTask.ImportIntoFileCompression == agentpb.GetImportFilesResp_UNKNOWN.String() {
					withOptions = append(withOptions, SplitFileOptionName, fmt.Sprintf("%s=%s", LinesTerminatorOptionName, val))
					// need to explicitly specify LINES_TERMINATED_BY for this option.
				} else {
					switch val {
					case "'\\r\\n'", "'\\r'", "'\\n'":
						// TiDB can adapt to avoid errors caused by default values being passed to the front end
						logger.Info("skip line terminator", zap.String("line", val))
					default:
						withOptions = append(withOptions, fmt.Sprintf("%s=%s", LinesTerminatorOptionName, val))
					}
				}
			}
			if customCSVFormat.GetNullValue() != "" {
				withOptions = append(withOptions, fmt.Sprintf("%s=%s", NullOptionName, FormatStringOptionValue(customCSVFormat.GetNullValue())))
			}
			if customCSVFormat.GetSkipRows() != 0 {
				withOptions = append(withOptions, fmt.Sprintf("%s=%d", SkipRowSOptionName, customCSVFormat.GetSkipRows()))
			}
		}
	}
	customColumns := ""
	if subTask.Columns != nil {
		safeColumns := make([]string, len(subTask.Columns))
		for i, column := range subTask.Columns {
			safeColumns[i] = fmt.Sprintf("`%s`", FormatStringOptionValue(column))
		}
		customColumns = fmt.Sprintf("(%s)", strings.Join(safeColumns, ","))
	}
	importUri, err := utils.MergeURIQueryFromSource(imp.SourceURI, subTask.ImportIntoFilesWildcard)
	if err != nil {
		return "", nil, err
	}
	parsedUri, err := url.Parse(importUri)
	if err != nil {
		return "", nil, err
	}
	if parsedUri.Scheme == "oss" {
		parsedUri.Scheme = "s3"
		importUri = parsedUri.String()
	}
	importIntoSql := fmt.Sprintf("IMPORT INTO `%s`.`%s`%s FROM '%s' WITH %s", subTask.DataBase, subTask.TableName, customColumns, importUri, strings.Join(withOptions, ","))
	var setupSQLs []string
	return importIntoSql, setupSQLs, nil
}

// Scaling-up helpers.
func (c *ImportIntoModel) syncImportIntoTiKVReservation(
	ctx context.Context,
	imp *entity.ImportIntoTask,
	agentClient *agent.NGRegionalClient,
	logger *zap.SugaredLogger,
) (bool, error) {
	sharedPoolID, err := getImportIntoSharedPoolID(ctx, agentClient, imp.ClusterID.Uint64())
	if err != nil {
		logger.Errorw("failed to get shared pool id", "error", err)
		return false, err
	}
	reservationID := buildImportIntoTiKVReservationID(imp.ID.Uint64())
	reservation, err := agentClient.GetTiKVCapacityReservation(ctx, sharedPoolID, reservationID)
	if err != nil {
		logger.Errorw("failed to get tikv capacity reservation",
			"sharedPoolID", sharedPoolID,
			"reservationID", reservationID,
			"error", err,
		)
		return false, err
	}
	switch reservation.Status {
	case agent.TiKVCapacityReservationStatusProvisioning:
		logger.Infow("waiting tikv capacity reservation", "sharedPoolID", sharedPoolID, "reservationID", reservationID)
		return false, nil
	case agent.TiKVCapacityReservationStatusReady:
		logger.Infow("tikv capacity reservation is ready",
			"sharedPoolID", sharedPoolID,
			"reservationID", reservationID,
			"needScale", reservation.NeedScale,
		)
		return true, nil
	case agent.TiKVCapacityReservationStatusReleased:
		applyImportIntoControlPlaneFailure(imp, "tikv capacity reservation was released before import started")
		return false, nil
	default:
		c.releaseTiKVCapacityReservation(ctx, logger, agentClient, imp.TenantID.Uint64(), imp.ProjectID.Uint64(), imp.ClusterID.Uint64(), imp.ID.Uint64())
		applyImportIntoControlPlaneFailure(imp, fmt.Sprintf("unexpected tikv capacity reservation status %s", reservation.Status))
		return false, nil
	}
}

func sanitizeImportIntoSubTaskForLog(st *entity.ImportIntoSubTask) entity.ImportIntoSubTask {
	cp := *st
	cp.ImportIntoFilesWildcard = utils.RedactSourceURI(cp.ImportIntoFilesWildcard)
	cp.ImportIntoSchemaFile = utils.RedactSourceURI(cp.ImportIntoSchemaFile)
	return cp
}

func applyImportIntoSubTaskFailure(task *entity.ImportIntoSubTask, msg string) {
	clearImportIntoPendingSubmitTask(task)
	setImportIntoSubTaskStatus(task, entity.ImportIntoJobStatusFailed)
	task.Message = msg
}

func setImportIntoSchemaOnlySubTaskSuccess(task *entity.ImportIntoSubTask) {
	setImportIntoSubTaskStatus(task, entity.ImportIntoJobStatusFinished)
	task.Result = entity.ImportTableResultWarning
	task.Message = entity.ImportIntoSchemaOnlySuccessMessage
}

func setImportIntoSubTaskStatus(task *entity.ImportIntoSubTask, status string) {
	task.ImportIntoJobStatus = status
	syncImportIntoSubTaskResult(task)
}

func syncImportIntoSubTaskResult(task *entity.ImportIntoSubTask) {
	switch task.ImportIntoJobStatus {
	case entity.ImportIntoJobStatusFinished:
		if task.Result != entity.ImportTableResultWarning {
			task.Result = entity.ImportTableResultSuccess
		}
	case entity.ImportIntoJobStatusFailed:
		task.Result = entity.ImportTableResultError
	default:
		task.Result = ""
	}
}

func markImportIntoControlPlaneFailed(imp *entity.ImportIntoTask) {
	if imp == nil {
		return
	}
	imp.ImportResults.ControlPlaneFailed = true
}

type importIntoJobSubmissionDecision int

const (
	importIntoJobSubmissionDecisionNoop importIntoJobSubmissionDecision = iota
	importIntoJobSubmissionDecisionWait
	importIntoJobSubmissionDecisionSubmitted
	importIntoJobSubmissionDecisionFailed
)

// importIntoSubmissionState is the task-level aggregate state after one
// ScalingUp submission batch. It is derived from per-subtask
// importIntoJobSubmissionDecision outcomes that have already been persisted into
// each subtask's job id and job status.
type importIntoSubmissionState int

const (
	importIntoSubmissionStateWaiting importIntoSubmissionState = iota
	importIntoSubmissionStateReadyForImporting
	importIntoSubmissionStateFailed
)

type importIntoPreparingDecision int

const (
	importIntoPreparingDecisionWaitFileDiscovery importIntoPreparingDecision = iota
	importIntoPreparingDecisionWaitSchemaPreparation
	importIntoPreparingDecisionEstimate
	importIntoPreparingDecisionScalingUp
	importIntoPreparingDecisionFailed
	importIntoPreparingDecisionCompleted
)

type importIntoImportingDecision int

const (
	importIntoImportingDecisionWait importIntoImportingDecision = iota
	importIntoImportingDecisionCompleted
	importIntoImportingDecisionFailed
)

type importIntoCancelingDecision int

const (
	importIntoCancelingDecisionWait importIntoCancelingDecision = iota
	importIntoCancelingDecisionCanceled
)

type importIntoPreparingHooks struct {
	prepareFileDiscovery func() error
	prepareSchemas       func() error
	prepareEstimate      func() (importIntoEstimateOutcome, uint64, error)
	createReservation    func(totalTiKVSize uint64) error
}

type importIntoEstimateOutcome int

const (
	importIntoEstimateOutcomeWait importIntoEstimateOutcome = iota
	importIntoEstimateOutcomeFailed
	importIntoEstimateOutcomeReady
)

func applyTerminalImportIntoPreparingDecision(imp *entity.ImportIntoTask, decision importIntoPreparingDecision) {
	if imp == nil {
		return
	}
	if imp.Status != entity.ImportIntoStatusPreparing {
		return
	}
	switch decision {
	case importIntoPreparingDecisionFailed:
		applyImportIntoBusinessFailure(imp, "Import failed")
	case importIntoPreparingDecisionCompleted:
		applyImportIntoCompletion(imp)
	}
}

func resolveImportIntoSubmissionState(imp *entity.ImportIntoTask) importIntoSubmissionState {
	hasSubmittedJob := false
	hasPendingSubmission := false
	hasFailedOrCanceled := false
	for i := range imp.ImportResults.SubTasks {
		task := imp.ImportResults.SubTasks[i]
		switch task.ImportIntoJobStatus {
		case entity.ImportIntoJobStatusPending:
			if task.ImportIntoJobID != 0 {
				hasSubmittedJob = true
			} else {
				hasPendingSubmission = true
			}
		case entity.ImportIntoJobStatusRunning, entity.ImportIntoJobStatusFinished:
			if task.ImportIntoJobID != 0 {
				hasSubmittedJob = true
			}
		case entity.ImportIntoJobStatusFailed, entity.ImportIntoJobStatusCanceled:
			hasFailedOrCanceled = true
		}
	}

	switch {
	// Priority:
	// 1. Wait if any subtask is still pending submission without a real job id.
	case hasPendingSubmission:
		return importIntoSubmissionStateWaiting
	// 2. Once all pending submissions resolve to real job ids, transition to Importing.
	case hasSubmittedJob:
		return importIntoSubmissionStateReadyForImporting
	// 3. If nothing started and at least one subtask already failed/canceled, fail the whole task.
	case hasFailedOrCanceled:
		return importIntoSubmissionStateFailed
	// 4. Default to Importing when all submissions are terminal and none failed.
	default:
		return importIntoSubmissionStateReadyForImporting
	}
}

func resolveImportIntoPreparingDecision(imp *entity.ImportIntoTask, reservationEnabled bool) importIntoPreparingDecision {
	// Priority:
	// 1. Keep a failed task failed.
	if imp.Status == entity.ImportIntoStatusFailed {
		return importIntoPreparingDecisionFailed
	}
	if imp.Status == entity.ImportIntoStatusCompleted {
		return importIntoPreparingDecisionCompleted
	}
	// 2. Wait for file discovery to complete for every subtask.
	if !isImportIntoFileDiscoveryPrepared(imp) {
		return importIntoPreparingDecisionWaitFileDiscovery
	}
	// 3. Decide what schema preparation produced before moving on.
	switch resolveImportIntoEstimatePreparationDecision(imp) {
	case importIntoEstimatePreparationDecisionWait:
		return importIntoPreparingDecisionWaitSchemaPreparation
	case importIntoEstimatePreparationDecisionFail:
		return importIntoPreparingDecisionFailed
	case importIntoEstimatePreparationDecisionComplete:
		return importIntoPreparingDecisionCompleted
	}
	// 4. When reservation is enabled, always drive the estimate phase and let
	// the estimate task outcome decide when Preparing can advance.
	if reservationEnabled {
		return importIntoPreparingDecisionEstimate
	}
	// 5. Otherwise the task is ready to enter ScalingUp.
	return importIntoPreparingDecisionScalingUp
}

func resolveImportIntoImportingDecision(imp *entity.ImportIntoTask) importIntoImportingDecision {
	allCompleted := true
	anyFailedOrCanceled := false

	for _, impTbl := range imp.ImportResults.SubTasks {
		switch impTbl.ImportIntoJobStatus {
		case entity.ImportIntoJobStatusRunning, entity.ImportIntoJobStatusPending:
			allCompleted = false
		case entity.ImportIntoJobStatusFailed, entity.ImportIntoJobStatusCanceled:
			anyFailedOrCanceled = true
		}
	}

	// Priority:
	// 1. Wait while any IMPORT INTO job is still pending or running.
	if !allCompleted {
		return importIntoImportingDecisionWait
	}
	// 2. Once all jobs are terminal, fail if any terminal state is failed/canceled.
	if anyFailedOrCanceled {
		return importIntoImportingDecisionFailed
	}
	// 3. Otherwise all jobs completed successfully.
	return importIntoImportingDecisionCompleted
}

func resolveImportIntoCancelingDecision(imp *entity.ImportIntoTask) importIntoCancelingDecision {
	// Priority:
	// 1. Keep waiting while any subtask is still pending/running.
	for _, impTbl := range imp.ImportResults.SubTasks {
		switch impTbl.ImportIntoJobStatus {
		case entity.ImportIntoJobStatusRunning, entity.ImportIntoJobStatusPending:
			return importIntoCancelingDecisionWait
		}
	}
	// 2. Once every subtask is terminal, mark the whole task canceled.
	return importIntoCancelingDecisionCanceled
}

func runImportIntoPreparingSteps(
	imp *entity.ImportIntoTask,
	reservationEnabled bool,
	applyTotalSize func(),
	hooks importIntoPreparingHooks,
) error {
	if err := hooks.prepareFileDiscovery(); err != nil {
		return err
	}
	decision := resolveImportIntoPreparingDecision(imp, reservationEnabled)
	switch decision {
	case importIntoPreparingDecisionWaitFileDiscovery,
		importIntoPreparingDecisionFailed,
		importIntoPreparingDecisionCompleted:
		applyTerminalImportIntoPreparingDecision(imp, decision)
		return nil
	}
	if decision == importIntoPreparingDecisionWaitSchemaPreparation {
		if err := hooks.prepareSchemas(); err != nil {
			return err
		}
		decision = resolveImportIntoPreparingDecision(imp, reservationEnabled)
	}
	switch decision {
	case importIntoPreparingDecisionWaitSchemaPreparation,
		importIntoPreparingDecisionFailed,
		importIntoPreparingDecisionCompleted:
		applyTerminalImportIntoPreparingDecision(imp, decision)
		return nil
	case importIntoPreparingDecisionEstimate:
		estimateOutcome, totalTiKVSize, err := hooks.prepareEstimate()
		if err != nil {
			return err
		}
		switch estimateOutcome {
		case importIntoEstimateOutcomeWait,
			importIntoEstimateOutcomeFailed:
			return nil
		case importIntoEstimateOutcomeReady:
			if err := hooks.createReservation(totalTiKVSize); err != nil {
				return err
			}
			imp.Status = entity.ImportIntoStatusScalingUp
			return nil
		default:
			return nil
		}
	case importIntoPreparingDecisionScalingUp:
		applyTotalSize()
		if resolveImportIntoPreparingDecision(imp, reservationEnabled) != importIntoPreparingDecisionScalingUp {
			return nil
		}
		imp.Status = entity.ImportIntoStatusScalingUp
		return nil
	default:
		return nil
	}
}

func (c *ImportIntoModel) resolveImportIntoJobSubmissionResponse(
	ctx context.Context,
	imp *entity.ImportIntoTask,
	task *entity.ImportIntoSubTask,
	submitTaskID string,
	submitResp *agentpb.GetImportIntoJobTaskResp,
	execSQL importIntoExecSQLFunc,
	recoveryLookup *lazyImportIntoRecoveryLookup,
	logger *zap.SugaredLogger,
) importIntoJobSubmissionDecision {
	switch resolveImportIntoJobTaskState(submitResp) {
	case importIntoAsyncTaskStateWait:
		logger.With(
			zap.Any("subTask", sanitizeImportIntoSubTaskForLog(task)),
			zap.String("taskID", submitTaskID),
		).Debug("waiting import-into job submission task")
		return importIntoJobSubmissionDecisionWait
	case importIntoAsyncTaskStateTerminalFailure:
		msg := buildImportIntoJobTaskFailureMessage(submitResp)
		if isRecoverableImportIntoSubmitFailureErr(msg) {
			// A recoverable submit failure means the async task result is not
			// trustworthy enough to decide the final state, so recover from TiDB's
			// current IMPORT JOB state instead of failing the subtask immediately.
			var jobID uint64
			var recoverErr error
			if recoveryLookup != nil {
				var lookup *importIntoJobStatusLookup
				lookup, recoverErr = recoveryLookup.Get()
				if recoverErr == nil {
					jobID, recoverErr = findRecoverableImportIntoJobIDFromLookup(lookup, imp, task)
				}
			} else {
				jobID, recoverErr = recoverImportIntoJobIDWithExec(ctx, imp, task, execSQL)
			}
			if recoverErr == nil {
				task.ImportIntoJobID = jobID
				logger.With(
					zap.Any("subTask", sanitizeImportIntoSubTaskForLog(task)),
					zap.String("taskID", submitTaskID),
					zap.Uint64("jobID", task.ImportIntoJobID),
				).Info("recovered import-into job id after recoverable submit failure")
				return importIntoJobSubmissionDecisionSubmitted
			}
			logger.With(
				zap.Any("subTask", sanitizeImportIntoSubTaskForLog(task)),
				zap.String("taskID", submitTaskID),
				zap.String("message", msg),
				zap.Error(recoverErr),
			).Info("failed to recover import-into job id after recoverable submit failure; keep original submit error")
			markImportIntoControlPlaneFailed(imp)
			applyImportIntoSubTaskFailure(task, msg)
			return importIntoJobSubmissionDecisionFailed
		}
		logger.With(
			zap.Any("subTask", sanitizeImportIntoSubTaskForLog(task)),
			zap.String("taskID", submitTaskID),
			zap.String("message", msg),
			zap.String("status", submitResp.GetStatus().String()),
		).Error("import-into job submission task failed")
		markImportIntoControlPlaneFailed(imp)
		applyImportIntoSubTaskFailure(task, msg)
		return importIntoJobSubmissionDecisionFailed
	case importIntoAsyncTaskStateSucceeded:
	default:
		logger.With(
			zap.Any("subTask", sanitizeImportIntoSubTaskForLog(task)),
			zap.String("taskID", submitTaskID),
			zap.String("status", submitResp.GetStatus().String()),
		).Error("unexpected import-into job submission task status")
		markImportIntoControlPlaneFailed(imp)
		applyImportIntoSubTaskFailure(task, "Internal Server Error")
		return importIntoJobSubmissionDecisionFailed
	}

	applyImportIntoJobTaskResult(task, submitResp)
	logger.With(
		zap.Any("subTask", sanitizeImportIntoSubTaskForLog(task)),
		zap.Uint64("jobID", task.ImportIntoJobID),
		zap.String("taskID", submitTaskID),
	).Info("success to submit import sql")
	return importIntoJobSubmissionDecisionSubmitted
}

type importIntoJobSubmissionCreateItem struct {
	subTaskIndex int
	taskID       string
	task         *entity.ImportIntoSubTask
	request      *agentpb.BatchCreateImportIntoJobTasksReq_Task
}

// lazyImportIntoRecoveryLookup loads the group-key job lookup only when a
// recoverable submit failure actually needs it.
// One batch round may recover multiple subtasks, so the first lookup result or
// first lookup error is reused for the remaining recoveries in that round.
type lazyImportIntoRecoveryLookup struct {
	ctx     context.Context
	model   *ImportIntoModel
	imp     *entity.ImportIntoTask
	execSQL importIntoExecSQLFunc

	loaded bool
	lookup *importIntoJobStatusLookup
	err    error
}

func newLazyImportIntoRecoveryLookup(
	ctx context.Context,
	model *ImportIntoModel,
	imp *entity.ImportIntoTask,
	execSQL importIntoExecSQLFunc,
) *lazyImportIntoRecoveryLookup {
	return &lazyImportIntoRecoveryLookup{
		ctx:     ctx,
		model:   model,
		imp:     imp,
		execSQL: execSQL,
	}
}

func (l *lazyImportIntoRecoveryLookup) Get() (*importIntoJobStatusLookup, error) {
	if l == nil {
		return nil, nil
	}
	if l.loaded {
		return l.lookup, l.err
	}
	l.loaded = true
	l.lookup, l.err = syncImportIntoProgressesWithExec(l.ctx, l.model, l.imp, l.execSQL)
	return l.lookup, l.err
}

// submissionBatchState keeps one ScalingUp batch-submit round in one place so
// readers do not need to reconstruct the flow from several side maps/helpers.
type submissionBatchState struct {
	model              *ImportIntoModel
	ctx                context.Context
	imp                *entity.ImportIntoTask
	agentClient        *agent.NGRegionalClient
	execSQL            importIntoExecSQLFunc
	logger             *zap.SugaredLogger
	submitRespByTaskID map[string]*agentpb.GetImportIntoJobTaskResp
	buildFailed        map[int]struct{}
	recoveryLookup     *lazyImportIntoRecoveryLookup
}

func newSubmissionBatchState(
	ctx context.Context,
	model *ImportIntoModel,
	imp *entity.ImportIntoTask,
	agentClient *agent.NGRegionalClient,
	execSQL importIntoExecSQLFunc,
	logger *zap.SugaredLogger,
) *submissionBatchState {
	return &submissionBatchState{
		model:              model,
		ctx:                ctx,
		imp:                imp,
		agentClient:        agentClient,
		execSQL:            execSQL,
		logger:             logger,
		submitRespByTaskID: make(map[string]*agentpb.GetImportIntoJobTaskResp),
		buildFailed:        make(map[int]struct{}),
		recoveryLookup:     newLazyImportIntoRecoveryLookup(ctx, model, imp, execSQL),
	}
}

func isImportIntoSubmitCandidate(task *entity.ImportIntoSubTask) bool {
	return task.ImportIntoJobStatus == entity.ImportIntoJobStatusPending && task.ImportIntoJobID == 0
}

func collectPendingImportIntoSubmitTaskIDs(imp *entity.ImportIntoTask) []string {
	taskIDs := make([]string, 0)
	for i := range imp.ImportResults.SubTasks {
		task := &imp.ImportResults.SubTasks[i]
		// Pending-submit phase means the service has accepted this subtask locally,
		// but still needs the agent-side async result to learn the detached Job_ID.
		if isImportIntoSubmitCandidate(task) && hasImportIntoPendingSubmitTask(task) {
			taskIDs = append(taskIDs, buildImportIntoSubmitTaskID(imp.ID.Uint64(), i))
		}
	}
	return taskIDs
}

func indexImportIntoJobTaskResponses(resp *agentpb.BatchGetImportIntoJobTasksResp) map[string]*agentpb.GetImportIntoJobTaskResp {
	result := make(map[string]*agentpb.GetImportIntoJobTaskResp, len(resp.GetTasks()))
	for _, taskResp := range resp.GetTasks() {
		result[taskResp.GetTaskId()] = taskResp
	}
	return result
}

func fetchImportIntoJobTaskResponses(
	ctx context.Context,
	agentClient *agent.NGRegionalClient,
	taskIDs []string,
) (map[string]*agentpb.GetImportIntoJobTaskResp, error) {
	if len(taskIDs) == 0 {
		return map[string]*agentpb.GetImportIntoJobTaskResp{}, nil
	}
	resp, err := agentClient.BatchGetImportIntoJobTasks(ctx, taskIDs)
	if err != nil {
		return nil, err
	}
	return indexImportIntoJobTaskResponses(resp), nil
}

func (c *ImportIntoModel) batchGetImportIntoJobTaskResponsesWithFallback(
	ctx context.Context,
	agentClient *agent.NGRegionalClient,
	taskIDs []string,
	logger *zap.SugaredLogger,
) map[string]*agentpb.GetImportIntoJobTaskResp {
	respByTaskID, err := fetchImportIntoJobTaskResponses(ctx, agentClient, taskIDs)
	if err != nil {
		logger.With(zap.Strings("taskIDs", taskIDs), zap.Error(err)).Error("failed to batch get import-into job tasks")
		return buildWaitImportIntoJobTaskResponses(taskIDs)
	}
	return respByTaskID
}

func (s *submissionBatchState) pollPendingSubmitTasks() {
	pendingTaskIDs := collectPendingImportIntoSubmitTaskIDs(s.imp)
	s.submitRespByTaskID = s.model.batchGetImportIntoJobTaskResponsesWithFallback(s.ctx, s.agentClient, pendingTaskIDs, s.logger)
}

func (s *submissionBatchState) collectCreateItems() []importIntoJobSubmissionCreateItem {
	createItems := make([]importIntoJobSubmissionCreateItem, 0)
	for i := range s.imp.ImportResults.SubTasks {
		task := &s.imp.ImportResults.SubTasks[i]
		if !isImportIntoSubmitCandidate(task) {
			continue
		}
		submitTaskID := buildImportIntoSubmitTaskID(s.imp.ID.Uint64(), i)
		if s.submitRespByTaskID[submitTaskID] != nil {
			continue
		}
		if hasImportIntoPendingSubmitTask(task) {
			// Batch polling is intentionally batch-only. If a previously created
			// submit task is omitted from the batch response, keep waiting for a
			// later batch loop instead of recreating or single-polling it here.
			continue
		}
		// Create deterministic submit tasks only for subtasks that have not yet
		// entered the pending-submit phase.
		importSQL, setupSQLs, err := s.model.buildImportIntoSQL(s.ctx, task, s.imp)
		if err != nil {
			s.logger.With(
				zap.Any("subTask", sanitizeImportIntoSubTaskForLog(task)),
				zap.String("taskID", submitTaskID),
				zap.Error(err),
			).Error("failed to build import sql")
			markImportIntoControlPlaneFailed(s.imp)
			applyImportIntoSubTaskFailure(task, importIntoSubtaskMessageFromErr(err))
			s.buildFailed[i] = struct{}{}
			continue
		}
		createItems = append(createItems, importIntoJobSubmissionCreateItem{
			subTaskIndex: i,
			taskID:       submitTaskID,
			task:         task,
			request: &agentpb.BatchCreateImportIntoJobTasksReq_Task{
				TaskId:    submitTaskID,
				Sql:       importSQL,
				SetupSqls: setupSQLs,
			},
		})
	}
	return createItems
}

func (s *submissionBatchState) createSubmitTasks(createItems []importIntoJobSubmissionCreateItem) {
	if len(createItems) == 0 {
		return
	}
	passwd, decryptErr := s.model.encryptor.DecryptBase64(s.imp.TargetPassword)
	if decryptErr != nil {
		for _, item := range createItems {
			s.logger.With(
				zap.Any("subTask", sanitizeImportIntoSubTaskForLog(item.task)),
				zap.String("taskID", item.taskID),
				zap.Error(decryptErr),
			).Error("failed to decrypt password for import sql")
			markImportIntoControlPlaneFailed(s.imp)
			applyImportIntoSubTaskFailure(item.task, importIntoSubtaskMessageFromErr(decryptErr))
			s.buildFailed[item.subTaskIndex] = struct{}{}
		}
		return
	}

	requests := make([]*agentpb.BatchCreateImportIntoJobTasksReq_Task, 0, len(createItems))
	for _, item := range createItems {
		requests = append(requests, item.request)
	}
	if _, err := s.agentClient.BatchCreateImportIntoJobTasks(s.ctx, s.imp.ClusterID.Uint64(), s.imp.TargetUsername, passwd, requests); err != nil {
		s.logger.With(zap.Int("taskCount", len(requests)), zap.Error(err)).Error("failed to batch create import-into job tasks")
		return
	}
	for _, item := range createItems {
		markImportIntoPendingSubmitTask(item.task)
		s.submitRespByTaskID[item.taskID] = buildPendingImportIntoJobTaskResp(item.taskID)
	}
}

func (s *submissionBatchState) applyDecision(subTaskIndex int) importIntoJobSubmissionDecision {
	task := &s.imp.ImportResults.SubTasks[subTaskIndex]
	if task.ImportIntoJobStatus != entity.ImportIntoJobStatusPending {
		s.logger.With(zap.Any("subTask", sanitizeImportIntoSubTaskForLog(task))).Info("task is not pending")
		return importIntoJobSubmissionDecisionNoop
	}
	if task.ImportIntoJobID != 0 {
		s.logger.With(
			zap.Any("subTask", sanitizeImportIntoSubTaskForLog(task)),
			zap.Uint64("jobID", task.ImportIntoJobID),
		).Info("import job has already been submitted")
		return importIntoJobSubmissionDecisionSubmitted
	}
	if _, ok := s.buildFailed[subTaskIndex]; ok {
		return importIntoJobSubmissionDecisionFailed
	}
	submitTaskID := buildImportIntoSubmitTaskID(s.imp.ID.Uint64(), subTaskIndex)
	submitResp := s.submitRespByTaskID[submitTaskID]
	if submitResp == nil {
		return importIntoJobSubmissionDecisionWait
	}
	return s.model.resolveImportIntoJobSubmissionResponse(s.ctx, s.imp, task, submitTaskID, submitResp, s.execSQL, s.recoveryLookup, s.logger)
}

func (s *submissionBatchState) collectDecisions() []importIntoJobSubmissionDecision {
	decisions := make([]importIntoJobSubmissionDecision, len(s.imp.ImportResults.SubTasks))
	for i := range s.imp.ImportResults.SubTasks {
		decisions[i] = s.applyDecision(i)
	}
	return decisions
}

func (s *submissionBatchState) applyAllDecisions() {
	for i := range s.imp.ImportResults.SubTasks {
		s.applyDecision(i)
	}
}

func (s *submissionBatchState) sync() {
	// Each ScalingUp round first polls previously created submit tasks, then
	// creates any deterministic tasks that are still missing, and finally
	// resolves the per-subtask decisions from the combined state.
	s.pollPendingSubmitTasks()
	s.createSubmitTasks(s.collectCreateItems())
	s.applyAllDecisions()
}

func (c *ImportIntoModel) syncImportIntoJobSubmissionsBatch(
	ctx context.Context,
	imp *entity.ImportIntoTask,
	agentClient *agent.NGRegionalClient,
	execSQL importIntoExecSQLFunc,
	logger *zap.SugaredLogger,
) {
	newSubmissionBatchState(ctx, c, imp, agentClient, execSQL, logger).sync()
}

func (c *ImportIntoModel) onImportIntoScalingUp(ctx context.Context, imp *entity.ImportIntoTask) (err error) {
	logger := c.sugaredTraceIDLogger(ctx, zap.Uint64("importID", imp.ID.Uint64()), zap.Uint64("clusterID", imp.ClusterID.Uint64()))
	logger.Info("begin to scaling IMPORT INTO task")
	agentClient, err := c.newNGRegionalClientByClusterID(ctx, imp.ClusterID.Uint64())
	if err != nil {
		return err
	}
	execSQLExecutor, err := c.newImportIntoSQLExecutor(ctx, imp, agentClient)
	if err != nil {
		return err
	}
	prevStatus := imp.Status
	defer func() {
		err = c.persistImportIntoTaskState(ctx, logger, imp, err, prevStatus)
	}()
	if c.isImportIntoTiKVCapacityReservationEnabled(imp.TenantID.Uint64(), imp.ProjectID.Uint64(), imp.ClusterID.Uint64()) {
		reservationReady, err := c.syncImportIntoTiKVReservation(ctx, imp, agentClient, logger)
		if err != nil {
			return err
		}
		if !reservationReady {
			return nil
		}
	}
	// syncImportIntoJobSubmissionsBatch updates per-subtask submission outcomes;
	// resolveImportIntoSubmissionState derives the task-level aggregate state for
	// this ScalingUp round from those persisted subtask fields.
	c.syncImportIntoJobSubmissionsBatch(ctx, imp, agentClient, execSQLExecutor.execute, logger)
	switch resolveImportIntoSubmissionState(imp) {
	case importIntoSubmissionStateWaiting:
		logger.Info("waiting for import job submission to finish")
		return nil
	case importIntoSubmissionStateFailed:
		imp.Status = entity.ImportIntoStatusFailed
		imp.Message = "All table(s) completed, but not all were successful"
		return nil
	default:
		imp.Status = entity.ImportIntoStatusImporting
		return nil
	}
}

type TaskStatus struct {
	Phase                        string    `json:"phase"`
	Status                       string    `json:"status"`
	ImportedRows                 int64     `json:"imported_rows"`
	ResultMessage                string    `json:"result_message"`
	CreateTime                   time.Time `json:"created_time"`
	StartTime                    time.Time `json:"started_time"`
	EndTime                      time.Time `json:"end_time"`
	CurrentStep                  string    `json:"current_step"`
	CurrentStepTotalSize         string    `json:"current_step_total_size"`
	CurrentStepProgressPercent   int32     `json:"current_step_progress_percent"`
	CurrentStepETA               string    `json:"current_step_eta"`
	CurrentStepProcessedDataSize string    `json:"current_step_processed_data_size"`
	CurrentStepSpeed             string    `json:"current_step_speed"`
}

type importIntoProgressSyncHooks struct {
	syncProgress func(subTask *entity.ImportIntoSubTask) (*TaskStatus, error)
}

func applyImportIntoProgressStatus(subTask *entity.ImportIntoSubTask, status *TaskStatus) {
	subTask.ImportIntoPhase = status.Phase
	subTask.Message = status.ResultMessage
	setImportIntoSubTaskStatus(subTask, status.Status)
	subTask.CurrentStep = status.CurrentStep
	subTask.CurrentStepTotalSize = status.CurrentStepTotalSize
	subTask.CurrentStepProgressPercent = status.CurrentStepProgressPercent
	subTask.CurrentStepETA = status.CurrentStepETA
	subTask.CurrentStepProcessedDataSize = status.CurrentStepProcessedDataSize
	subTask.CurrentStepSpeed = status.CurrentStepSpeed
	subTask.ImportedRows = status.ImportedRows
	if !status.StartTime.IsZero() {
		subTask.ImportStartedAt = &status.StartTime
	}
	if !status.EndTime.IsZero() {
		subTask.ImportFinishedAt = &status.EndTime
	}
}

func runImportIntoImportingProgressSync(
	imp *entity.ImportIntoTask,
	logger *zap.SugaredLogger,
	hooks importIntoProgressSyncHooks,
) {
	var syncErrCount int
	var missingJobIDCount int
	for i := 0; i != len(imp.ImportResults.SubTasks); i++ {
		subTask := &imp.ImportResults.SubTasks[i]
		switch subTask.ImportIntoJobStatus {
		case entity.ImportIntoJobStatusFinished, entity.ImportIntoJobStatusCanceled, entity.ImportIntoJobStatusFailed:
			continue
		}
		if subTask.ImportIntoJobID == 0 {
			missingJobIDCount++
			continue
		}
		status, err := hooks.syncProgress(subTask)
		if err != nil {
			syncErrCount++
			logger.Errorw("failed to sync import into progress",
				"database", subTask.DataBase,
				"table", subTask.TableName,
				"jobID", subTask.ImportIntoJobID,
				"error", err,
			)
			continue
		}
		applyImportIntoProgressStatus(subTask, status)
		if status.Status == entity.ImportIntoJobStatusCanceled {
			logger.Errorw("job has been canceled on importing",
				"clusterID", imp.ClusterID.Uint64(),
				"jobID", subTask.ImportIntoJobID,
			)
		}
	}

	var pendingCount int
	var runningCount int
	var finishedCount int
	var failedCount int
	for i := 0; i != len(imp.ImportResults.SubTasks); i++ {
		switch imp.ImportResults.SubTasks[i].ImportIntoJobStatus {
		case entity.ImportIntoJobStatusPending:
			pendingCount++
		case entity.ImportIntoJobStatusRunning:
			runningCount++
		case entity.ImportIntoJobStatusFinished:
			finishedCount++
		case entity.ImportIntoJobStatusFailed:
			failedCount++
		}
	}
	logger.Infow("synced import into progress summary",
		"totalSubtasks", len(imp.ImportResults.SubTasks),
		"pendingSubtasks", pendingCount,
		"runningSubtasks", runningCount,
		"finishedSubtasks", finishedCount,
		"failedSubtasks", failedCount,
		"missingJobIDSubtasks", missingJobIDCount,
		"syncErrorSubtasks", syncErrCount,
	)
}

func markCancelingSubtasks(imp *entity.ImportIntoTask) {
	for i := 0; i != len(imp.ImportResults.SubTasks); i++ {
		subTask := &imp.ImportResults.SubTasks[i]
		switch subTask.ImportIntoJobStatus {
		case entity.ImportIntoJobStatusFinished, entity.ImportIntoJobStatusCanceled, entity.ImportIntoJobStatusFailed:
			continue
		case entity.ImportIntoJobStatusPending, entity.ImportIntoJobStatusRunning:
			if subTask.ImportIntoJobID == 0 {
				if hasImportIntoPendingSubmitTask(subTask) {
					continue
				}
				clearImportIntoPendingSubmitTask(subTask)
				setImportIntoSubTaskStatus(subTask, entity.ImportIntoJobStatusCanceled)
				continue
			}
		}
	}
}

func syncCancelingSubtasks(
	imp *entity.ImportIntoTask,
	logger *zap.Logger,
	syncProgress func(subTask *entity.ImportIntoSubTask) (*TaskStatus, error),
) {
	for i := 0; i != len(imp.ImportResults.SubTasks); i++ {
		subTask := &imp.ImportResults.SubTasks[i]
		switch subTask.ImportIntoJobStatus {
		case entity.ImportIntoJobStatusFinished, entity.ImportIntoJobStatusCanceled, entity.ImportIntoJobStatusFailed:
			continue
		case entity.ImportIntoJobStatusPending, entity.ImportIntoJobStatusRunning:
			if subTask.ImportIntoJobID == 0 {
				continue
			}
			status, err := syncProgress(subTask)
			if err != nil {
				logger.With(
					zap.String("database", subTask.DataBase),
					zap.String("table", subTask.TableName),
					zap.Uint64("jobID", subTask.ImportIntoJobID),
					zap.Error(err),
				).Error("failed to sync import into progress")
				continue
			}
			applyImportIntoProgressStatus(subTask, status)
		}
	}
}

func shouldSyncImportIntoCancelingProgressBatch(imp *entity.ImportIntoTask) bool {
	for i := range imp.ImportResults.SubTasks {
		subTask := &imp.ImportResults.SubTasks[i]
		switch subTask.ImportIntoJobStatus {
		case entity.ImportIntoJobStatusFinished, entity.ImportIntoJobStatusCanceled, entity.ImportIntoJobStatusFailed:
			continue
		}
		if subTask.ImportIntoJobID != 0 {
			return true
		}
	}
	return false
}

func (c *ImportIntoModel) syncPendingImportIntoSubmitTasksForCancel(
	ctx context.Context,
	imp *entity.ImportIntoTask,
	agentClient *agent.NGRegionalClient,
	logger *zap.Logger,
) error {
	taskIDs := make([]string, 0)
	taskByID := make(map[string]*entity.ImportIntoSubTask)
	for i := range imp.ImportResults.SubTasks {
		subTask := &imp.ImportResults.SubTasks[i]
		if !hasImportIntoPendingSubmitTask(subTask) {
			continue
		}
		// Canceling may race with async submit. Poll the submit task first so a
		// just-created Job_ID is canceled instead of being marked locally canceled.
		taskID := buildImportIntoSubmitTaskID(imp.ID.Uint64(), i)
		taskIDs = append(taskIDs, taskID)
		taskByID[taskID] = subTask
	}
	if len(taskIDs) == 0 {
		return nil
	}
	respByTaskID := c.batchGetImportIntoJobTaskResponsesWithFallback(ctx, agentClient, taskIDs, logger.Sugar())
	for _, taskID := range taskIDs {
		taskResp := respByTaskID[taskID]
		if taskResp == nil {
			// Keep omitted pending-submit tasks unresolved for the next batch loop.
			continue
		}
		subTask := taskByID[taskID]
		if subTask == nil {
			continue
		}
		switch resolveImportIntoJobTaskState(taskResp) {
		case importIntoAsyncTaskStateWait:
			continue
		case importIntoAsyncTaskStateSucceeded:
			applyImportIntoJobTaskResult(subTask, taskResp)
		case importIntoAsyncTaskStateTerminalFailure:
			applyImportIntoSubTaskFailure(subTask, buildImportIntoJobTaskFailureMessage(taskResp))
		default:
			applyImportIntoSubTaskFailure(subTask, "Internal Server Error")
		}
	}
	return nil
}

func (c *ImportIntoModel) batchCancelImportIntoJobs(
	ctx context.Context,
	imp *entity.ImportIntoTask,
	agentClient *agent.NGRegionalClient,
	logger *zap.Logger,
) error {
	jobIDs := make([]uint64, 0)
	for i := range imp.ImportResults.SubTasks {
		subTask := &imp.ImportResults.SubTasks[i]
		switch subTask.ImportIntoJobStatus {
		case entity.ImportIntoJobStatusPending, entity.ImportIntoJobStatusRunning:
			if subTask.ImportIntoJobID != 0 {
				jobIDs = append(jobIDs, subTask.ImportIntoJobID)
			}
		}
	}
	if len(jobIDs) == 0 {
		return nil
	}
	passwd, err := c.encryptor.DecryptBase64(imp.TargetPassword)
	if err != nil {
		return err
	}
	resp, cancelErr := agentClient.BatchCancelImportIntoJobs(ctx, imp.ClusterID.Uint64(), imp.TargetUsername, passwd, jobIDs)
	if cancelErr != nil {
		return cancelErr
	}
	for _, result := range resp.GetResults() {
		if result.GetErrorMsg() == "" {
			continue
		}
		logger.With(
			zap.Uint64("jobID", result.GetJobId()),
			zap.String("errorMsg", result.GetErrorMsg()),
		).Error("failed to execute cancel SQL")
	}
	return nil
}

func syncImportIntoProgressesWithExec(
	ctx context.Context,
	c *ImportIntoModel,
	imp *entity.ImportIntoTask,
	execSQL importIntoExecSQLFunc,
) (*importIntoJobStatusLookup, error) {
	logger := c.sugaredTraceIDLogger(ctx)
	if imp.GroupKey == "" {
		return nil, pingErr.New("import into group key is empty")
	}
	showSQL := buildShowImportJobsByGroupKeySQL(imp.GroupKey)
	sqlResp, sqlErr := execSQL(ctx, imp, showSQL, nil, 18)
	if sqlErr != nil {
		logger.Error("failed to query import jobs status", zap.Error(sqlErr), zap.String("sql", showSQL))
		return nil, pingErr.Annotate(sqlErr, "failed to query import jobs status")
	}
	lookup, err := findImportIntoJobStatusesFromResp(sqlResp)
	if err != nil {
		logger.With(zap.Any("resp", sqlResp), zap.Error(err)).Error("failed to parse import jobs status")
		return nil, err
	}
	return lookup, nil
}

func syncImportIntoProgressesWithExecutor(
	ctx context.Context,
	c *ImportIntoModel,
	imp *entity.ImportIntoTask,
	executor *importIntoSQLExecutor,
) (*importIntoJobStatusLookup, error) {
	if executor == nil {
		return nil, pingErr.New("import into sql executor is nil")
	}
	return syncImportIntoProgressesWithExec(ctx, c, imp, executor.execute)
}

func syncImportIntoProgressWithLookup(subTask *entity.ImportIntoSubTask, lookup *importIntoJobStatusLookup) (*TaskStatus, error) {
	status, ok := lookup.findByJobID(subTask.ImportIntoJobID)
	if ok {
		return status, nil
	}
	targetTable := importIntoTargetTableName(subTask)
	if _, hasTableStatuses := lookup.findByTargetTable(targetTable); hasTableStatuses {
		return nil, fmt.Errorf("import job %d for target table %s not found in batch result", subTask.ImportIntoJobID, targetTable)
	}
	return nil, fmt.Errorf("import job %d not found in batch result", subTask.ImportIntoJobID)
}

func parseImportIntoCurrentStepProgressPercent(raw string) (uint64, error) {
	raw = strings.TrimSpace(raw)
	switch raw {
	case "", "N/A", "/":
		return 0, nil
	default:
		return strconv.ParseUint(raw, 10, 64)
	}
}

func getColumnValue(resp *agentpb.ExecSQLResp, columnName string) (*agentpb.ColumnValue, error) {
	result, err := getColumnValues(resp, columnName)
	if err != nil {
		return nil, err
	}
	return result[columnName], nil
}

func getColumnValues(resp *agentpb.ExecSQLResp, columnNames ...string) (map[string]*agentpb.ColumnValue, error) {
	if resp == nil || len(resp.Rows) == 0 || len(resp.Rows[0].TypedColumnValues) == 0 || len(resp.ColumnNames) != len(resp.Rows[0].TypedColumnValues) {
		return nil, pingErr.New("failed to exec sql")
	}
	indexMap := make(map[string]int, len(columnNames))
	for _, col := range columnNames {
		indexMap[col] = -1
	}
	for i, cname := range resp.ColumnNames {
		if _, ok := indexMap[cname]; ok {
			indexMap[cname] = i
		}
	}
	result := make(map[string]*agentpb.ColumnValue, len(columnNames))
	for _, col := range columnNames {
		if indexMap[col] == -1 {
			return nil, fmt.Errorf("not found the column %s", col)
		}
		result[col] = resp.Rows[0].TypedColumnValues[indexMap[col]]
	}
	return result, nil
}

// onImportIntoImporting monitors the progress of all active IMPORT INTO jobs with one batched SHOW IMPORT JOBS query.
func (c *ImportIntoModel) onImportIntoImporting(ctx context.Context, imp *entity.ImportIntoTask) error {
	logger := c.sugaredTraceIDLogger(ctx, zap.Uint64("importID", imp.ID.Uint64()))
	logger.Info("monitoring IMPORT INTO task progress with batched SHOW IMPORT JOBS query")
	prevStatus := imp.Status

	agentClient, err := c.newNGRegionalClientByClusterID(ctx, imp.ClusterID.Uint64())
	if err != nil {
		return err
	}
	executor, err := c.newImportIntoSQLExecutor(ctx, imp, agentClient)
	if err != nil {
		return err
	}
	statusLookup, err := syncImportIntoProgressesWithExecutor(ctx, c, imp, executor)
	if err != nil {
		return err
	}
	runImportIntoImportingProgressSync(imp, logger, importIntoProgressSyncHooks{
		syncProgress: func(subTask *entity.ImportIntoSubTask) (*TaskStatus, error) {
			return syncImportIntoProgressWithLookup(subTask, statusLookup)
		},
	})

	switch resolveImportIntoImportingDecision(imp) {
	case importIntoImportingDecisionCompleted:
		applyImportIntoCompletion(imp)
		now := time.Now()
		imp.FinishedAt = &now
		logger.Info("Clear user password on task completed", zap.Uint64("clusterID", imp.ClusterID.Uint64()), zap.Uint64("impId", imp.ID.Uint64()))
		c.releaseTiKVCapacityReservation(ctx, logger, agentClient, imp.TenantID.Uint64(), imp.ProjectID.Uint64(), imp.ClusterID.Uint64(), imp.ID.Uint64())
	case importIntoImportingDecisionFailed:
		applyImportIntoSQLFailure(imp, "All table(s) completed, but not all were successful")
		now := time.Now()
		imp.FinishedAt = &now
		logger.Info("Clear user password on task completed", zap.Uint64("clusterID", imp.ClusterID.Uint64()), zap.Uint64("impId", imp.ID.Uint64()))
		imp.TargetPassword = ""
		c.releaseTiKVCapacityReservation(ctx, logger, agentClient, imp.TenantID.Uint64(), imp.ProjectID.Uint64(), imp.ClusterID.Uint64(), imp.ID.Uint64())
	}
	txErr := db.TransactionWithContext(ctx, func(tx *gorm.DB) error {
		updateFields := []entity.ImportIntoTaskDBSchemaField{
			entity.ImportIntoTaskDBSchema.ImportResults,
			entity.ImportIntoTaskDBSchema.Message,
			entity.ImportIntoTaskDBSchema.TargetPassword,
		}
		if imp.Status == entity.ImportIntoStatusCompleted || imp.Status == entity.ImportIntoStatusFailed {
			updateFields = append(updateFields,
				entity.ImportIntoTaskDBSchema.Status,
				entity.ImportIntoTaskDBSchema.FinishedAt,
				entity.ImportIntoTaskDBSchema.SourceURI,
			)
		}
		return imp.Update(tx, updateFields...)
	})
	if txErr != nil {
		logger.Error("failed to update overall import progress in DB", zap.Error(txErr))
		return pingErr.AddStack(txErr)
	}
	publishImportIntoEventIfStatusChanged(ctx, c.eventPublisher, imp, prevStatus)
	if isTerminalImportIntoStatus(imp.Status) {
		if cleanupErr := c.cleanupImportClusterJob(ctx, logger, imp); cleanupErr != nil {
			logger.With(zap.Error(cleanupErr)).Error("failed to clean up Import Into custom status")
		}
	}
	return nil
}

// onImportIntoCanceling handles the cancellation of multiple IMPORT INTO jobs.
func (c *ImportIntoModel) onImportIntoCanceling(ctx context.Context, imp *entity.ImportIntoTask) (err error) {
	logger := c.traceIDLogger(ctx, zap.Uint64("importID", imp.ID.Uint64()))
	logger.Info("begin to cancel IMPORT INTO task and its sub-jobs")
	// Cancel path only needs target credentials for batch cancel; strip source URIs from DB immediately.
	cleanImportIntoStoredURIs(imp)
	agentClient, err := c.newNGRegionalClientByClusterID(ctx, imp.ClusterID.Uint64())
	if err != nil {
		return err
	}
	if err := c.syncPendingImportIntoSubmitTasksForCancel(ctx, imp, agentClient, logger); err != nil {
		logger.With(zap.Error(err)).Error("failed to prefetch pending import-into submit tasks while canceling")
	}
	if err := c.batchCancelImportIntoJobs(ctx, imp, agentClient, logger); err != nil {
		logger.With(zap.Error(err)).Error("failed to batch cancel import-into jobs")
	}
	markCancelingSubtasks(imp)
	if shouldSyncImportIntoCancelingProgressBatch(imp) {
		executor, err := c.newImportIntoSQLExecutor(ctx, imp, agentClient)
		if err != nil {
			return err
		}
		statusLookup, err := syncImportIntoProgressesWithExecutor(ctx, c, imp, executor)
		if err != nil {
			return err
		}
		syncCancelingSubtasks(imp, logger, func(subTask *entity.ImportIntoSubTask) (*TaskStatus, error) {
			return syncImportIntoProgressWithLookup(subTask, statusLookup)
		})
	}
	canceled := resolveImportIntoCancelingDecision(imp) == importIntoCancelingDecisionCanceled
	prevStatus := imp.Status
	if canceled {
		imp.Status = entity.ImportIntoStatusCanceled
		now := time.Now()
		imp.FinishedAt = &now
		imp.TargetPassword = ""
		logger.Info("Clear user password on task canceled", zap.Uint64("clusterID", imp.ClusterID.Uint64()), zap.Uint64("impId", imp.ID.Uint64()))
		c.releaseTiKVCapacityReservation(ctx, logger.Sugar(), agentClient, imp.TenantID.Uint64(), imp.ProjectID.Uint64(), imp.ClusterID.Uint64(), imp.ID.Uint64())
	}
	txErr := db.TransactionWithContext(ctx, func(tx *gorm.DB) error {
		updateFields := []entity.ImportIntoTaskDBSchemaField{
			entity.ImportIntoTaskDBSchema.ImportResults,
			entity.ImportIntoTaskDBSchema.Message,
			entity.ImportIntoTaskDBSchema.SourceURI,
		}
		if canceled {
			updateFields = append(updateFields,
				entity.ImportIntoTaskDBSchema.Status,
				entity.ImportIntoTaskDBSchema.TargetPassword,
				entity.ImportIntoTaskDBSchema.FinishedAt,
			)
		}
		return imp.Update(tx, updateFields...)
	})
	if txErr != nil {
		logger.With(zap.Error(txErr)).Error("failed to update import into on canceling")
		return nil
	}
	publishImportIntoEventIfStatusChanged(ctx, c.eventPublisher, imp, prevStatus)
	if canceled {
		if cleanupErr := c.cleanupImportClusterJob(ctx, logger.Sugar(), imp); cleanupErr != nil {
			logger.With(zap.Error(cleanupErr)).Error("failed to clean up Import Into custom status")
		}
	}
	return nil
}