package importadmin

import (
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"net/url"
	"regexp"
	"slices"
	"strconv"
	"strings"
	"time"

	"github.com/avast/retry-go/v4"
	pinerrors "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"
	pimport "github.com/tidbcloud/dataflow-service-ng/app/models/import_into"
	"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/bizutils"
	"github.com/tidbcloud/dataflow-service-ng/internal/config"
	"github.com/tidbcloud/dataflow-service-ng/internal/db"
	"github.com/tidbcloud/dataflow-service-ng/pkg/account/accountdb"
	"github.com/tidbcloud/dataflow-service-ng/pkg/agent"
	"github.com/tidbcloud/dataflow-service-ng/pkg/azblob"
	"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"
	pkglog "github.com/tidbcloud/dataflow-service-ng/pkg/log"
	agentpb "github.com/tidbcloud/dataflow-service/dataflow-agent-ng/idl/pbgen/proto/ng_agent"
	"github.com/tidbcloud/pkgv2/bizerr"
	"github.com/tidbcloud/pkgv2/import/cloud_storage"
	mgmt "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
	"github.com/tidbcloud/tiinf/sdk/auditlog"
	"go.uber.org/zap"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/types/known/emptypb"
	"google.golang.org/protobuf/types/known/timestamppb"
	"gorm.io/gorm"
)

// pemChangedError is a custom error indicating that the PEM key has changed and a retry is needed.
type pemChangedError struct{}

func (e *pemChangedError) Error() string {
	return "PEM has changed, need to retry"
}

const checkDatabaseConnectivityFallbackMsg = "failed to check database connectivity"

const importClusterJobCleanupTimeout = 10 * time.Second

var iamRoleARNPattern = regexp.MustCompile(`^arn:aws(?:-[a-z0-9-]+)?:iam::[0-9]{12}:role/(?:[\x21-\x2E\x30-\x7E]+/)*[A-Za-z0-9_+=,.@-]+$`)

type ImportIntoService struct {
	importadminpb.UnimplementedImportIntoServiceServer

	mgmtSvcGrpcClient        grpcli.MgmtSvcGrpcClient
	encryptor                encryptor_wrapper.Encryptor
	ngRegionalClientProvider agent.NGRegionalClientProvider
	eventPublisher           center.EventCenterPublisher
	clusterJob               cluster_job.Model
}

type checkImportSourceTaskRequest struct {
	clusterID  uint64
	sourceURL  *url.URL
	sourceType importadminpb.ImportSourceType_Enum
	inMetas    []*agentpb.GetImportFilesReq_TableMeta
	tableMetas entity.ImportIntoPrecheckTableMetas
}

type checkImportSourceResult struct {
	success             bool
	errorMsg            *string
	tableMetas          []*importadminpb.ImportTargetTableInfo
	totalScannedObjects int32
	status              importadminpb.CheckImportSourceStatus_Enum
}

const importIntoLegacyWriteFrozenErrMsg = "import-into write path is frozen during migration"

func NewImportIntoService(
	mgmtSvcGrpcClient grpcli.MgmtSvcGrpcClient,
	encryptor encryptor_wrapper.Encryptor,
	ngRegionalClientProvider agent.NGRegionalClientProvider,
	eventPublisher center.EventCenterPublisher,
	clusterJob cluster_job.Model,
) *ImportIntoService {
	return &ImportIntoService{
		mgmtSvcGrpcClient:        mgmtSvcGrpcClient,
		encryptor:                encryptor,
		ngRegionalClientProvider: ngRegionalClientProvider,
		eventPublisher:           eventPublisher,
		clusterJob:               clusterJob,
	}
}

type scopedImport struct {
	UserID    uint64
	TenantID  uint64
	ProjectID uint64
	ClusterID uint64

	Cluster *mgmt.Cluster
	User    *accountdb.User
}

func (s *ImportIntoService) ensureImportIntoWriteEnabled(ctx context.Context, cid any) (*scopedImport, bizerr.BizError) {
	var clusterID uint64
	switch cid := cid.(type) {
	case uint64:
		clusterID = cid
	case string:
		var err error
		clusterID, err = strconv.ParseUint(cid, 10, 64)
		if err != nil {
			return nil, bizerr.New(bizerr.ErrCodeParams, "invalid cluster ID")
		}
	default:
		return nil, bizerr.New(bizerr.ErrCodeParams, "invalid cluster ID type")
	}

	userInCtx, uerr := bizutils.GetUserFromContext(ctx)
	if uerr != nil {
		return nil, bizerr.New(bizerr.ErrCodePermission, "user not in context")
	}
	cluster, err := s.mgmtSvcGrpcClient.GetClusterByID(ctx, clusterID)
	if err != nil {
		return nil, bizerr.New(bizerr.ErrCodeParams, "failed to get cluster by ID")
	}
	if _, enabled := config.Global().ImportConfig.WriteEnabled.Enabled(config.FeatureGateQueryParams{
		Provider:  cluster.GetCloudProvider(),
		Region:    cluster.GetRegion(),
		TenantID:  userInCtx.TenantID,
		ProjectID: cluster.GetProjectId(),
		ClusterID: cluster.GetClusterId(),
	}); !enabled {
		return nil, bizerr.New(bizerr.ErrCodePermission, importIntoLegacyWriteFrozenErrMsg)
	}

	return &scopedImport{
		UserID:    userInCtx.ID,
		TenantID:  userInCtx.TenantID,
		ProjectID: cluster.GetProjectId(),
		ClusterID: cluster.GetClusterId(),
		Cluster:   cluster,
		User:      userInCtx,
	}, nil
}

func payloadErrorOrFallback(payloadErr, fallback string) string {
	payloadErr = strings.TrimSpace(payloadErr)
	if payloadErr != "" {
		return payloadErr
	}
	return fallback
}

func checkDatabaseConnectivityFailureResp(payloadErr string) *importadminpb.CheckDatabaseConnectivityResp {
	errMsg := payloadErrorOrFallback(payloadErr, checkDatabaseConnectivityFallbackMsg)
	return &importadminpb.CheckDatabaseConnectivityResp{
		Success:  false,
		ErrorMsg: &errMsg,
	}
}

func validateS3SourceAuth(source *importadminpb.S3Source) error {
	if source == nil {
		return pinerrors.New("S3 source is required")
	}

	switch source.GetAuthType() {
	case importadminpb.ImportS3AuthType_ENUM_ROLE_ARN:
		if _, ok := source.GetAuth().(*importadminpb.S3Source_RoleArn); !ok {
			return pinerrors.New("S3 auth_type ROLE_ARN requires role_arn")
		}
		if !iamRoleARNPattern.MatchString(source.GetRoleArn()) {
			return pinerrors.New("invalid AWS IAM role ARN")
		}
	case importadminpb.ImportS3AuthType_ENUM_ACCESS_KEY:
		if _, ok := source.GetAuth().(*importadminpb.S3Source_AccessKey_); !ok {
			return pinerrors.New("S3 auth_type ACCESS_KEY requires access_key")
		}
	default:
		return pinerrors.New("unsupported S3 auth_type")
	}

	return nil
}

func validateImportSourceAuth(source *importadminpb.ImportSource) error {
	if source.GetType() != importadminpb.ImportSourceType_ENUM_S3 {
		return nil
	}
	return validateS3SourceAuth(source.GetS3())
}

func FromSourceURI(sourceUri string) (*importadminpb.ImportSource, error) {
	u, err := url.Parse(sourceUri)
	if err != nil {
		return nil, pinerrors.Annotate(err, "failed to parse source url")
	}
	queryParamVals := u.Query()
	u.RawQuery = ""
	result := &importadminpb.ImportSource{Type: importadminpb.ImportSourceType_ENUM_UNSPECIFIED, Uri: u.String()}
	switch u.Scheme {
	case "s3":
		s3 := &importadminpb.S3Source{Uri: u.String()}
		result.Source = &importadminpb.ImportSource_S3{S3: s3}
		result.Type = importadminpb.ImportSourceType_ENUM_S3
		if len(queryParamVals.Get("role-arn")) > 0 {
			s3.AuthType = importadminpb.ImportS3AuthType_ENUM_ROLE_ARN
			s3.Auth = &importadminpb.S3Source_RoleArn{RoleArn: queryParamVals.Get("role-arn")}
		} else {
			s3.AuthType = importadminpb.ImportS3AuthType_ENUM_ACCESS_KEY
			s3.Auth = &importadminpb.S3Source_AccessKey_{AccessKey: &importadminpb.S3Source_AccessKey{Id: queryParamVals.Get("access-key")}}
		}
	case "oss":
		oss := &importadminpb.OSSSource{Uri: u.String()}
		result.Source = &importadminpb.ImportSource_Oss{Oss: oss}
		result.Type = importadminpb.ImportSourceType_ENUM_OSS
		oss.AuthType = importadminpb.ImportOSSAuthType_ENUM_ACCESS_KEY
		oss.Auth = &importadminpb.OSSSource_AccessKey_{AccessKey: &importadminpb.OSSSource_AccessKey{Id: queryParamVals.Get("access-key")}}
	case "azure":
		azureBlob := &importadminpb.AzureBlobSource{Uri: u.String()}
		result.Source = &importadminpb.ImportSource_AzureBlob{AzureBlob: azureBlob}
		result.Type = importadminpb.ImportSourceType_ENUM_AZURE_BLOB
	}
	return result, nil
}

func mergeQueryWithRefPrecedence(baseQuery, refQuery url.Values) url.Values {
	q := make(url.Values, len(baseQuery)+len(refQuery))
	for key, values := range baseQuery {
		q[key] = append([]string(nil), values...)
	}
	for key, values := range refQuery {
		q[key] = append([]string(nil), values...)
	}
	return q
}

func JoinSourceURI(baseURI, ref string) (string, error) {
	baseURL, err := url.Parse(baseURI)
	if err != nil {
		return "", err
	}
	refURL, err := url.Parse(ref)
	if err != nil {
		return "", err
	}

	baseQuery := baseURL.Query()
	refQuery := refURL.Query()
	var res *url.URL
	if refURL.Scheme == "" && refURL.Host == "" {
		res = baseURL.JoinPath(strings.TrimPrefix(refURL.Path, "/"))
	} else {
		res = baseURL.ResolveReference(refURL)
	}

	res.RawQuery = mergeQueryWithRefPrecedence(baseQuery, refQuery).Encode()
	return res.String(), nil
}

func JoinS3Uri(a, b string) (string, error) {
	urlA, err := url.Parse(a)
	if err != nil {
		return "", err
	}
	urlB, err := url.Parse(b)
	if err != nil {
		return "", err
	}
	var res *url.URL
	if (!strings.HasSuffix(urlA.Path, "/")) && urlB.Scheme == "" {
		res = urlA.JoinPath(urlB.Path)
		// a := "s3://bucket"; b := "/tz*.csv"
		q := res.Query()
		for key, values := range urlB.Query() {
			for _, val := range values {
				q.Add(key, val)
			}
		}
		res.RawQuery = q.Encode()
	} else {
		// a := "s3://bucket"; b := "s3://bucket/tz*.csv"
		// a := "s3://bucket/"; b := "tz*.csv"
		// a := "s3://bucket/"; b := "/tz*.csv"
		res = urlA.ResolveReference(urlB)
		q := res.Query()
		for key, values := range urlA.Query() {
			for _, val := range values {
				q.Add(key, val)
			}
		}
		res.RawQuery = q.Encode()
	}
	return res.String(), nil
}

func buildCheckImportSourceSuccessResp(
	req *importadminpb.CheckImportSourceReq,
	resp *agentpb.GetImportFilesResp,
	rewriteToOSSScheme bool,
) (*importadminpb.CheckImportSourceResp, error) {
	result, err := buildCheckImportSourceSuccessResult(len(req.GetTableMetas()) != 0, resp, rewriteToOSSScheme)
	if err != nil {
		return nil, err
	}
	if result.totalScannedObjects >= 1000 && len(req.GetTableMetas()) == 0 {
		result.status = importadminpb.CheckImportSourceStatus_ENUM_TRUNCATED
	}
	return &importadminpb.CheckImportSourceResp{
		Success:             result.success,
		ErrorMsg:            result.errorMsg,
		TableMetas:          result.tableMetas,
		TotalScannedObjects: result.totalScannedObjects,
		Status:              result.status,
	}, nil
}

// Normalize wildcard paths for API responses by dropping query params and restoring the OSS scheme when needed.
func formatCheckImportSourceWildcardPath(wildcardPath string, rewriteToOSSScheme bool) (string, error) {
	parsedWildPath, err := url.Parse(wildcardPath)
	if err != nil {
		return "", err
	}
	parsedWildPath.RawQuery = ""
	if rewriteToOSSScheme && wildcardPath != "" {
		parsedWildPath.Scheme = "oss"
	}
	return parsedWildPath.String(), nil
}

func buildCheckImportSourceSuccessResult(
	isExplicitTableCheck bool,
	resp *agentpb.GetImportFilesResp,
	rewriteToOSSScheme bool,
) (*checkImportSourceResult, error) {
	tables, totalScanned, errMsg, err := buildCheckImportSourceTableMetas(isExplicitTableCheck, resp.GetTableMeta(), rewriteToOSSScheme)
	if err != nil {
		return nil, err
	}
	return &checkImportSourceResult{
		success:             true,
		errorMsg:            errMsg,
		tableMetas:          tables,
		totalScannedObjects: int32(totalScanned),
		status:              importadminpb.CheckImportSourceStatus_ENUM_SUCCESS,
	}, nil
}

func buildCheckImportSourceTableMetas(
	isExplicitTableCheck bool,
	respTableMetas []*agentpb.GetImportFilesResp_TableMeta,
	rewriteToOSSScheme bool,
) ([]*importadminpb.ImportTargetTableInfo, int, *string, error) {
	totalScanned := 0
	tables := make([]*importadminpb.ImportTargetTableInfo, 0, len(respTableMetas))
	var errMsg *string
	for _, tbl := range respTableMetas {
		formattedWildcardPath := ""
		if len(tbl.GetDataFiles()) != 0 {
			var err error
			formattedWildcardPath, err = formatCheckImportSourceWildcardPath(tbl.GetWildcardPath(), rewriteToOSSScheme)
			if err != nil {
				return nil, 0, nil, err
			}
		}
		tables = append(tables, &importadminpb.ImportTargetTableInfo{
			Database:      tbl.GetDatabase(),
			Table:         tbl.GetTable(),
			FilesWildcard: formattedWildcardPath,
		})
		if len(tbl.GetDataFiles()) == 0 && isExplicitTableCheck {
			// An empty dataset is a completed check result, not an execution failure.
			msg := fmt.Sprintf("empty data set: %s", tbl.GetWildcardPath())
			errMsg = &msg
		}
		totalScanned += len(tbl.GetDataFiles())
		if tbl.GetSchemaFile() != "" {
			totalScanned++
		}
	}
	return tables, totalScanned, errMsg, nil
}

func buildGetCheckImportSourceTaskSuccessResp(
	taskID string,
	task *entity.ImportIntoPrecheckTask,
	resp *agentpb.GetImportFilesResp,
	rewriteToOSSScheme bool,
) (*importadminpb.GetCheckImportSourceTaskResp, error) {
	result, err := buildCheckImportSourceSuccessResult(len(task.TableMetas) != 0, resp, rewriteToOSSScheme)
	if err != nil {
		return nil, err
	}
	return &importadminpb.GetCheckImportSourceTaskResp{
		TaskId:              taskID,
		Success:             result.success,
		ErrorMsg:            result.errorMsg,
		TableMetas:          result.tableMetas,
		TotalScannedObjects: result.totalScannedObjects,
		Status:              result.status,
	}, nil
}

func (s *ImportIntoService) buildCheckImportSourceTaskRequest(
	ctx context.Context,
	tidbID string,
	source *importadminpb.ImportSource,
	tableMetas []*importadminpb.ImportTargetTableInfo,
) (*checkImportSourceTaskRequest, error) {
	clusterID, err := strconv.ParseUint(tidbID, 10, 64)
	if err != nil {
		return nil, err
	}
	sourceURI, err := s.ToSourceURI(ctx, clusterID, source)
	if err != nil {
		return nil, err
	}
	parsedSourceURL, err := url.Parse(sourceURI)
	if err != nil {
		return nil, err
	}
	if !strings.HasSuffix(parsedSourceURL.Path, "/") {
		parsedSourceURL.Path += "/"
	}
	query := parsedSourceURL.Query()
	if parsedSourceURL.Scheme == "s3" || parsedSourceURL.Scheme == "oss" {
		query.Set("external-id", tidbID) // todo only for next gen
	}
	oss := parsedSourceURL.Scheme == "oss"
	if oss {
		parsedSourceURL.Scheme = "s3"
	}
	parsedSourceURL.RawQuery = query.Encode()

	inMetas := make([]*agentpb.GetImportFilesReq_TableMeta, 0, len(tableMetas))
	persistedTableMetas := make(entity.ImportIntoPrecheckTableMetas, 0, len(tableMetas))
	for _, inMeta := range tableMetas {
		resolvedWildcard, _, err := resolveImportIntoFilesWildcard(
			parsedSourceURL.String(),
			inMeta.GetDatabase(),
			inMeta.GetTable(),
			strings.TrimPrefix(inMeta.GetFilesWildcard(), "/"),
		)
		if err != nil {
			return nil, err
		}
		parsedWildcard, err := url.Parse(resolvedWildcard)
		if err != nil {
			return nil, err
		}
		parsedWildcard.RawQuery = ""
		inMetas = append(inMetas, &agentpb.GetImportFilesReq_TableMeta{
			Database:     inMeta.GetDatabase(),
			Table:        inMeta.GetTable(),
			WildcardPath: parsedWildcard.String(),
		})
		persistedTableMetas = append(persistedTableMetas, entity.ImportIntoPrecheckTableMeta{
			Database:      inMeta.GetDatabase(),
			Table:         inMeta.GetTable(),
			FilesWildcard: inMeta.GetFilesWildcard(),
		})
	}

	return &checkImportSourceTaskRequest{
		clusterID:  clusterID,
		sourceURL:  parsedSourceURL,
		sourceType: source.GetType(),
		inMetas:    inMetas,
		tableMetas: persistedTableMetas,
	}, nil
}

func (s *ImportIntoService) createCheckImportSourceTask(
	ctx context.Context,
	agentTaskID string,
	taskReq *checkImportSourceTaskRequest,
) error {
	agentCli, err := s.ngRegionalClientProvider.NewByClusterID(ctx, taskReq.clusterID)
	if err != nil {
		return pinerrors.Annotate(err, "failed to create agent client")
	}
	_, err = agentCli.CreateImportFilesTask(
		ctx,
		agentTaskID,
		taskReq.clusterID,
		taskReq.sourceURL.String(),
		-1,
		taskReq.inMetas...,
	)
	if err != nil {
		sanitized := *taskReq.sourceURL
		sanitized.RawQuery = ""
		sanitized.Fragment = ""
		log.Error("failed to create import files task", zap.Error(err), zap.String("url", sanitized.String()), zap.String("taskID", agentTaskID))
	}
	return err
}

func (s *ImportIntoService) getCheckImportSourceTask(
	ctx context.Context,
	clusterID uint64,
	taskID string,
) (*agentpb.GetImportFilesTaskResp, error) {
	agentCli, err := s.ngRegionalClientProvider.NewByClusterID(ctx, clusterID)
	if err != nil {
		return nil, pinerrors.Annotate(err, "failed to create agent client")
	}
	return agentCli.GetImportFilesTask(ctx, taskID)
}

func createImportIntoPrecheckTask(ctx context.Context, taskReq *checkImportSourceTaskRequest) (*entity.ImportIntoPrecheckTask, error) {
	precheckTask := &entity.ImportIntoPrecheckTask{
		ClusterID:  entity.ID(taskReq.clusterID),
		SourceType: taskReq.sourceType.String(),
		TableMetas: taskReq.tableMetas,
	}
	err := db.TransactionWithContext(ctx, func(tx *gorm.DB) error {
		return tx.Create(precheckTask).Error
	})
	if err != nil {
		return nil, err
	}
	return precheckTask, nil
}

func getImportIntoPrecheckTask(ctx context.Context, clusterID uint64, taskID string) (*entity.ImportIntoPrecheckTask, error) {
	taskIDUint64, err := strconv.ParseUint(taskID, 10, 64)
	if err != nil {
		return nil, err
	}
	var precheckTask entity.ImportIntoPrecheckTask
	err = db.TransactionWithContext(ctx, func(tx *gorm.DB) error {
		return tx.Model(&entity.ImportIntoPrecheckTask{}).
			Where("id = ? and cluster_id = ?", taskIDUint64, clusterID).
			First(&precheckTask).Error
	})
	if err != nil {
		return nil, err
	}
	return &precheckTask, nil
}

func deleteImportIntoPrecheckTask(ctx context.Context, clusterID uint64, taskID entity.ID) error {
	return db.TransactionWithContext(ctx, func(tx *gorm.DB) error {
		return tx.Where("id = ? and cluster_id = ?", taskID, clusterID).Delete(&entity.ImportIntoPrecheckTask{}).Error
	})
}

func wrapGetCheckImportSourceTaskErr(err error) error {
	if err == nil {
		return nil
	}
	if errors.Is(err, gorm.ErrRecordNotFound) {
		return bizerr.New(bizerr.ErrCodeRecordNotFound, "task not found")
	}
	return err
}

func wrapCreateCheckImportSourceTaskErr(err error, msg string) error {
	if err == nil {
		return nil
	}
	return pinerrors.Annotate(err, msg)
}

func normalizeImportIntoURIForComparison(rawURI string) string {
	parsedURI, err := url.Parse(rawURI)
	if err != nil {
		return rawURI
	}

	parsedURI.Path = strings.TrimRight(parsedURI.Path, "/")
	return parsedURI.String()
}

func isRootFilesWildcard(sourceURI, resolvedFilesWildcard string) bool {
	return normalizeImportIntoURIForComparison(sourceURI) == normalizeImportIntoURIForComparison(resolvedFilesWildcard)
}

func rewriteRootFilesWildcard(database, table string) string {
	return fmt.Sprintf("%s.%s.*.sql", database, table)
}

func resolveImportIntoFilesWildcard(sourceURI, database, table, filesWildcard string) (resolvedFilesWildcard string, rewrittenFilesWildcard string, err error) {
	resolvedFilesWildcard, err = JoinSourceURI(sourceURI, filesWildcard)
	if err != nil {
		return "", "", err
	}
	if !isRootFilesWildcard(sourceURI, resolvedFilesWildcard) {
		return resolvedFilesWildcard, "", nil
	}

	rewrittenFilesWildcard = rewriteRootFilesWildcard(database, table)
	resolvedFilesWildcard, err = JoinSourceURI(sourceURI, rewrittenFilesWildcard)
	if err != nil {
		return "", "", err
	}
	return resolvedFilesWildcard, rewrittenFilesWildcard, nil
}

func (s *ImportIntoService) ToSourceURI(ctx context.Context, clusterID uint64, source *importadminpb.ImportSource) (string, error) {
	var err error
	var uri string
	query := make(url.Values)
	switch source.GetType() {
	case importadminpb.ImportSourceType_ENUM_S3:
		uri = source.GetS3().GetUri()
		switch source.GetS3().GetAuthType() {
		case importadminpb.ImportS3AuthType_ENUM_ROLE_ARN:
			query.Set("role-arn", source.GetS3().GetRoleArn())
		case importadminpb.ImportS3AuthType_ENUM_ACCESS_KEY:
			query.Set("access-key", source.GetS3().GetAccessKey().GetId())
			query.Set("secret-access-key", source.GetS3().GetAccessKey().GetSecret())
		}
	case importadminpb.ImportSourceType_ENUM_OSS:
		uri = source.GetOss().GetUri()
		if uri == "" {
			uri = "oss://" + source.GetOss().GetBucket() + "/" + source.GetOss().GetPrefix()
		}
		endpoint := source.GetOss().GetEndpoint()
		ackId := source.GetOss().GetAccessKey().GetId()
		secret := source.GetOss().GetAccessKey().GetSecret()
		if endpoint == "" {
			_, endpoint, err = s.getAliOssRegionAndEndpoint(ctx, clusterID, uri, ackId, secret)
			if err != nil {
				return "", err
			}
		}
		switch source.GetOss().GetAuthType() {
		case importadminpb.ImportOSSAuthType_ENUM_ACCESS_KEY:
			query.Set("access-key", source.GetOss().GetAccessKey().GetId())
			query.Set("secret-access-key", source.GetOss().GetAccessKey().GetSecret())
			query.Set("endpoint", endpoint)
			query.Set("provider", "alibaba")
		}
	case importadminpb.ImportSourceType_ENUM_AZURE_BLOB:
		uri = source.GetAzureBlob().GetUri()
		if uri == "" {
			return "", pinerrors.New("azure blob uri is required")
		}
		parsedURI, err := url.Parse(uri)
		if err != nil {
			return "", pinerrors.Annotate(err, "failed to parse azure blob uri")
		}
		if parsedURI.Scheme != "https" {
			return "", pinerrors.New("azure blob uri must use https scheme")
		}
		sasToken := strings.TrimPrefix(source.GetAzureBlob().GetSasToken(), "?")
		if sasToken == "" {
			return "", pinerrors.New("azure blob sas token is required")
		}
		parsedURI.RawQuery = ""
		q := parsedURI.Query()
		q.Set("sas-token", sasToken)
		parsedURI.RawQuery = q.Encode()
		return azblob.ConvertAzblobURL(parsedURI.String(), true)
	default:
		return "", pinerrors.New("unsupported source type")
	}
	parsedUri, err := url.Parse(uri)
	if err != nil {
		return "", pinerrors.Annotate(err, "failed to parse source uri")
	}
	parsedUri.RawQuery = query.Encode()
	return parsedUri.String(), nil
}

func (s *ImportIntoService) convertImportToImportAdminImportIntoTask(impdb *entity.ImportIntoTask) (*importadminpb.ImportIntoTask, error) {
	var state importadminpb.ImportState_Enum
	switch impdb.Status {
	case entity.ImportIntoStatusPreparing:
		state = importadminpb.ImportState_ENUM_PREPARING
	case entity.ImportIntoStatusImporting:
		state = importadminpb.ImportState_ENUM_IMPORTING
	case entity.ImportIntoStatusCanceling:
		state = importadminpb.ImportState_ENUM_CANCELING
	case entity.ImportIntoStatusCanceled:
		state = importadminpb.ImportState_ENUM_CANCELED
	case entity.ImportIntoStatusCompleted:
		state = importadminpb.ImportState_ENUM_COMPLETED
	case entity.ImportIntoStatusFailed:
		state = importadminpb.ImportState_ENUM_FAILED
	case entity.ImportIntoStatusScalingUp:
		state = importadminpb.ImportState_ENUM_SCALING_UP
	default:
		state = importadminpb.ImportState_ENUM_UNSPECIFIED
	}
	importOptions := &importadminpb.ImportOptions{}
	if impdb.CustomCSVFormat != "" {
		var customCSVFormat importadminpb.CSVFormat
		err := protojson.Unmarshal([]byte(impdb.CustomCSVFormat), &customCSVFormat)
		if err != nil {
			return nil, err
		}
		importOptions.CsvFormat = &customCSVFormat
	}

	importSource, err := FromSourceURI(impdb.SourceURI)
	if err != nil {
		return nil, err
	}

	subTasks := make([]*importadminpb.ImportIntoSubTask, len(impdb.ImportResults.SubTasks))
	targetTables := make([]*importadminpb.ImportTargetTableInfo, len(impdb.ImportResults.SubTasks))
	for i, result := range impdb.ImportResults.SubTasks {
		path, err := url.Parse(result.ImportIntoFilesWildcard)
		if err != nil {
			return nil, err
		}
		path.RawQuery = ""
		targetTables[i] = &importadminpb.ImportTargetTableInfo{
			Database:      result.DataBase,
			Table:         result.TableName,
			FilesWildcard: path.String(),
		}
		curStep := result.CurrentStep
		if curStep == "" && !slices.Contains([]string{entity.ImportIntoJobStatusFinished, entity.ImportIntoJobStatusCanceled, entity.ImportIntoJobStatusFailed}, result.ImportIntoJobStatus) {
			curStep = "prevalidate"
		}
		subTasks[i] = &importadminpb.ImportIntoSubTask{
			Database:                     result.DataBase,
			Table:                        result.TableName,
			FilesWildcard:                path.String(),
			ImportedRows:                 uint64(result.ImportedRows),
			TotalSize:                    result.ImportIntoTotalSize,
			JobId:                        result.ImportIntoJobID,
			Columns:                      result.Columns,
			CurrentStep:                  curStep,
			CurrentStepTotalSize:         result.CurrentStepTotalSize,
			CurrentStepProgressPercent:   result.CurrentStepProgressPercent,
			CurrentStepEta:               result.CurrentStepETA,
			CurrentStepProcessedDataSize: result.CurrentStepProcessedDataSize,
			CurrentStepSpeed:             result.CurrentStepSpeed,
		}
		if result.ImportStartedAt != nil {
			subTasks[i].StartedTime = timestamppb.New(*result.ImportStartedAt)
			subTasks[i].Phase = result.ImportIntoPhase
		}
		subTasks[i].ResultMessage = result.Message
		if result.ImportFinishedAt != nil {
			subTasks[i].FinishedTime = timestamppb.New(*result.ImportFinishedAt)
		}
		subTasks[i].Status = importIntoSubTaskStatusToPB(result)
	}

	creationDetails := &importadminpb.CreationDetails{
		ImportOptions:    importOptions,
		Source:           importSource,
		TargetTableInfos: targetTables,
	}

	var completeTime *timestamppb.Timestamp
	if impdb.FinishedAt != nil {
		completeTime = timestamppb.New(*impdb.FinishedAt)
	}

	task := &importadminpb.ImportIntoTask{
		TidbId:          strconv.FormatUint(impdb.ClusterID.Uint64(), 10),
		CreateTime:      timestamppb.New(impdb.CreatedAt),
		CompleteTime:    completeTime,
		State:           state,
		Message:         impdb.Message,
		CreationDetails: creationDetails,
		ImportId:        strconv.FormatUint(impdb.ID.Uint64(), 10),
		Subtasks:        subTasks,
		Creator:         impdb.Creator,
	}
	if impdb.GroupKey != "" {
		groupKey := impdb.GroupKey
		task.GroupKey = &groupKey
	}
	task.TotalSize = int64(impdb.TotalSize)
	return task, nil
}

func importIntoSubTaskStatusToPB(result entity.ImportIntoSubTask) importadminpb.ImportSubTaskStatus_Enum {
	switch result.ImportIntoJobStatus {
	case entity.ImportIntoJobStatusPending:
		return importadminpb.ImportSubTaskStatus_ENUM_PENDING
	case entity.ImportIntoJobStatusRunning:
		return importadminpb.ImportSubTaskStatus_ENUM_RUNNING
	case entity.ImportIntoJobStatusCanceled:
		return importadminpb.ImportSubTaskStatus_ENUM_CANCELED
	case entity.ImportIntoJobStatusFailed:
		return importadminpb.ImportSubTaskStatus_ENUM_FAILED
	case entity.ImportIntoJobStatusFinished:
		switch result.Result {
		case entity.ImportTableResultSuccess:
			return importadminpb.ImportSubTaskStatus_ENUM_FINISHED
		case entity.ImportTableResultWarning:
			return importadminpb.ImportSubTaskStatus_ENUM_WARNING
		default:
			return importadminpb.ImportSubTaskStatus_ENUM_FINISHED
		}
	default:
		return importadminpb.ImportSubTaskStatus_ENUM_UNSPECIFIED
	}
}

func importStateToImportIntoStatus(state importadminpb.ImportState_Enum) (entity.ImportIntoStatus, error) {
	switch state {
	case importadminpb.ImportState_ENUM_UNSPECIFIED:
		return entity.ImportIntoStatusUnknown, nil
	case importadminpb.ImportState_ENUM_PREPARING:
		return entity.ImportIntoStatusPreparing, nil
	case importadminpb.ImportState_ENUM_IMPORTING:
		return entity.ImportIntoStatusImporting, nil
	case importadminpb.ImportState_ENUM_COMPLETED:
		return entity.ImportIntoStatusCompleted, nil
	case importadminpb.ImportState_ENUM_FAILED:
		return entity.ImportIntoStatusFailed, nil
	case importadminpb.ImportState_ENUM_CANCELING:
		return entity.ImportIntoStatusCanceling, nil
	case importadminpb.ImportState_ENUM_CANCELED:
		return entity.ImportIntoStatusCanceled, nil
	case importadminpb.ImportState_ENUM_SCALING_UP:
		return entity.ImportIntoStatusScalingUp, nil
	default:
		return "", bizerr.New(bizerr.ErrCodeParams, "invalid import state: %s", state.String())
	}
}

func (s *ImportIntoService) recordAuditLogForImport(ctx context.Context, clusterID uint64, req *importadminpb.CreateImportIntoTaskRequest, retErr error) {
	var err error
	listResp, err := s.mgmtSvcGrpcClient.ListClusters(ctx, &mgmt.ListClustersRequest{Filter: &mgmt.ClusterFilter{ClusterIds: []uint64{clusterID}}})
	if err != nil || len(listResp.Clusters) == 0 {
		log.Error("list mgmt clusters error", zap.Error(err), zap.Uint64("clusterID", clusterID))
		return
	}
	cluster := listResp.Clusters[0]
	source := req.GetImport().GetCreationDetails().GetSource()
	var eventType auditlog.EventType

	details := auditlog.Details{
		auditlog.KeyTargetClusterName: cluster.GetClusterName(),
		auditlog.KeyTargetRegion:      cluster.GetRegion(),
		auditlog.KeyError:             err,
	}
	switch source.GetType() {
	case importadminpb.ImportSourceType_ENUM_S3:
		details[auditlog.KeySourceDataSource] = "AWS"
		details[auditlog.KeyBucketUri] = utils.RedactSourceURI(source.GetS3().GetUri())
		eventType = auditlog.EventImportDataFromAWS
	case importadminpb.ImportSourceType_ENUM_OSS:
		details[auditlog.KeySourceDataSource] = "ALIBABA"
		details[auditlog.KeyBucketUri] = utils.RedactSourceURI(source.GetOss().GetUri())
		eventType = auditlog.EventImportDataFromAlibabaCloud
	case importadminpb.ImportSourceType_ENUM_AZURE_BLOB:
		details[auditlog.KeySourceDataSource] = "AZURE"
		bucketURI := source.GetAzureBlob().GetUri()
		if httpsURI, _, nerr := normalizeAzureBlobInput(bucketURI, source.GetAzureBlob().GetSasToken()); nerr == nil {
			bucketURI = httpsURI
		} else {
			if u, perr := url.Parse(bucketURI); perr == nil {
				u.RawQuery = ""
				u.Fragment = ""
				bucketURI = u.String()
			} else if idx := strings.Index(bucketURI, "?"); idx >= 0 {
				bucketURI = bucketURI[:idx]
			}
		}
		details[auditlog.KeyBucketUri] = utils.RedactSourceURI(bucketURI)
		eventType = auditlog.EventImportDataFromAzure
	default:
		log.Error("unsupport type")
		return
	}
	err = auditlog.NewEvent(ctx, eventType).
		SetResult(retErr == nil).
		SetOrgID(cluster.GetOrganizationId()).
		SetProjectID(cluster.GetProjectId()).
		SetDetails(details.Build()).
		Submit()
	if err != nil {
		log.Warn("failed to submit audit log", zap.Error(err))
		return
	}
}

func buildImportIntoGroupKey(importID entity.ID) string {
	return fmt.Sprintf("import-into-%d", importID.Uint64())
}

// CreateImportIntoTask Create an import job
func (s *ImportIntoService) CreateImportIntoTask(ctx context.Context, req *importadminpb.CreateImportIntoTaskRequest) (res *importadminpb.ImportIntoTask, retErr error) {
	sourceType := req.GetImport().GetCreationDetails().GetSource().GetType().String()
	log.L().Info("CreateImportIntoTask called", zap.String("tidbID", req.GetTidbId()), zap.String("name", req.GetImport().GetName()), zap.String("sourceType", sourceType))
	clusterID, err := strconv.ParseUint(req.GetTidbId(), 10, 64)
	if err != nil {
		return nil, err
	}
	scoped, err := s.ensureImportIntoWriteEnabled(ctx, clusterID)
	if err != nil {
		return nil, err
	}
	defer func() {
		s.recordAuditLogForImport(ctx, clusterID, req, retErr)
	}()
	ent := entity.ImportIntoTask{
		ClusterID:      entity.ID(clusterID),
		ProjectID:      entity.ID(scoped.ProjectID),
		TenantID:       entity.ID(scoped.TenantID),
		ServerTier:     entity.ServerTierNextGeneration,
		Status:         entity.ImportIntoStatusPreparing,
		ArchType:       entity.ThreeLayerArch,
		TargetUsername: req.GetImport().GetUsername(),
		CreatorID:      scoped.User.ID,
		Creator:        scoped.User.Email,
		Name:           req.GetImport().GetName(),
	}
	if req.GetImport().GetPassword() != "" {
		passwd, err := s.encryptor.EncryptBase64(req.GetImport().GetPassword())
		if err != nil {
			log.Error("failed to decrypt cluster password", zap.Error(err))
			return nil, pinerrors.Trace(err)
		}
		ent.TargetPassword = passwd
	}
	ent.SourceURI, err = s.ToSourceURI(ctx, clusterID, req.GetImport().GetCreationDetails().GetSource())
	if err != nil {
		return nil, err
	}
	csvFormat := req.GetImport().GetCreationDetails().GetImportOptions().GetCsvFormat()
	if csvFormat != nil {
		encodeCSVFormat, err := protojson.Marshal(csvFormat)
		if err != nil {
			return nil, err
		}
		ent.CustomCSVFormat = string(encodeCSVFormat)
		log.L().Info("encoded custom csv format", zap.String("csv_format", string(encodeCSVFormat)))
	}
	subTasks := entity.ImportIntoSubTasks{SubTasks: make([]entity.ImportIntoSubTask, 0, len(req.GetImport().GetCreationDetails().GetTargetTableInfos()))}
	for _, t := range req.GetImport().GetCreationDetails().GetTargetTableInfos() {
		resolvedFilesWildcard, rewrittenFilesWildcard, err := resolveImportIntoFilesWildcard(ent.SourceURI, t.GetDatabase(), t.GetTable(), t.GetFilesWildcard())
		if err != nil {
			return nil, bizerr.New(bizerr.ErrCodeParams, "invalid wildcard format: %s", t.GetFilesWildcard())
		}
		if rewrittenFilesWildcard != "" {
			log.L().Warn("rewrite target table info whose files wildcard equals import source uri",
				zap.String("tidbID", req.GetTidbId()),
				zap.String("database", t.GetDatabase()),
				zap.String("table", t.GetTable()),
				zap.String("filesWildcard", utils.RedactSourceURI(t.GetFilesWildcard())),
				zap.String("rewrittenFilesWildcard", utils.RedactSourceURI(rewrittenFilesWildcard)),
				zap.String("resolvedFilesWildcard", utils.RedactSourceURI(resolvedFilesWildcard)),
				zap.String("sourceURI", utils.RedactSourceURI(ent.SourceURI)),
			)
		}
		storedWildcard, stripErr := utils.StripURIQuery(resolvedFilesWildcard)
		if stripErr != nil {
			return nil, bizerr.New(bizerr.ErrCodeParams, "invalid files wildcard uri: %s", utils.RedactSourceURI(resolvedFilesWildcard))
		}
		subTasks.SubTasks = append(subTasks.SubTasks, entity.ImportIntoSubTask{
			TableName:               t.GetTable(),
			DataBase:                t.GetDatabase(),
			ImportIntoFilesWildcard: storedWildcard,
		})
	}
	ent.ImportResults = subTasks
	clusterJobRegistered := false
	if err = db.TransactionWithContext(ctx, func(tx *gorm.DB) error {
		if createErr := tx.Create(&ent).Error; createErr != nil {
			return createErr
		}
		ent.GroupKey = buildImportIntoGroupKey(ent.ID)
		if updateErr := tx.Model(&ent).Update("group_key", ent.GroupKey).Error; updateErr != nil {
			return updateErr
		}
		if registerErr := s.clusterJob.EnsureImportClusterJob(ctx, clusterID, ent.ID.Uint64()); registerErr != nil {
			return registerErr.Chain("register cluster job")
		}
		clusterJobRegistered = true
		ent.CustomStatusRegistered = true
		return tx.Model(&ent).Update("custom_status_registered", true).Error
	}); err != nil {
		if clusterJobRegistered {
			cleanupCtx, cleanupCancel := context.WithTimeout(context.WithoutCancel(ctx), importClusterJobCleanupTimeout)
			defer cleanupCancel()
			if deregisterErr := s.clusterJob.DeregisterImportClusterJob(
				cleanupCtx, clusterID, ent.ID.Uint64(),
			); deregisterErr != nil {
				pkglog.GetLogger(ctx).Error("failed to compensate Import Into custom status registration",
					zap.Uint64("clusterID", clusterID),
					zap.Uint64("importID", ent.ID.Uint64()),
					zap.Error(deregisterErr),
				)
			}
		}
		return nil, err
	}
	pimport.PublishImportIntoEvent(ctx, s.eventPublisher, &ent)
	return s.convertImportToImportAdminImportIntoTask(&ent)
}

// ListImportIntoTasks Provides a list of import jobs
func (s *ImportIntoService) ListImportIntoTasks(ctx context.Context, req *importadminpb.ListImportIntoTasksRequest) (*importadminpb.ListImportIntoTasksResponse, error) {
	var entities []*entity.ImportIntoTask
	var totalSize int64
	pageSize := int(req.GetPageSize())
	pageToken := req.GetPageToken()

	if pageSize == 0 {
		pageSize = 10
	} else if pageSize > 100 {
		pageSize = 100
	}
	err := db.TransactionWithContext(ctx, func(tx *gorm.DB) error {
		q := tx.Model(&entity.ImportIntoTask{}).
			Where("cluster_id = ?", req.GetTidbId())
		if req.GetQ() != "" {
			q = q.Where(tx.Where("name LIKE ?", "%"+req.GetQ()+"%").Or("group_key LIKE ?", "%"+req.GetQ()+"%"))
		}
		if req.GetStatus() != importadminpb.ImportState_ENUM_UNSPECIFIED {
			status, err := importStateToImportIntoStatus(req.GetStatus())
			if err != nil {
				return err
			}
			q = q.Where("status = ?", status)
		}
		if err := q.Count(&totalSize).Error; err != nil {
			return err
		}
		query := q.Limit(pageSize).Order("created_at desc")
		_, err := base64.StdEncoding.DecodeString(pageToken) // todo remove if FE adopted
		if pageToken != "" && err == nil {                   // for compatible for FE
			offset, err := base64.StdEncoding.DecodeString(pageToken)
			if err != nil {
				return fmt.Errorf("invalid page token: %w", err)
			}
			offsetInt, err := strconv.Atoi(string(offset))
			if err != nil {
				return fmt.Errorf("invalid page token format: %w", err)
			}
			query = query.Offset(offsetInt)
		}
		if err := query.Scan(&entities).Error; err != nil {
			return err
		}
		return nil
	})
	if err != nil {
		return nil, err
	}
	var res importadminpb.ListImportIntoTasksResponse
	res.Imports = make([]*importadminpb.ImportIntoTask, len(entities))
	for i, ent := range entities {
		imp, err := s.convertImportToImportAdminImportIntoTask(ent)
		if err != nil {
			return nil, err
		}
		res.Imports[i] = imp
	}
	res.TotalSize = int32(totalSize)

	if int(totalSize) > len(entities) {
		nextOffset := 0
		if pageToken != "" {
			offset, _ := base64.StdEncoding.DecodeString(pageToken)
			offsetInt, _ := strconv.Atoi(string(offset))
			nextOffset = offsetInt + len(entities)
		} else {
			nextOffset = len(entities)
		}
		res.NextPageToken = base64.StdEncoding.EncodeToString([]byte(strconv.Itoa(nextOffset)))
	}
	return &res, nil
}

// GetImportIntoTask Get an import detailed information
func (s *ImportIntoService) GetImportIntoTask(ctx context.Context, req *importadminpb.GetImportIntoTaskRequest) (*importadminpb.ImportIntoTask, error) {
	var ent entity.ImportIntoTask
	err := db.TransactionWithContext(ctx, func(tx *gorm.DB) error {
		err := tx.Model(&entity.ImportIntoTask{}).Where("id = ? and cluster_id = ?", req.GetImportId(), req.GetTidbId()).First(&ent).Error
		if err != nil {
			return err
		}
		return nil
	})
	if err != nil {
		return nil, err
	}
	return s.convertImportToImportAdminImportIntoTask(&ent)
}

func (s *ImportIntoService) recordAuditLogForCancelImport(ctx context.Context, req *importadminpb.CancelImportIntoTaskRequest, retErr error) {
	var err error
	var ent entity.ImportIntoTask
	err = db.TransactionWithContext(ctx, func(tx *gorm.DB) error {
		err := tx.Model(&entity.ImportIntoTask{}).Where("id = ? and cluster_id = ?", req.GetImportId(), req.GetTidbId()).First(&ent).Error
		if err != nil {
			return err
		}
		return nil
	})
	if err != nil {
		return
	}
	imp, err := s.convertImportToImportAdminImportIntoTask(&ent)
	if err != nil {
		return
	}
	clusterID, err := strconv.ParseUint(req.GetTidbId(), 10, 64)
	if err != nil {
		return
	}
	listResp, err := s.mgmtSvcGrpcClient.ListClusters(ctx, &mgmt.ListClustersRequest{Filter: &mgmt.ClusterFilter{ClusterIds: []uint64{clusterID}}})
	if err != nil || len(listResp.Clusters) == 0 {
		log.Error("list mgmt clusters error", zap.Error(err), zap.Uint64("clusterID", clusterID))
		return
	}
	cluster := listResp.Clusters[0]
	eventType := auditlog.EventImportDataCanceled

	details := auditlog.Details{
		auditlog.KeyTargetClusterName: cluster.GetClusterName(),
		auditlog.KeyTargetRegion:      cluster.GetRegion(),
		auditlog.KeyError:             err,
	}
	details[auditlog.KeyBucketUri] = utils.RedactSourceURI(imp.CreationDetails.Source.GetUri())
	switch imp.CreationDetails.Source.Type {
	case importadminpb.ImportSourceType_ENUM_S3:
		details[auditlog.KeySourceDataSource] = "AWS"
	case importadminpb.ImportSourceType_ENUM_OSS:
		details[auditlog.KeySourceDataSource] = "ALIBABA"
	case importadminpb.ImportSourceType_ENUM_AZURE_BLOB:
		details[auditlog.KeySourceDataSource] = "AZURE"
	default:
		log.Error("unsupport type")
		return
	}
	err = auditlog.NewEvent(ctx, eventType).
		SetResult(retErr == nil).
		SetOrgID(cluster.GetOrganizationId()).
		SetProjectID(cluster.GetProjectId()).
		SetDetails(details.Build()).
		Submit()
	if err != nil {
		log.Warn("failed to submit audit log", zap.Error(err))
		return
	}
}

// CancelImportIntoTask Cancel an import job
func (s *ImportIntoService) CancelImportIntoTask(ctx context.Context, req *importadminpb.CancelImportIntoTaskRequest) (res *emptypb.Empty, retErr error) {
	if _, err := s.ensureImportIntoWriteEnabled(ctx, req.GetTidbId()); err != nil {
		return nil, err
	}
	defer func() {
		s.recordAuditLogForCancelImport(ctx, req, retErr)
	}()
	err := db.TransactionWithContext(ctx, func(tx *gorm.DB) error {
		err := tx.Model(&entity.ImportIntoTask{}).Where("id = ? and cluster_id = ? and status in ?", req.GetImportId(), req.GetTidbId(), []entity.ImportIntoStatus{entity.ImportIntoStatusImporting, entity.ImportIntoStatusPreparing, entity.ImportIntoStatusScalingUp}).Update("status", entity.ImportIntoStatusCanceling).Error
		if err != nil {
			return err
		}
		return nil
	})
	return nil, err
}

func (s *ImportIntoService) CheckImportSource(ctx context.Context, req *importadminpb.CheckImportSourceReq) (*importadminpb.CheckImportSourceResp, error) {
	clusterID, err := strconv.ParseUint(req.GetTidbId(), 10, 64)
	if err != nil {
		return nil, err
	}
	agentCli, err := s.ngRegionalClientProvider.NewByClusterID(ctx, clusterID)
	if err != nil {
		return nil, pinerrors.Annotate(err, "failed to create agent client")
	}
	sourceURI, err := s.ToSourceURI(ctx, clusterID, req.GetSource())
	if err != nil {
		return nil, err
	}
	parsedSourceUrl, err := url.Parse(sourceURI)
	if err != nil {
		return nil, err
	}
	if !strings.HasSuffix(parsedSourceUrl.Path, "/") {
		parsedSourceUrl.Path = parsedSourceUrl.Path + "/"
	}
	query := parsedSourceUrl.Query()
	if parsedSourceUrl.Scheme == "s3" || parsedSourceUrl.Scheme == "oss" {
		query.Set("external-id", fmt.Sprint(req.GetTidbId())) // todo only for next gen
	}
	oss := parsedSourceUrl.Scheme == "oss"
	if oss {
		parsedSourceUrl.Scheme = "s3"
	}
	parsedSourceUrl.RawQuery = query.Encode()
	var inMetas []*agentpb.GetImportFilesReq_TableMeta
	for _, inMeta := range req.GetTableMetas() {
		wildcard := inMeta.GetFilesWildcard()
		wildcard = strings.TrimPrefix(wildcard, "/")
		parsedWildcard := parsedSourceUrl.JoinPath(wildcard)
		parsedWildcard.RawQuery = ""
		inMetas = append(inMetas, &agentpb.GetImportFilesReq_TableMeta{
			Database:     inMeta.GetDatabase(),
			Table:        inMeta.GetTable(),
			WildcardPath: parsedWildcard.String(),
		})
	}
	resp, err := agentCli.GetImportFiles(ctx, clusterID, parsedSourceUrl.String(), 2000, inMetas...)
	if err != nil {
		sanitized := *parsedSourceUrl
		sanitized.RawQuery = ""
		sanitized.Fragment = ""
		log.Error("failed to call get import files", zap.Error(err), zap.String("url", sanitized.String()))
		errMsg := "Please ensure import files is validated."
		return &importadminpb.CheckImportSourceResp{Success: false, ErrorMsg: &errMsg}, nil
	}
	return buildCheckImportSourceSuccessResp(req, resp, oss)
}

func (s *ImportIntoService) CreateCheckImportSourceTask(ctx context.Context, req *importadminpb.CreateCheckImportSourceTaskReq) (*importadminpb.CreateCheckImportSourceTaskResp, error) {
	if _, err := s.ensureImportIntoWriteEnabled(ctx, req.GetTidbId()); err != nil {
		return nil, err
	}
	taskReq, err := s.buildCheckImportSourceTaskRequest(ctx, req.GetTidbId(), req.GetSource(), req.GetTableMetas())
	if err != nil {
		return nil, err
	}
	precheckTask, err := createImportIntoPrecheckTask(ctx, taskReq)
	if err != nil {
		return nil, wrapCreateCheckImportSourceTaskErr(err, "failed to persist check import source task")
	}
	agentTaskID := strconv.FormatUint(uint64(precheckTask.ID), 10)
	if err := s.createCheckImportSourceTask(ctx, agentTaskID, taskReq); err != nil {
		if deleteErr := deleteImportIntoPrecheckTask(ctx, taskReq.clusterID, precheckTask.ID); deleteErr != nil {
			log.Error("failed to clean up precheck task after agent task creation failure", zap.Error(deleteErr), zap.Uint64("clusterID", taskReq.clusterID), zap.Uint64("taskID", uint64(precheckTask.ID)))
		}
		return nil, wrapCreateCheckImportSourceTaskErr(err, "failed to create check import source task on agent")
	}
	return &importadminpb.CreateCheckImportSourceTaskResp{TaskId: agentTaskID}, nil
}

func (s *ImportIntoService) GetCheckImportSourceTask(ctx context.Context, req *importadminpb.GetCheckImportSourceTaskReq) (*importadminpb.GetCheckImportSourceTaskResp, error) {
	clusterID, err := strconv.ParseUint(req.GetTidbId(), 10, 64)
	if err != nil {
		return nil, err
	}
	precheckTask, err := getImportIntoPrecheckTask(ctx, clusterID, req.GetTaskId())
	if err != nil {
		return nil, wrapGetCheckImportSourceTaskErr(err)
	}
	taskResp, err := s.getCheckImportSourceTask(ctx, clusterID, strconv.FormatUint(uint64(precheckTask.ID), 10))
	if err != nil {
		return nil, err
	}
	switch taskResp.GetStatus() {
	case agentpb.ImportFilesTaskStatus_IMPORT_FILES_TASK_STATUS_RUNNING,
		agentpb.ImportFilesTaskStatus_IMPORT_FILES_TASK_STATUS_UNSPECIFIED:
		return &importadminpb.GetCheckImportSourceTaskResp{
			TaskId:  req.GetTaskId(),
			Success: true,
			Status:  importadminpb.CheckImportSourceStatus_ENUM_RUNNING,
		}, nil
	case agentpb.ImportFilesTaskStatus_IMPORT_FILES_TASK_STATUS_FAILED:
		errMsg := payloadErrorOrFallback(taskResp.GetErrorMsg(), "Please ensure import files is validated.")
		return &importadminpb.GetCheckImportSourceTaskResp{
			TaskId:   req.GetTaskId(),
			Success:  false,
			ErrorMsg: &errMsg,
		}, nil
	case agentpb.ImportFilesTaskStatus_IMPORT_FILES_TASK_STATUS_SUCCEEDED:
	default:
		errMsg := "Please ensure import files is validated."
		return &importadminpb.GetCheckImportSourceTaskResp{
			TaskId:   req.GetTaskId(),
			Success:  false,
			ErrorMsg: &errMsg,
		}, nil
	}
	if taskResp.GetResult() == nil {
		errMsg := "Please ensure import files is validated."
		return &importadminpb.GetCheckImportSourceTaskResp{
			TaskId:   req.GetTaskId(),
			Success:  false,
			ErrorMsg: &errMsg,
		}, nil
	}
	return buildGetCheckImportSourceTaskSuccessResp(
		req.GetTaskId(),
		precheckTask,
		taskResp.GetResult(),
		precheckTask.SourceType == importadminpb.ImportSourceType_ENUM_OSS.String(),
	)
}

func (s *ImportIntoService) CheckDatabaseConnectivity(ctx context.Context, req *importadminpb.CheckDatabaseConnectivityReq) (*importadminpb.CheckDatabaseConnectivityResp, error) {
	clusterID, err := strconv.ParseUint(req.GetTidbId(), 10, 64)
	if err != nil {
		return nil, err
	}
	agentCli, err := s.ngRegionalClientProvider.NewByClusterID(ctx, clusterID)
	if err != nil {
		return nil, pinerrors.Annotate(err, "failed to create agent client")
	}
	var payloadErrMsg string
	err = retry.Do(func() error {
		pemResp, err := agentCli.GetPublicKey(ctx, &agentpb.GetPublicKeyRequest{})
		if err != nil {
			return err
		}
		pem := pemResp.GetPem()
		cipherPasswd, err := crypto.EncryptWithPEM(pem, []byte(req.GetPassword()))
		if err != nil {
			return retry.Unrecoverable(err)
		}
		execReq := &agentpb.ExecSQLReq{
			Sql: "SELECT 1",
			Cluster: &agentpb.Cluster{
				ClusterId:         clusterID,
				User:              req.Username,
				PasswordEncrypted: true,
				EncryptedPassword: &agentpb.EncryptedData{
					Pem:  pem,
					Data: cipherPasswd,
				},
			},
			TypedColumns: true,
			ColumnNum:    1,
		}
		resp, err := agentCli.ExecuteSQL(ctx, execReq)
		if err != nil {
			return retry.Unrecoverable(err)
		}
		if resp.GetPemChanged() {
			return agent.PEMChangedError
		}
		if resp.GetErrorMsg() != "" {
			payloadErrMsg = resp.GetErrorMsg()
			return retry.Unrecoverable(errors.New(payloadErrMsg))
		}
		return nil
	},
		retry.Attempts(config.Global().ImportConfig.MaxRetryTimesOnExecImportingQuerySQL),
		retry.RetryIf(func(err error) bool {
			return errors.Is(err, agent.PEMChangedError)
		}),
	)

	if err != nil {
		// If there's an error, the connection failed.
		log.Info("failed to call execute sql in check database connectivity", zap.Error(err), zap.String("url", req.GetTidbId()))
		return checkDatabaseConnectivityFailureResp(payloadErrMsg), nil
	}
	return &importadminpb.CheckDatabaseConnectivityResp{Success: true}, nil
}

func (s *ImportIntoService) getAliOssRegionAndEndpoint(ctx context.Context, clusterID uint64, uri, ackID, ackSec string) (string, string, error) {
	resp, err := s.mgmtSvcGrpcClient.ListClusters(ctx, &mgmt.ListClustersRequest{Filter: &mgmt.ClusterFilter{ClusterIds: []uint64{clusterID}}})
	if err != nil {
		return "", "", err
	}
	if len(resp.Clusters) == 0 {
		return "", "", fmt.Errorf("no cluster")
	}
	cluster := resp.Clusters[0]
	region, err := cloud_storage.GetOssBucketRegion(ctx, uri, cloud_storage.WithOSSAccessKey(ackID, ackSec))
	if err != nil {
		return "", "", err
	}
	if cluster.GetCloudProvider() == "alicloud" && strings.HasPrefix(cluster.GetRegion(), region) {
		return region, cloud_storage.GenerateOssInternalEndpoint(region), nil
	}
	log.Info("provider and region not match", zap.Uint64("clusterID", clusterID), zap.String("BucketRegion", region), zap.String("ClusterRegion", cluster.GetRegion()), zap.String("clusterProvider", cluster.GetCloudProvider()))
	return region, cloud_storage.GenerateOssEndpoint(region), nil
}

func normalizeAzureBlobInput(uri string, sasToken string) (httpsURI string, rawSas string, err error) {
	rawSas = strings.TrimPrefix(sasToken, "?")
	if uri == "" {
		return "", "", pinerrors.New("azure blob uri is required")
	}
	if rawSas == "" {
		return "", "", pinerrors.New("azure blob sas token is required")
	}

	u, err := url.Parse(uri)
	if err != nil {
		return "", "", pinerrors.Annotate(err, "failed to parse azure blob uri")
	}
	if u.Scheme != "https" {
		return "", "", pinerrors.New("azure blob uri must use https scheme")
	}

	sanitized := *u
	sanitized.RawQuery = ""
	sanitized.Fragment = ""
	httpsURI = sanitized.String()

	validateURL := *u
	validateURL.RawQuery = ""
	validateURL.Fragment = ""
	q := make(url.Values)
	q.Set("sas-token", "redacted")
	validateURL.RawQuery = q.Encode()
	if _, err := azblob.ParseAzblobURL(validateURL.String()); err != nil {
		return "", "", pinerrors.New("invalid azure blob uri")
	}

	return httpsURI, rawSas, nil
}

func (s *ImportIntoService) CheckImportSourceConnection(ctx context.Context, req *importadminpb.CheckImportSourceConnectionReq) (*importadminpb.CheckImportSourceConnectionResp, error) {
	clusterID, err := strconv.ParseUint(req.GetTidbId(), 10, 64)
	if err != nil {
		return nil, err
	}
	if err := validateImportSourceAuth(req.GetSource()); err != nil {
		return &importadminpb.CheckImportSourceConnectionResp{Success: false, ErrorMsg: err.Error()}, nil
	}
	agentCli, err := s.ngRegionalClientProvider.NewByClusterID(ctx, clusterID)
	if err != nil {
		return nil, pinerrors.Annotate(err, "failed to create agent client")
	}
	switch req.GetSource().GetType() {
	case importadminpb.ImportSourceType_ENUM_S3:
		log.Info("start to check s3 access", zap.Uint64("clusterID", clusterID), zap.String("uri", utils.RedactSourceURI(req.GetSource().GetS3().GetUri())))
		res, err := agentCli.CheckS3Access(ctx, &agentpb.CheckS3AccessReq{
			RoleArn:            req.GetSource().GetS3().GetRoleArn(),
			AwsAccessKeyId:     req.GetSource().GetS3().GetAccessKey().GetId(),
			AwsSecretAccessKey: req.GetSource().GetS3().GetAccessKey().GetSecret(),
			S3Url:              req.GetSource().GetS3().Uri,
			ExternalId:         req.GetTidbId(),
		})
		errMsg := "failed to check s3 access"
		if err != nil {
			log.Error(errMsg, zap.Error(err))
			return &importadminpb.CheckImportSourceConnectionResp{Success: false, ErrorMsg: errMsg}, nil
		}
		if res.GetErrorMsg() != "" {
			return &importadminpb.CheckImportSourceConnectionResp{Success: false, ErrorMsg: res.GetErrorMsg()}, nil
		}
		log.Info("success to check s3 access", zap.Uint64("clusterID", clusterID), zap.String("uri", utils.RedactSourceURI(req.GetSource().GetS3().GetUri())))
		return &importadminpb.CheckImportSourceConnectionResp{Success: true, Region: res.GetRegion()}, nil
	case importadminpb.ImportSourceType_ENUM_OSS:
		var bizErrorMsg string
		uri := req.GetSource().GetOss().GetUri()
		log.Info("start to check oss access", zap.Uint64("clusterID", clusterID), zap.String("uri", utils.RedactSourceURI(uri)))
		if uri == "" {
			uri = "oss://" + req.GetSource().GetOss().GetBucket() + "/" + req.GetSource().GetOss().GetPrefix()
		}
		ackId := req.GetSource().GetOss().GetAccessKey().GetId()
		secret := req.GetSource().GetOss().GetAccessKey().GetSecret()
		region, endpoint, err := s.getAliOssRegionAndEndpoint(ctx, clusterID, uri, ackId, secret)
		if err != nil {
			return nil, err
		}
		if req.GetSource().GetOss().GetEndpoint() != "" {
			endpoint = req.GetSource().GetOss().GetEndpoint()
		}
		// Retry the entire process of fetching PEM, encrypting, and checking access.
		err = retry.Do(
			func() error {
				// 1. Get the public key (PEM)
				pemResp, err := agentCli.GetPublicKey(ctx, &agentpb.GetPublicKeyRequest{})
				if err != nil {
					return retry.Unrecoverable(err) // Unrecoverable on RPC error
				}
				pem := pemResp.GetPem()
				// 2. Encrypt the secret using the new PEM
				cipherPasswd, err := crypto.EncryptWithPEM(pem, []byte(secret))
				if err != nil {
					return retry.Unrecoverable(err) // Unrecoverable on encryption failure
				}
				// 3. Call CheckAliOSSAccess
				res, err := agentCli.CheckAliOSSAccess(ctx, &agentpb.CheckAliOSSAccessReq{
					OssUri:         uri,
					OssAccessKeyId: ackId,
					OssAccessKeySecret: &agentpb.EncryptedData{
						Pem:  pem,
						Data: cipherPasswd,
					},
					Endpoint: endpoint,
				})
				if err != nil {
					return err // Default retryable for RPC errors
				}

				// Check for PemChanged flag
				if res.GetPemChanged() {
					// Return custom error to trigger a retry
					return &pemChangedError{}
				}

				// Check for business error message
				if res.GetErrorMsg() != "" {
					bizErrorMsg = res.GetErrorMsg()
					return retry.Unrecoverable(pinerrors.New(bizErrorMsg)) // Unrecoverable on business error
				}

				return nil
			},
			// Retry only if the specific pemChangedError is encountered
			retry.RetryIf(func(err error) bool {
				var pErr *pemChangedError
				return errors.As(err, &pErr)
			}),
			retry.Attempts(3),
		)

		// Handle the final error from retry.Do
		if err != nil {
			// If the error is a wrapped business error, return the specific response
			if bizErrorMsg != "" {
				return &importadminpb.CheckImportSourceConnectionResp{Success: false, ErrorMsg: bizErrorMsg}, nil
			}

			// For other errors (RPC failure, GetPublicKey failure, etc.), return the error
			return nil, err
		}

		// Successful check
		log.Info("success to check oss access", zap.Uint64("clusterID", clusterID), zap.String("uri", utils.RedactSourceURI(uri)))
		return &importadminpb.CheckImportSourceConnectionResp{Success: true, Region: region}, nil
	case importadminpb.ImportSourceType_ENUM_AZURE_BLOB:
		httpsURI, rawSas, err := normalizeAzureBlobInput(req.GetSource().GetAzureBlob().GetUri(), req.GetSource().GetAzureBlob().GetSasToken())
		if err != nil {
			return &importadminpb.CheckImportSourceConnectionResp{Success: false, ErrorMsg: err.Error()}, nil
		}
		log.Info("start to check azblob access", zap.Uint64("clusterID", clusterID), zap.String("uri", utils.RedactSourceURI(httpsURI)))
		res, err := agentCli.CheckAzblobAccess(ctx, &agentpb.CheckAzblobAccessReq{AzblobUrl: httpsURI, SasToken: rawSas})
		errMsg := "failed to check azblob access"
		if err != nil {
			log.Error(errMsg, zap.Error(err))
			return &importadminpb.CheckImportSourceConnectionResp{Success: false, ErrorMsg: errMsg}, nil
		}
		if res.GetErrorMsg() != "" {
			return &importadminpb.CheckImportSourceConnectionResp{Success: false, ErrorMsg: res.GetErrorMsg()}, nil
		}
		log.Info("success to check azblob access", zap.Uint64("clusterID", clusterID), zap.String("uri", utils.RedactSourceURI(httpsURI)))
		return &importadminpb.CheckImportSourceConnectionResp{Success: true}, nil
	}
	return &importadminpb.CheckImportSourceConnectionResp{Success: false}, nil
}