package grpc

import (
	"bytes"
	"context"
	"crypto/tls"
	"crypto/x509"
	"database/sql"
	"encoding/json"
	goerr "errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"path"
	"path/filepath"
	"strings"
	"time"

	"github.com/go-sql-driver/mysql"
	"github.com/google/uuid"
	"github.com/pingcap/errors"
	plog "github.com/pingcap/log"
	"github.com/pingcap/tidb/pkg/importsdk"
	lightningconfig "github.com/pingcap/tidb/pkg/lightning/config"
	"github.com/pingcap/tidb/pkg/lightning/log"
	"github.com/pingcap/tidb/pkg/lightning/mydump"
	"github.com/pingcap/tidb/pkg/util/redact"
	dmchecker "github.com/pingcap/tiflow/dm/checker"
	dmconfig "github.com/pingcap/tiflow/dm/config"
	pchecker "github.com/pingcap/tiflow/dm/pkg/checker"
	"github.com/tidbcloud/dataflow-service/dataflow-agent-ng/agent/tidbcloud"
	agent "github.com/tidbcloud/dataflow-service/dataflow-agent-ng/idl/pbgen/proto/ng_agent"
	"github.com/tidbcloud/dataflow-service/dataflow-agent-ng/pkg/alioss"
	"github.com/tidbcloud/dataflow-service/dataflow-agent-ng/pkg/azblob"
	"github.com/tidbcloud/dataflow-service/dataflow-agent-ng/pkg/crypto"
	"github.com/tidbcloud/dataflow-service/dataflow-agent-ng/pkg/gcs"
	"github.com/tidbcloud/dataflow-service/dataflow-agent-ng/pkg/loader"
	pkglog "github.com/tidbcloud/dataflow-service/dataflow-agent-ng/pkg/log"
	"github.com/tidbcloud/dataflow-service/dataflow-agent-ng/pkg/s3"
	"github.com/tidbcloud/pkgv2/bizerr"
	"github.com/tidwall/sjson"
	"go.uber.org/zap"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/types/known/timestamppb"
	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
	"k8s.io/utils/pointer"
)

// wildcardDirURL returns a URL pointing at the directory that contains the
// given wildcard, preserving scheme/host/query. The returned URL ends with "/"
// so JoinPath produces "<dir>/<name>" rather than treating the last segment as
// a file. For wildcards already ending in "/" the same directory is returned.
func wildcardDirURL(u *url.URL) *url.URL {
	res := *u
	dir := path.Dir(res.Path)
	if !strings.HasSuffix(dir, "/") {
		dir += "/"
	}
	res.Path = dir
	// Clear RawPath so String() re-derives the correctly escaped path from Path.
	res.RawPath = ""
	return &res
}

func matchSchemaFileName(path, schemaFileName string) (suffix string, ok bool) {
	baseName := filepath.Base(path)
	if baseName == schemaFileName {
		return "", true
	}
	prefix := schemaFileName + "."
	if !strings.HasPrefix(baseName, prefix) {
		return "", false
	}
	// Keep in sync with mydump compression parser behavior:
	// only names like *.sql.<valid_compression> are accepted.
	if mydump.ParseCompressionOnFileExtension(baseName) == mydump.CompressionNone {
		return "", false
	}
	return strings.TrimPrefix(baseName, schemaFileName), true
}

func init() {
	conf := &plog.Config{Level: "info", File: plog.FileLogConfig{}}
	plogLogger, plogProps, err := plog.InitLogger(conf)
	if err == nil {
		plog.ReplaceGlobals(plogLogger, plogProps)
	}
	// Initialize redact to redact sensitive data in importSDK error messages.
	redact.InitRedact(true)
}

func sqlPayloadErrorMessage(err error) string {
	if err == nil {
		return ""
	}

	var mysqlErr *mysql.MySQLError
	if goerr.As(err, &mysqlErr) {
		return mysqlErr.Error()
	}

	return ""
}

func sourceDBPayloadErrorMessage(err error) string {
	if err == nil {
		return ""
	}
	if goerr.Is(err, context.Canceled) {
		return ""
	}
	if msg := sqlPayloadErrorMessage(err); msg != "" {
		return msg
	}

	errMsg := strings.ToLower(err.Error())
	if goerr.Is(err, context.DeadlineExceeded) &&
		!strings.Contains(errMsg, "dial tcp") &&
		!strings.Contains(errMsg, "i/o timeout") &&
		!strings.Contains(errMsg, "connection refused") &&
		!strings.Contains(errMsg, "connection reset") &&
		!strings.Contains(errMsg, "no such host") &&
		!strings.Contains(errMsg, "network is unreachable") {
		return ""
	}

	var netErr net.Error
	if goerr.As(err, &netErr) {
		return err.Error()
	}

	if strings.Contains(errMsg, "dial tcp") ||
		strings.Contains(errMsg, "i/o timeout") ||
		strings.Contains(errMsg, "connection refused") ||
		strings.Contains(errMsg, "connection reset") ||
		strings.Contains(errMsg, "no such host") ||
		strings.Contains(errMsg, "network is unreachable") {
		return err.Error()
	}

	return ""
}

// Async import task polling uses a two-lane error contract:
//   - bizerr is considered a user-facing/business failure and is surfaced via
//     the task response payload's error_msg field.
//   - all other errors are treated as internal/transport failures and are
//     returned from the RPC directly so callers can handle them separately.
func bizErrMessage(err error) *string {
	if err == nil {
		return nil
	}

	var bErr bizerr.BizError
	if goerr.As(err, &bErr) {
		return pointer.String(bErr.ErrMsg())
	}

	return nil
}

type schemaCreateSDK interface {
	CreateSchemasAndTables(ctx context.Context) error
	CreateSchemaAndTableByName(ctx context.Context, database, table string) error
}

func schemaCreationDBOptions() *tidbcloud.DBOptions {
	return &tidbcloud.DBOptions{
		Params: map[string]string{
			"foreign_key_checks": "0",
		},
	}
}

func executeCreateSchemas(
	ctx context.Context,
	logger *zap.Logger,
	sdk schemaCreateSDK,
	schemas []*agent.CreateSchemasReq_Schema,
) *agent.CreateSchemasResp {
	resp := &agent.CreateSchemasResp{PemChanged: false}
	if len(schemas) == 0 {
		err := sdk.CreateSchemasAndTables(ctx)
		if err != nil {
			logger.Error("failed to create schemas and tables", zap.Error(err))
			failedMsg := fmt.Sprintf("failed to create schemas: err=%v", err)
			resp.ErrorMsg = &failedMsg
		}
		return resp
	}

	for _, schema := range schemas {
		result := &agent.CreateSchemasResp_TableResult{
			Database: schema.GetDatabase(),
			Table:    schema.GetTable(),
		}
		if err := sdk.CreateSchemaAndTableByName(ctx, schema.GetDatabase(), schema.GetTable()); err != nil {
			logger.Error("failed to create schema and table", zap.Error(err))
			failedMsg := fmt.Sprintf(
				"failed to create schemas: schema=%s, table=%s, err=%v",
				schema.GetDatabase(),
				schema.GetTable(),
				err,
			)
			result.Status = agent.CreateSchemasResp_TableResult_STATUS_FAILED
			result.ErrorMsg = &failedMsg
			resp.TableResults = append(resp.TableResults, result)
			continue
		}
		result.Status = agent.CreateSchemasResp_TableResult_STATUS_SUCCEEDED
		resp.TableResults = append(resp.TableResults, result)
	}
	return resp
}

func (s *AgentServer) CheckS3Access(ctx context.Context, req *agent.CheckS3AccessReq) (
	*agent.CheckS3AccessResp, error,
) {
	resp := &agent.CheckS3AccessResp{}
	if err := s3.CheckS3Access(
		ctx,
		req.S3Url, req.RoleArn, req.ExternalId,
		req.AwsAccessKeyId,
		req.AwsSecretAccessKey,
	); err != nil {
		resp.ErrorMsg = err.Error()
		return resp, nil
	}
	u, err := s3.NewS3URL(req.S3Url)
	if err != nil {
		resp.ErrorMsg = err.Error()
		return resp, nil
	}
	resp.Region, err = u.GetBucketRegion(ctx)
	if err != nil {
		resp.ErrorMsg = err.Error()
		return resp, nil
	}
	return resp, nil
}

func (s *AgentServer) CheckGcsAccess(ctx context.Context, req *agent.CheckGcsAccessReq) (
	*agent.CheckGcsAccessResp, error,
) {
	resp := &agent.CheckGcsAccessResp{}
	if err := gcs.CheckGcsAccess(ctx, req.GcsUrl, req.GetGcsOptions().GetImpersonateServiceAccount()); err != nil {
		resp.ErrorMsg = err.Error()
	}
	return resp, nil
}

func (s *AgentServer) CheckAzblobAccess(ctx context.Context, req *agent.CheckAzblobAccessReq) (
	*agent.CheckAzblobAccessResp, error,
) {
	resp := &agent.CheckAzblobAccessResp{}
	if err := azblob.CheckAzblobAccess(ctx, req.AzblobUrl, req.SasToken); err != nil {
		resp.ErrorMsg = err.Error()
	}
	return resp, nil
}

func (s *AgentServer) CheckAliOSSAccess(ctx context.Context, req *agent.CheckAliOSSAccessReq) (
	*agent.CheckAliOSSAccessResp, error,
) {
	resp := &agent.CheckAliOSSAccessResp{}
	// -------------------------------------------------------------------------
	// 1. Decrypt Secret Key and Verify PEM
	// -------------------------------------------------------------------------
	encryptedSecret := req.GetOssAccessKeySecret()
	if encryptedSecret == nil || encryptedSecret.GetData() == nil {
		resp.ErrorMsg = "Encrypted access key secret is missing."
		return resp, nil
	}

	pem := encryptedSecret.GetPem()
	verified, err := s.encryptor.VerifyKeyMatch(pem)
	if err != nil {
		resp.ErrorMsg = errors.Annotate(err, "failed to verify encryption key match").Error()
		return resp, nil
	}
	if !verified {
		resp.PemChanged = true
		resp.ErrorMsg = "Encryption key has been updated. Please re-enter the access key secret."
		return resp, nil
	}

	data := encryptedSecret.GetData()
	secretKey, err := s.encryptor.Decrypt(data)
	if err != nil {
		resp.ErrorMsg = errors.Annotate(err, "failed to decrypt OSS access key secret").Error()
		return resp, nil
	}

	// -------------------------------------------------------------------------
	// 2. Parse OSS URL
	// -------------------------------------------------------------------------
	bucket, pathPrefix, err := alioss.ParseOssUrl(req.GetOssUri())
	if err != nil {
		resp.ErrorMsg = errors.Annotatef(err, "invalid OSS URL: %s", req.GetOssUri()).Error()
		return resp, nil
	}

	// 3. Check OSS Access (List/Get Permissions)
	err = alioss.CheckOSSAccess(
		ctx,
		bucket,
		pathPrefix,
		alioss.WithAccessKey(req.GetOssAccessKeyId(), string(secretKey)),
		alioss.WithEndpoint(req.GetEndpoint()),
	)

	if err != nil {
		resp.ErrorMsg = errors.Annotate(err, "failed to check OSS access permissions (ListObjects/GetObject)").Error()
		return resp, nil
	}

	// 3. Success
	resp.ErrorMsg = ""
	return resp, nil
}

func (s *AgentServer) getClusterPassword(ctx context.Context, cluster *agent.Cluster) (string, error) {
	if !cluster.GetPasswordEncrypted() {
		return cluster.GetPassword(), nil
	}
	pem := cluster.GetEncryptedPassword().GetPem()
	verified, err := s.encryptor.VerifyKeyMatch(pem)
	if err != nil {
		return "", errors.Trace(err)
	}
	if !verified {
		return "", crypto.PEMUpdatedError
	}
	data := cluster.GetEncryptedPassword().GetData()
	passwd, err := s.encryptor.Decrypt(data)
	if err != nil {
		return "", errors.Trace(err)
	}
	return string(passwd), nil
}

func (s *AgentServer) runExecSQLOnCluster(ctx context.Context, req *agent.ExecSQLReq) (
	*agent.ExecSQLResp, error,
) {
	if s.execSQLOnClusterFn != nil {
		return s.execSQLOnClusterFn(ctx, req)
	}
	logger := pkglog.GetLogger(ctx)
	resp := &agent.ExecSQLResp{
		Rows: make([]*agent.Row, 0),
	}
	opts := &tidbcloud.DBOptions{
		ReadTimeout:  req.Cluster.ReadTimeout,
		WriteTimeout: req.Cluster.WriteTimeout,
	}
	passwd, err := s.getClusterPassword(ctx, req.Cluster)
	if err != nil {
		if goerr.Is(err, crypto.PEMUpdatedError) {
			resp.PemChanged = true
			return resp, nil
		}
		return nil, err
	}
	db, err := s.clusterProvider.GetDB(ctx, fmt.Sprint(req.Cluster.ClusterId), req.Cluster.User, passwd, opts)
	if err != nil {
		return nil, err
	}
	defer db.Close() // nolint: errcheck
	for _, sql := range req.SetupSqls {
		logger.With(zap.String("sql", sql)).Info("execute setup SQL")
		_, err = db.ExecContext(ctx, sql)
		if err != nil {
			logger.With(zap.String("sql", sql), zap.Error(err)).Error("exec setup SQL error")
			return nil, err
		}
	}

	// this represents need not return rows, so use exec
	if req.ColumnNum == 0 {
		_, err := db.ExecContext(ctx, req.GetSql())
		if err != nil {
			logger.With(zap.String("sql", req.GetSql()), zap.Error(err)).Error("exec SQL error")
			if msg := sqlPayloadErrorMessage(err); msg != "" {
				resp.ErrorMsg = msg
				return resp, nil
			}
			return nil, err
		}
		return resp, nil
	}
	rows, err := db.QueryContext(ctx, req.Sql)
	if err != nil {
		logger.With(zap.String("sql", req.GetSql()), zap.Error(err)).Error("query SQL error")
		if msg := sqlPayloadErrorMessage(err); msg != "" {
			resp.ErrorMsg = msg
			return resp, nil
		}
		return nil, err
	}
	defer rows.Close() // nolint: errcheck

	resp.ColumnNames, err = rows.Columns()
	if err != nil {
		logger.With(zap.String("sql", req.GetSql()), zap.Error(err)).Error("get columns error")
		return nil, err
	}

	if req.TypedColumns {

		for rows.Next() {
			row := &agent.Row{}
			values := make([]interface{}, len(resp.ColumnNames))
			valuePtrs := make([]interface{}, len(resp.ColumnNames))
			for i := range values {
				valuePtrs[i] = &values[i]
			}
			if err := rows.Scan(valuePtrs...); err != nil {
				logger.With(zap.String("sql", req.GetSql()), zap.Error(err)).Error("scan SQL error")
				return nil, err
			}
			for _, value := range values {
				columnValue := &agent.ColumnValue{}
				switch v := value.(type) {
				case nil:
					columnValue.Value = nil
				case uint, uint8, uint16, uint32, uint64:
					columnValue.Value = &agent.ColumnValue_Int64Value{Int64Value: int64(v.(uint64))}
				case int, int8, int16, int32, int64:
					columnValue.Value = &agent.ColumnValue_Int64Value{Int64Value: int64(v.(int64))}
				case float32, float64:
					columnValue.Value = &agent.ColumnValue_DoubleValue{DoubleValue: float64(v.(float64))}
				case bool:
					columnValue.Value = &agent.ColumnValue_BoolValue{BoolValue: v}
				case string:
					columnValue.Value = &agent.ColumnValue_StringValue{StringValue: v}
				case []uint8:
					columnValue.Value = &agent.ColumnValue_StringValue{StringValue: string(v)}
				case time.Time:
					columnValue.Value = &agent.ColumnValue_TimestampValue{TimestampValue: timestamppb.New(v)}
				default:
					logger.With(zap.String("sql", req.GetSql()), zap.String("columnType", fmt.Sprintf("%T", v))).Error("unsupported typed SQL column value")
					return nil, errors.Errorf("unsupported column type: %T", v)
				}
				row.TypedColumnValues = append(row.TypedColumnValues, columnValue)
			}
			resp.Rows = append(resp.Rows, row)
		}
		return resp, nil
	}

	for rows.Next() {
		row := &agent.Row{
			ColumnValues: make([]string, req.ColumnNum),
		}
		columns := make([]any, req.ColumnNum)
		for i := range columns {
			columns[i] = any(&row.ColumnValues[i])
		}
		if err := rows.Scan(columns...); err != nil {
			logger.With(zap.String("sql", req.GetSql()), zap.Error(err)).Error("scan SQL error")
			return nil, err
		}
		resp.Rows = append(resp.Rows, row)
	}

	return resp, nil
}

func (s *AgentServer) ExecSQLOnCluster(ctx context.Context, req *agent.ExecSQLReq) (
	*agent.ExecSQLResp, error,
) {
	return s.runExecSQLOnCluster(ctx, req)
}

func isImportIntoSQL(sql string) bool {
	trimmed := strings.TrimSpace(sql)
	return strings.HasPrefix(strings.ToUpper(trimmed), "IMPORT INTO ")
}

func importJobIDFromExecSQLResp(resp *agent.ExecSQLResp) (int64, error) {
	if resp == nil || len(resp.GetRows()) == 0 || len(resp.GetColumnNames()) == 0 {
		return 0, errors.New("failed to parse import job id from empty response")
	}
	if len(resp.GetRows()[0].GetTypedColumnValues()) != len(resp.GetColumnNames()) {
		return 0, errors.New("failed to parse import job id from malformed response")
	}
	for i, name := range resp.GetColumnNames() {
		if name != "Job_ID" {
			continue
		}
		return resp.GetRows()[0].GetTypedColumnValues()[i].GetInt64Value(), nil
	}
	return 0, errors.New("Job_ID column not found")
}

func validateCreateImportIntoJobTask(taskID string, cluster *agent.Cluster, sql string) error {
	if taskID == "" {
		return status.Error(codes.InvalidArgument, "task_id is required")
	}
	if cluster == nil {
		return status.Error(codes.InvalidArgument, "cluster is required")
	}
	if !isImportIntoSQL(sql) {
		return status.Error(codes.InvalidArgument, "only IMPORT INTO sql is supported")
	}
	return nil
}

func (s *AgentServer) createImportIntoJobTask(
	ctx context.Context,
	taskID string,
	cluster *agent.Cluster,
	sql string,
	setupSQLs []string,
) *agent.CreateImportIntoJobTaskResp {
	s.importTaskManager.CreateIfAbsent(ctx, taskID, func(taskCtx context.Context) (any, error) {
		logger := pkglog.GetLogger(taskCtx).With(zap.String("task_id", taskID))
		resp, err := s.runExecSQLOnCluster(taskCtx, &agent.ExecSQLReq{
			Cluster:      cluster,
			Sql:          sql,
			SetupSqls:    setupSQLs,
			TypedColumns: true,
			ColumnNum:    13,
		})
		if err != nil {
			logger.Error("failed to run import into job", zap.Error(err))
			return nil, err
		}
		if resp.GetErrorMsg() != "" {
			logger.Error("failed to execute import into job", zap.String("error_msg", resp.GetErrorMsg()))
			return nil, bizerr.New(bizerr.ErrCodeGeneralBusiness, "failed to execute import into job: %s", resp.GetErrorMsg())
		}
		jobID, err := importJobIDFromExecSQLResp(resp)
		if err != nil {
			logger.Error("failed to parse import job id from exec SQL response", zap.Error(err))
			return nil, err
		}
		return &agent.GetImportIntoJobTaskResp_Result{JobId: jobID}, nil
	})
	return &agent.CreateImportIntoJobTaskResp{TaskId: taskID}
}

func (s *AgentServer) CreateImportIntoJobTask(ctx context.Context, req *agent.CreateImportIntoJobTaskReq) (*agent.CreateImportIntoJobTaskResp, error) {
	if err := validateCreateImportIntoJobTask(req.GetTaskId(), req.GetCluster(), req.GetSql()); err != nil {
		return nil, err
	}
	return s.createImportIntoJobTask(ctx, req.GetTaskId(), req.GetCluster(), req.GetSql(), req.GetSetupSqls()), nil
}

func (s *AgentServer) BatchCreateImportIntoJobTasks(ctx context.Context, req *agent.BatchCreateImportIntoJobTasksReq) (*agent.BatchCreateImportIntoJobTasksResp, error) {
	if req.GetCluster() == nil {
		return nil, status.Error(codes.InvalidArgument, "cluster is required")
	}
	for _, task := range req.GetTasks() {
		if err := validateCreateImportIntoJobTask(task.GetTaskId(), req.GetCluster(), task.GetSql()); err != nil {
			return nil, err
		}
	}
	resp := &agent.BatchCreateImportIntoJobTasksResp{
		Tasks: make([]*agent.CreateImportIntoJobTaskResp, 0, len(req.GetTasks())),
	}
	for _, task := range req.GetTasks() {
		resp.Tasks = append(resp.Tasks, s.createImportIntoJobTask(ctx, task.GetTaskId(), req.GetCluster(), task.GetSql(), task.GetSetupSqls()))
	}
	return resp, nil
}

func (s *AgentServer) GetImportIntoJobTask(ctx context.Context, req *agent.GetImportIntoJobTaskReq) (*agent.GetImportIntoJobTaskResp, error) {
	resp, ok, err := s.getImportIntoJobTask(req.GetTaskId())
	if !ok {
		return nil, status.Error(codes.NotFound, "import into job task not found")
	}
	return resp, err
}

func (s *AgentServer) getImportIntoJobTask(taskID string) (*agent.GetImportIntoJobTaskResp, bool, error) {
	task, ok := s.importTaskManager.Get(taskID)
	if !ok {
		return nil, false, nil
	}

	resp := &agent.GetImportIntoJobTaskResp{
		TaskId: taskID,
		Status: toProtoImportIntoJobTaskStatus(task.Status),
	}
	if msg := bizErrMessage(task.Error); msg != nil {
		resp.ErrorMsg = msg
		return resp, true, nil
	}
	if task.Error != nil {
		return nil, true, task.Error
	}
	if result, ok := task.Result.(*agent.GetImportIntoJobTaskResp_Result); ok {
		resp.Result = result
	}
	return resp, true, nil
}

func (s *AgentServer) BatchGetImportIntoJobTasks(_ context.Context, req *agent.BatchGetImportIntoJobTasksReq) (*agent.BatchGetImportIntoJobTasksResp, error) {
	resp := &agent.BatchGetImportIntoJobTasksResp{
		Tasks: make([]*agent.GetImportIntoJobTaskResp, 0, len(req.GetTaskIds())),
	}
	for _, taskID := range req.GetTaskIds() {
		taskResp, ok, err := s.getImportIntoJobTask(taskID)
		if !ok {
			continue
		}
		if err != nil {
			continue
		}
		resp.Tasks = append(resp.Tasks, taskResp)
	}
	return resp, nil
}

func (s *AgentServer) BatchCancelImportIntoJobs(ctx context.Context, req *agent.BatchCancelImportIntoJobsReq) (*agent.BatchCancelImportIntoJobsResp, error) {
	if req.GetCluster() == nil {
		return nil, status.Error(codes.InvalidArgument, "cluster is required")
	}
	resp := &agent.BatchCancelImportIntoJobsResp{
		Results: make([]*agent.BatchCancelImportIntoJobsResp_Result, 0, len(req.GetJobIds())),
	}
	for _, jobID := range req.GetJobIds() {
		result := &agent.BatchCancelImportIntoJobsResp_Result{JobId: jobID}
		sql := fmt.Sprintf("CANCEL IMPORT JOB %d", jobID)
		sqlResp, err := s.runExecSQLOnCluster(ctx, &agent.ExecSQLReq{
			Cluster: req.GetCluster(),
			Sql:     sql,
		})
		if err != nil {
			return nil, err
		}
		if sqlResp.GetPemChanged() {
			return nil, status.Error(codes.Aborted, "cluster pem changed")
		}
		if sqlResp.GetErrorMsg() != "" {
			result.ErrorMsg = pointer.String(sqlResp.GetErrorMsg())
		}
		resp.Results = append(resp.Results, result)
	}
	return resp, nil
}

func (s *AgentServer) GetImportSchemasMeta(ctx context.Context, req *agent.GetImportSchemasMetaReq) (
	*agent.GetImportSchemasMetaResp, error,
) {
	return nil, errors.New("not implemented")
}

func (s *AgentServer) getImportFiles(ctx context.Context, req *agent.GetImportFilesReq) (*agent.GetImportFilesResp, error) {
	logger := plog.L().With(zap.Uint64("cluster-id", req.ClusterId))
	sourceUri := req.SourceUri
	var tableMetas []*agent.GetImportFilesResp_TableMeta
	if len(req.GetTableMeta()) == 0 {
		options := []importsdk.SDKOption{
			importsdk.WithEstimateRealSize(false),
			importsdk.WithLogger(log.Wrap(logger)),
		}
		if req.GetMaxScanFiles() > 0 {
			options = append(options, importsdk.WithMaxScanFiles(int(req.GetMaxScanFiles())))
		}
		sdk, err := importsdk.NewImportSDK(ctx, sourceUri, nil, options...)
		if err != nil {
			logger.Error("failed to create import SDK", zap.Error(err))
			return nil, bizerr.New(bizerr.ErrCodeGeneralBusiness, "failed to create import SDK: %v", err)
		}
		metas, err := sdk.GetTableMetas(ctx)
		if err != nil {
			logger.Error("failed to get table metas", zap.Error(err))
			return nil, bizerr.New(bizerr.ErrCodeGeneralBusiness, "failed to get table metas: %v", err)
		}
		for _, meta := range metas {
			dataFiles := []*agent.GetImportFilesResp_DataFileMeta{}
			for _, f := range meta.DataFiles {
				df := &agent.GetImportFilesResp_DataFileMeta{
					Path:        f.Path,
					Size:        uint64(f.Size),
					SourceType:  agent.GetImportFilesResp_SourceType(f.Format),
					Compression: agent.GetImportFilesResp_Compression(f.Compression),
				}
				dataFiles = append(dataFiles, df)
			}
			tableMeta := &agent.GetImportFilesResp_TableMeta{
				Database:     meta.Database,
				Table:        meta.Table,
				DataFiles:    dataFiles,
				TotalSize:    uint64(meta.TotalSize),
				WildcardPath: meta.WildcardPath,
				SchemaFile:   meta.SchemaFile,
			}
			tableMetas = append(tableMetas, tableMeta)
		}
	} else {
		parsedSourceUrl, err := url.Parse(req.SourceUri)
		if err != nil {
			return nil, err
		}
		for _, inMeta := range req.GetTableMeta() {
			customWildcardParsed, err := url.Parse(inMeta.GetWildcardPath())
			if err != nil {
				return nil, err
			}
			customWildcardParsed.RawQuery = parsedSourceUrl.RawQuery
			if customWildcardParsed.Scheme == "oss" {
				customWildcardParsed.Scheme = "s3"
			}
			files, fileSize, err := s.scanDataFiles(ctx, customWildcardParsed.String(), int(req.GetMaxScanFiles()))
			if err != nil {
				return nil, err
			}
			dataFiles := []*agent.GetImportFilesResp_DataFileMeta{}
			for _, f := range files {
				df := &agent.GetImportFilesResp_DataFileMeta{
					Path:        f.Path,
					Size:        uint64(f.FileSize),
					SourceType:  agent.GetImportFilesResp_SourceType(f.Type),
					Compression: agent.GetImportFilesResp_Compression(f.Compression),
				}
				dataFiles = append(dataFiles, df)
			}
			tableMeta := &agent.GetImportFilesResp_TableMeta{
				Database:     inMeta.Database,
				Table:        inMeta.Table,
				DataFiles:    dataFiles,
				TotalSize:    uint64(fileSize),
				WildcardPath: inMeta.WildcardPath,
			}
			// Schema files live next to the data files, i.e. in the directory
			// the wildcard resolves into. Deriving the search base from the
			// source URI root only works when the wildcard is at that root and
			// misses schemas whenever the user points the wildcard at a
			// subdirectory. This mirrors how the auto-detect branch (via
			// importsdk / mydump loader) pairs schema and data within the
			// same directory.
			schemaBase := wildcardDirURL(customWildcardParsed)
			schemaFileName := fmt.Sprintf("%s.%s-schema.sql", inMeta.Database, inMeta.Table)
			baseSchemaURI := schemaBase.JoinPath(schemaFileName).String()
			schemaPatternURI := schemaBase.JoinPath(schemaFileName + "*").String()
			logger.With(zap.String("schemaPatternUri", schemaPatternURI)).Info("searching schema by prefix")
			schemaFiles, _, err := s.scanDataFiles(ctx, schemaPatternURI, int(req.GetMaxScanFiles()))
			if err != nil {
				if !goerr.Is(err, loader.ErrFileNotFound) {
					return nil, err
				}
			} else {
				for _, schemaFile := range schemaFiles {
					suffix, matched := matchSchemaFileName(schemaFile.Path, schemaFileName)
					if !matched {
						continue
					}
					if suffix == "" {
						tableMeta.SchemaFile = baseSchemaURI
					} else {
						tableMeta.SchemaFile = schemaBase.JoinPath(schemaFileName + suffix).String()
					}
					break
				}
			}
			tableMetas = append(tableMetas, tableMeta)
		}
	}
	resp := &agent.GetImportFilesResp{
		TableMeta: tableMetas,
	}
	return resp, nil
}

func (s *AgentServer) GetImportFiles(ctx context.Context, req *agent.GetImportFilesReq) (*agent.GetImportFilesResp, error) {
	return s.getImportFiles(ctx, req)
}

func (s *AgentServer) CreateImportFilesTask(ctx context.Context, req *agent.CreateImportFilesTaskReq) (*agent.CreateImportFilesTaskResp, error) {
	if req.GetTaskId() == "" {
		return nil, status.Error(codes.InvalidArgument, "task_id is required")
	}
	s.importTaskManager.CreateIfAbsent(ctx, req.GetTaskId(), func(taskCtx context.Context) (any, error) {
		return s.getImportFiles(taskCtx, &agent.GetImportFilesReq{
			ClusterId:    req.GetClusterId(),
			SourceUri:    req.GetSourceUri(),
			MaxScanFiles: req.GetMaxScanFiles(),
			TableMeta:    req.GetTableMeta(),
		})
	})
	return &agent.CreateImportFilesTaskResp{TaskId: req.GetTaskId()}, nil
}

func (s *AgentServer) GetImportFilesTask(ctx context.Context, req *agent.GetImportFilesTaskReq) (*agent.GetImportFilesTaskResp, error) {
	task, ok := s.importTaskManager.Get(req.GetTaskId())
	if !ok {
		return nil, status.Error(codes.NotFound, "import files task not found")
	}

	resp := &agent.GetImportFilesTaskResp{
		TaskId: req.GetTaskId(),
		Status: toProtoImportFilesTaskStatus(task.Status),
	}
	if msg := bizErrMessage(task.Error); msg != nil {
		resp.ErrorMsg = msg
		return resp, nil
	}
	if task.Error != nil {
		return nil, task.Error
	}
	if result, ok := task.Result.(*agent.GetImportFilesResp); ok {
		resp.Result = result
	}
	return resp, nil
}

func (s *AgentServer) runCreateSchemas(ctx context.Context, req *agent.CreateSchemasReq) (*agent.CreateSchemasResp, error) {
	if s.createSchemasFn != nil {
		return s.createSchemasFn(ctx, req)
	}
	logger := pkglog.GetLogger(ctx).With(zap.Uint64("cluster-id", req.ClusterId))
	var db *sql.DB
	var err error
	dbOpts := schemaCreationDBOptions()
	if req.GetCluster() == nil {
		db, err = s.clusterProvider.GetDB(ctx, fmt.Sprint(req.ClusterId), req.User, req.Password, dbOpts)
	} else {
		user := req.GetCluster().GetUser()
		passwd, err := s.getClusterPassword(ctx, req.Cluster)
		if err != nil {
			if goerr.Is(err, crypto.PEMUpdatedError) {
				return &agent.CreateSchemasResp{PemChanged: true}, nil
			}
			return nil, err
		}
		db, err = s.clusterProvider.GetDB(ctx, fmt.Sprint(req.ClusterId), user, passwd, dbOpts)
	}
	if err != nil {
		return nil, err
	}
	defer db.Close() // nolint: errcheck

	options := []importsdk.SDKOption{
		importsdk.WithEstimateRealSize(false),
		importsdk.WithLogger(log.Wrap(logger)),
	}
	sourceUri := req.SourceUri
	sdk, err := importsdk.NewImportSDK(ctx, sourceUri, db, options...)
	if err != nil {
		return nil, bizerr.New(bizerr.ErrCodeGeneralBusiness, "failed to create import SDK: %v", err)
	}
	return executeCreateSchemas(ctx, logger, sdk, req.GetSchemas()), nil
}

func (s *AgentServer) CreateSchemas(ctx context.Context, req *agent.CreateSchemasReq) (*agent.CreateSchemasResp, error) {
	return s.runCreateSchemas(ctx, req)
}

func (s *AgentServer) CreateSchemasTask(ctx context.Context, req *agent.CreateSchemasTaskReq) (*agent.CreateSchemasTaskResp, error) {
	if req.GetTaskId() == "" {
		return nil, status.Error(codes.InvalidArgument, "task_id is required")
	}
	if req.GetCluster() == nil {
		return nil, status.Error(codes.InvalidArgument, "cluster is required")
	}
	s.importTaskManager.CreateIfAbsent(ctx, req.GetTaskId(), func(taskCtx context.Context) (any, error) {
		return s.runCreateSchemas(taskCtx, &agent.CreateSchemasReq{
			ClusterId: req.GetClusterId(),
			SourceUri: req.GetSourceUri(),
			Schemas:   req.GetSchemas(),
			Cluster:   req.GetCluster(),
		})
	})
	return &agent.CreateSchemasTaskResp{TaskId: req.GetTaskId()}, nil
}

func (s *AgentServer) GetSchemasTask(ctx context.Context, req *agent.GetSchemasTaskReq) (*agent.GetSchemasTaskResp, error) {
	task, ok := s.importTaskManager.Get(req.GetTaskId())
	if !ok {
		return nil, status.Error(codes.NotFound, "create schemas task not found")
	}

	resp := &agent.GetSchemasTaskResp{
		TaskId: req.GetTaskId(),
		Status: toProtoCreateSchemasTaskStatus(task.Status),
	}
	if msg := bizErrMessage(task.Error); msg != nil {
		resp.ErrorMsg = msg
		return resp, nil
	}
	if task.Error != nil {
		return nil, task.Error
	}
	if result, ok := task.Result.(*agent.CreateSchemasResp); ok {
		resp.Result = result
		if resp.ErrorMsg == nil && result.GetErrorMsg() != "" {
			resp.ErrorMsg = result.ErrorMsg
		}
	}
	return resp, nil
}

func (s *AgentServer) CreateImportDataSizeEstimateTask(ctx context.Context, req *agent.CreateImportDataSizeEstimateTaskReq) (*agent.CreateImportDataSizeEstimateTaskResp, error) {
	if req.GetTaskId() == "" {
		return nil, status.Error(codes.InvalidArgument, "task_id is required")
	}
	s.importTaskManager.CreateIfAbsent(ctx, req.GetTaskId(), func(taskCtx context.Context) (any, error) {
		return estimateImportDataSize(taskCtx, req)
	})
	return &agent.CreateImportDataSizeEstimateTaskResp{TaskId: req.GetTaskId()}, nil
}

func (s *AgentServer) GetImportDataSizeEstimateTask(ctx context.Context, req *agent.GetImportDataSizeEstimateTaskReq) (*agent.GetImportDataSizeEstimateTaskResp, error) {
	task, ok := s.importTaskManager.Get(req.GetTaskId())
	if !ok {
		return nil, status.Error(codes.NotFound, "estimate task not found")
	}

	resp := &agent.GetImportDataSizeEstimateTaskResp{
		TaskId: req.GetTaskId(),
		Status: toProtoImportAsyncTaskStatus(task.Status),
	}
	if msg := bizErrMessage(task.Error); msg != nil {
		resp.ErrorMsg = msg
		return resp, nil
	}
	if task.Error != nil {
		return nil, task.Error
	}
	if result, ok := task.Result.(*agent.GetImportDataSizeEstimateTaskResp_Result); ok {
		resp.Result = result
	}
	return resp, nil
}

func (s *AgentServer) GetPublicKey(context.Context, *agent.GetPublicKeyRequest) (*agent.GetPublicKeyResponse, error) {
	pem, err := s.encryptor.ExportPublicKeyAsPEM()
	if err != nil {
		return nil, err
	}
	return &agent.GetPublicKeyResponse{Pem: pem}, nil
}

func getDMMasterHost(ctx context.Context, namespace string) (string, error) {
	kubeconfig, err := rest.InClusterConfig()
	if err != nil {
		return "", err
	}
	kubeCli, err := kubernetes.NewForConfig(kubeconfig)
	if err != nil {
		return "", err
	}

	svcLabels := &v1.LabelSelector{
		MatchLabels: map[string]string{
			"app.kubernetes.io/component": "dm-master",
		},
	}

	selector, err := v1.LabelSelectorAsSelector(svcLabels)
	if err != nil {
		return "", errors.Trace(err)
	}

	svcList, err := kubeCli.CoreV1().Services(namespace).List(ctx, v1.ListOptions{
		LabelSelector: selector.String(),
	})
	if err != nil {
		return "", errors.Trace(err)
	}

	var targetSvcName string
	var targetNamespace string
	targetPort := int32(8261)

	for _, svc := range svcList.Items {
		for _, port := range svc.Spec.Ports {
			if port.Port == targetPort {
				targetSvcName = svc.Name
				targetNamespace = svc.Namespace
				break
			}
		}
		if targetSvcName != "" {
			break
		}
	}

	if targetSvcName == "" {
		return "", errors.Errorf("no service found with port %d", targetPort)
	}
	targetHost := fmt.Sprintf("%s.%s.svc", targetSvcName, targetNamespace)

	return targetHost, nil
}

func estimateImportDataSize(ctx context.Context, req *agent.CreateImportDataSizeEstimateTaskReq) (*agent.GetImportDataSizeEstimateTaskResp_Result, error) {
	logger := pkglog.GetLogger(ctx).With(zap.String("task_id", req.GetTaskId()))
	options := []importsdk.SDKOption{
		importsdk.WithEstimateRealSize(false),
		importsdk.WithLogger(log.Wrap(pkglog.GetLogger(ctx))),
	}
	if req.GetMaxScanFiles() > 0 {
		options = append(options, importsdk.WithMaxScanFiles(int(req.GetMaxScanFiles())))
	}
	if csvCfg := req.GetCsvConfig(); csvCfg != nil {
		options = append(options, importsdk.WithCSVConfig(protoCSVConfigToSDKConfig(csvCfg)))
	}

	sdk, err := importsdk.NewImportSDK(ctx, req.GetSourceUri(), nil, options...)
	if err != nil {
		logger.Error("failed to create import SDK", zap.Error(err))
		return nil, bizerr.New(bizerr.ErrCodeGeneralBusiness, "failed to create import SDK: %v", err)
	}
	defer sdk.Close() // nolint: errcheck

	estimate, err := sdk.EstimateImportDataSize(ctx)
	if err != nil {
		logger.Error("failed to estimate import data size", zap.Error(err))
		return nil, bizerr.New(bizerr.ErrCodeGeneralBusiness, "failed to estimate import data size: %v", err)
	}
	return filterImportDataSizeEstimateResult(convertImportDataSizeEstimate(estimate), req.GetTableMeta()), nil
}

func protoCSVConfigToSDKConfig(pb *agent.CsvConfig) lightningconfig.CSVConfig {
	// Start from the default CSV config so that unset fields retain safe
	// defaults (e.g. separator=",", delimiter='"') instead of empty strings
	// which would cause a panic in NewCSVParser.
	defaultCfg := lightningconfig.NewConfig()
	cfg := defaultCfg.Mydumper.CSV

	if pb.GetSeparator() != "" {
		cfg.FieldsTerminatedBy = pb.GetSeparator()
	}
	if pb.GetDelimiter() != "" {
		cfg.FieldsEnclosedBy = pb.GetDelimiter()
	}
	cfg.Header = pb.GetHeader()
	cfg.BackslashEscape = pb.GetBackslashEscape()
	cfg.TrimLastEmptyField = pb.GetTrimLastSeparator()
	if pb.NotNull != nil {
		cfg.NotNull = pb.GetNotNull()
	}
	if pb.NullValue != nil {
		cfg.FieldNullDefinedBy = lightningconfig.StringOrStringSlice{pb.GetNullValue()}
	}
	if pb.EscapeChar != nil {
		cfg.FieldsEscapedBy = pb.GetEscapeChar()
	}
	if pb.LinesTerminatedBy != nil {
		cfg.LinesTerminatedBy = pb.GetLinesTerminatedBy()
	}
	return cfg
}

func convertImportDataSizeEstimate(estimate *importsdk.ImportDataSizeEstimate) *agent.GetImportDataSizeEstimateTaskResp_Result {
	if estimate == nil {
		return nil
	}
	resp := &agent.GetImportDataSizeEstimateTaskResp_Result{
		TotalSourceSize: uint64(estimate.TotalSourceSize),
		TotalTikvSize:   uint64(estimate.TotalTiKVSize),
		Tables:          make([]*agent.GetImportDataSizeEstimateTaskResp_TableEstimate, 0, len(estimate.Tables)),
	}
	for _, tbl := range estimate.Tables {
		resp.Tables = append(resp.Tables, &agent.GetImportDataSizeEstimateTaskResp_TableEstimate{
			Database:   tbl.Database,
			Table:      tbl.Table,
			SourceSize: uint64(tbl.SourceSize),
			TikvSize:   uint64(tbl.TiKVSize),
		})
	}
	return resp
}

func filterImportDataSizeEstimateResult(
	result *agent.GetImportDataSizeEstimateTaskResp_Result,
	tableMetas []*agent.CreateImportDataSizeEstimateTaskReq_TableMeta,
) *agent.GetImportDataSizeEstimateTaskResp_Result {
	if result == nil || len(tableMetas) == 0 {
		return result
	}

	selectedTables := make(map[string]struct{}, len(tableMetas))
	for _, tableMeta := range tableMetas {
		if tableMeta == nil {
			continue
		}
		selectedTables[buildImportEstimateTableKey(tableMeta.GetDatabase(), tableMeta.GetTable())] = struct{}{}
	}
	if len(selectedTables) == 0 {
		return result
	}

	filtered := &agent.GetImportDataSizeEstimateTaskResp_Result{
		Tables: make([]*agent.GetImportDataSizeEstimateTaskResp_TableEstimate, 0, len(result.GetTables())),
	}
	for _, table := range result.GetTables() {
		if _, ok := selectedTables[buildImportEstimateTableKey(table.GetDatabase(), table.GetTable())]; !ok {
			continue
		}
		filtered.TotalSourceSize += table.GetSourceSize()
		filtered.TotalTikvSize += table.GetTikvSize()
		filtered.Tables = append(filtered.Tables, table)
	}
	return filtered
}

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

func toProtoImportAsyncTaskStatus(status importAsyncTaskStatus) agent.ImportDataSizeEstimateTaskStatus {
	switch status {
	case importAsyncTaskStatusRunning:
		return agent.ImportDataSizeEstimateTaskStatus_IMPORT_DATA_SIZE_ESTIMATE_TASK_STATUS_RUNNING
	case importAsyncTaskStatusSucceeded:
		return agent.ImportDataSizeEstimateTaskStatus_IMPORT_DATA_SIZE_ESTIMATE_TASK_STATUS_SUCCEEDED
	case importAsyncTaskStatusFailed:
		return agent.ImportDataSizeEstimateTaskStatus_IMPORT_DATA_SIZE_ESTIMATE_TASK_STATUS_FAILED
	default:
		return agent.ImportDataSizeEstimateTaskStatus_IMPORT_DATA_SIZE_ESTIMATE_TASK_STATUS_UNSPECIFIED
	}
}

func toProtoImportFilesTaskStatus(status importAsyncTaskStatus) agent.ImportFilesTaskStatus {
	switch status {
	case importAsyncTaskStatusRunning:
		return agent.ImportFilesTaskStatus_IMPORT_FILES_TASK_STATUS_RUNNING
	case importAsyncTaskStatusSucceeded:
		return agent.ImportFilesTaskStatus_IMPORT_FILES_TASK_STATUS_SUCCEEDED
	case importAsyncTaskStatusFailed:
		return agent.ImportFilesTaskStatus_IMPORT_FILES_TASK_STATUS_FAILED
	default:
		return agent.ImportFilesTaskStatus_IMPORT_FILES_TASK_STATUS_UNSPECIFIED
	}
}

func toProtoCreateSchemasTaskStatus(status importAsyncTaskStatus) agent.CreateSchemasTaskStatus {
	switch status {
	case importAsyncTaskStatusRunning:
		return agent.CreateSchemasTaskStatus_CREATE_SCHEMAS_TASK_STATUS_RUNNING
	case importAsyncTaskStatusSucceeded:
		return agent.CreateSchemasTaskStatus_CREATE_SCHEMAS_TASK_STATUS_SUCCEEDED
	case importAsyncTaskStatusFailed:
		return agent.CreateSchemasTaskStatus_CREATE_SCHEMAS_TASK_STATUS_FAILED
	default:
		return agent.CreateSchemasTaskStatus_CREATE_SCHEMAS_TASK_STATUS_UNSPECIFIED
	}
}

func toProtoImportIntoJobTaskStatus(status importAsyncTaskStatus) agent.ImportIntoJobTaskStatus {
	switch status {
	case importAsyncTaskStatusRunning:
		return agent.ImportIntoJobTaskStatus_IMPORT_INTO_JOB_TASK_STATUS_RUNNING
	case importAsyncTaskStatusSucceeded:
		return agent.ImportIntoJobTaskStatus_IMPORT_INTO_JOB_TASK_STATUS_SUCCEEDED
	case importAsyncTaskStatusFailed:
		return agent.ImportIntoJobTaskStatus_IMPORT_INTO_JOB_TASK_STATUS_FAILED
	default:
		return agent.ImportIntoJobTaskStatus_IMPORT_INTO_JOB_TASK_STATUS_UNSPECIFIED
	}
}

var resolveDMTargetTiDBHostFunc = resolveDMTargetTiDBHost

func patchDMTargetConn(ctx context.Context, body []byte, clusterId string) ([]byte, error) {
	if clusterId == "" {
		return body, nil
	}
	host, err := resolveDMTargetTiDBHostFunc(ctx, clusterId)
	if err != nil {
		return nil, err
	}
	res, err := sjson.SetBytes(body, "task.target_config.host", host)
	if err != nil {
		return nil, err
	}
	return res, nil
}

func resolveDMTargetTiDBHost(ctx context.Context, clusterId string) (string, error) {
	if clusterId == "" {
		return "", errors.New("clusterId is required")
	}
	kubeconfig, err := rest.InClusterConfig()
	if err != nil {
		return "", err
	}
	kubeCli, err := kubernetes.NewForConfig(kubeconfig)
	if err != nil {
		return "", err
	}
	podLabels := &v1.LabelSelector{
		MatchLabels: map[string]string{
			"pingcap.com/component":      "tidb",
			"tags.tidbcloud.com/cluster": fmt.Sprint(clusterId),
		},
	}
	selector, err := v1.LabelSelectorAsSelector(podLabels)
	if err != nil {
		return "", errors.Trace(err)
	}
	podList, err := kubeCli.CoreV1().Pods("").List(ctx, v1.ListOptions{LabelSelector: selector.String()})
	if err != nil {
		return "", errors.Trace(err)
	}
	if len(podList.Items) == 0 {
		return "", errors.Errorf("the user cluster has no tidb pods yet")
	}
	// a user cluster may contain several tidb groups
	// any one can be used, for simplicity we choose the first group
	pod := podList.Items[0]
	ns := pod.Namespace
	groupName := pod.Labels["pingcap.com/group"]
	svcName := groupName + "-tidb"
	host := fmt.Sprintf("%s.%s.svc", svcName, ns)
	return host, nil
}

func (s *AgentServer) GetTiDBHost(ctx context.Context, req *agent.GetTiDBHostRequest) (*agent.GetTiDBHostResponse, error) {
	host, err := resolveDMTargetTiDBHostFunc(ctx, req.GetClusterId())
	if err != nil {
		return nil, err
	}
	return &agent.GetTiDBHostResponse{Host: host}, nil
}

func (s *AgentServer) DMOpenAPIProxy(ctx context.Context, req *agent.DMOpenAPIProxyRequest) (*agent.DMOpenAPIProxyResponse, error) {
	host, err := getDMMasterHost(ctx, req.GetDmNamespace())
	if err != nil {
		return nil, err
	}
	method := req.GetMethod()
	requestUrl := fmt.Sprintf("http://%s:8261%s", host, req.GetPath())
	logger := pkglog.GetLogger(ctx).With(zap.String("method", req.GetMethod()), zap.String("url", requestUrl))
	body := req.GetBody()
	if req.GetPatchTargetConn() {
		logger.Info("patching target conn")
		var err error
		body, err = patchDMTargetConn(ctx, body, req.GetClusterId())
		if err != nil {
			return nil, err
		}
	}
	logger.With(zap.ByteString("body", body)).Info("start request")
	httpReq, err := http.NewRequest(method, requestUrl, bytes.NewBuffer(body))
	if err != nil {
		return nil, err
	}
	httpReq.Header.Set("Content-Type", "application/json")
	client := &http.Client{Timeout: 10 * time.Second}
	resp, err := client.Do(httpReq)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	logger.With(zap.ByteString("body", respBody)).Info("success to get response body")
	return &agent.DMOpenAPIProxyResponse{
		Body:   respBody,
		Status: int64(resp.StatusCode),
	}, nil
}

func (s *AgentServer) PrecheckDMJob(ctx context.Context, req *agent.PrecheckDMJobReq) (*agent.PrecheckDMJobResp, error) {
	logger := pkglog.GetLogger(ctx)
	logger.Info("validate dm job",
		zap.String("req_id", req.GetReqId()),
		zap.Int("task_config_count", len(req.GetTaskConfig())),
	)

	resp := &agent.PrecheckDMJobResp{}
	targetHost, err := resolveDMTargetTiDBHost(ctx, req.GetClusterId())
	if err != nil {
		logger.Error("failed to resolve target host for precheck", zap.Error(err))
		return nil, err
	}

	configs, err := buildPrecheckConfigs(req, targetHost)
	if err != nil {
		logger.Error("fail to build precheck configs", zap.Error(err))
		return nil, precheckDMJobClientError(err)
	}
	if err := validateExtraCheckingItems(req.GetExtraCheckingItems()); err != nil {
		logger.Error("invalid extra dm precheck item", zap.Error(err))
		return nil, precheckDMJobClientError(err)
	}

	results, err := dmchecker.RunCheckOnConfigs(
		ctx,
		configs,
		req.GetFullInstanceMigrate(),
		10,
		10,
		req.GetExtraCheckingItems()...,
	)
	if err != nil {
		logger.Error("fail to validate dm job", zap.Error(err))
		return precheckDMJobFailureResp(err)
	}

	logger.Info("precheck", zap.Reflect("precheck result", results))
	resultJSON, err := json.Marshal(results)
	if err != nil {
		logger.Error("fail to marshal check results", zap.Error(err))
		return nil, err
	}
	resp.CheckResults = resultJSON
	return resp, nil
}

// dmPrecheckInitializationDesc is a wire marker consumed by dataflow-service-ng
// to recognize agent-side initialization failures encoded in CheckResults.
const dmPrecheckInitializationDesc = "precheck initialization"

func precheckDMJobFailureResp(err error) (*agent.PrecheckDMJobResp, error) {
	if err == nil {
		return &agent.PrecheckDMJobResp{}, nil
	}
	resultJSON, marshalErr := json.Marshal(&pchecker.Results{
		Results: []*pchecker.Result{
			{
				Desc:  dmPrecheckInitializationDesc,
				State: pchecker.StateFailure,
				Errors: []*pchecker.Error{
					{
						Severity: pchecker.StateFailure,
						ShortErr: err.Error(),
					},
				},
			},
		},
		Summary: &pchecker.ResultSummary{
			Passed: false,
			Total:  1,
			Failed: 1,
		},
	})
	if marshalErr != nil {
		return nil, marshalErr
	}
	return &agent.PrecheckDMJobResp{CheckResults: resultJSON}, nil
}

func precheckDMJobClientError(err error) error {
	if err == nil {
		return nil
	}
	return status.Error(codes.InvalidArgument, err.Error())
}

func validateExtraCheckingItems(items []string) error {
	for _, item := range items {
		if err := dmconfig.ValidateCheckingItem(item); err != nil {
			return err
		}
	}
	return nil
}

func buildPrecheckConfigs(req *agent.PrecheckDMJobReq, targetHost string) ([]*dmconfig.SubTaskConfig, error) {
	if len(req.GetSourcePasswords()) != len(req.GetTaskConfig()) {
		return nil, errors.Errorf(
			"source_passwords count %d does not match task_config count %d",
			len(req.GetSourcePasswords()),
			len(req.GetTaskConfig()),
		)
	}

	configs := make([]*dmconfig.SubTaskConfig, len(req.GetTaskConfig()))
	for i, taskCfg := range req.GetTaskConfig() {
		subCfg := &dmconfig.SubTaskConfig{}
		if err := subCfg.Decode(taskCfg, false); err != nil {
			return nil, err
		}
		if subCfg.To.Host == "" {
			subCfg.To.Host = targetHost
		}
		subCfg.From.Password = req.GetSourcePasswords()[i]
		subCfg.To.Password = req.GetClusterPassword()
		configs[i] = subCfg
	}

	return configs, nil
}

// ExecSQLOnDMSource executes SQL on a data source defined by ConnProfile
func (s *AgentServer) ExecSQLOnDMSource(ctx context.Context, req *agent.ExecSQLOnDMSourceRequest) (
	*agent.ExecSQLOnDMSourceResponse, error,
) {
	logger := pkglog.GetLogger(ctx).With(zap.String("sql", req.GetSql()))
	resp := &agent.ExecSQLOnDMSourceResponse{
		Rows: make([]*agent.Row, 0),
	}

	profile := req.GetConnProfile()
	if profile == nil {
		return nil, errors.New("connection profile is required")
	}

	// 1. Decrypt database password
	passwd, err := s.getDMSourcePassword(ctx, profile)
	if err != nil {
		if errors.ErrorEqual(err, crypto.PEMUpdatedError) {
			resp.PemChanged = true
			return resp, nil
		}
		return nil, errors.Trace(err)
	}

	// 2. Setup TLS configuration if Security info is provided
	tlsConfigName := "false" // Default to no TLS
	if profile.Security != nil && len(profile.Security.SslCaContent) > 0 {
		// Use a unique name to avoid global collision and pollution
		tlsConfigName = fmt.Sprintf("tls_dm_%s", uuid.NewString())
		if err := s.registerTLS(tlsConfigName, profile.Security); err != nil {
			return nil, errors.Annotate(err, "register TLS failed")
		}
		// Ensure the global registry is cleaned up after connection is established
		defer mysql.DeregisterTLSConfig(tlsConfigName)
	}

	// 3. Resolve Host/Endpoint
	host := profile.GetHost()
	//if host == "" && profile.GetConnType() == agent.ConnType_PRIVATE_LINK {
	//	host = profile.GetEndpointId()
	//}

	// 4. Construct MySQL DSN
	// parseTime=true is essential for scanning time.Time correctly
	dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/?tls=%s&parseTime=true&timeout=10s&interpolateParams=true",
		profile.User, passwd, host, profile.Port, tlsConfigName)

	db, err := sql.Open("mysql", dsn)
	if err != nil {
		return nil, errors.Trace(err)
	}
	defer db.Close()

	// 5. Execute Query
	rows, err := db.QueryContext(ctx, req.GetSql())
	if err != nil {
		logger.Error("execute query failed", zap.Error(err))
		if msg := sourceDBPayloadErrorMessage(err); msg != "" {
			resp.ErrorMsg = msg
			return resp, nil
		}
		return nil, err
	}
	defer rows.Close()

	// 6. Process Metadata and Rows
	resp.ColumnNames, _ = rows.Columns()
	colCount := len(resp.ColumnNames)

	for rows.Next() {
		row := &agent.Row{}
		values := make([]interface{}, colCount)
		valuePtrs := make([]interface{}, colCount)
		for i := range values {
			valuePtrs[i] = &values[i]
		}

		if err := rows.Scan(valuePtrs...); err != nil {
			return nil, errors.Trace(err)
		}

		for _, val := range values {
			row.TypedColumnValues = append(row.TypedColumnValues, s.convertToColumnValue(val))
		}
		resp.Rows = append(resp.Rows, row)
	}

	return resp, nil
}

// getDMSourcePassword logic adapted from ExecSQLOnCluster
func (s *AgentServer) getDMSourcePassword(ctx context.Context, profile *agent.ConnProfile) (string, error) {
	enc := profile.GetEncryptedPassword()
	if enc == nil {
		return "", errors.New("encrypted password is empty")
	}

	verified, err := s.encryptor.VerifyKeyMatch(enc.GetPem())
	if err != nil {
		return "", errors.Trace(err)
	}
	if !verified {
		return "", crypto.PEMUpdatedError
	}

	passwd, err := s.encryptor.Decrypt(enc.GetData())
	if err != nil {
		return "", errors.Trace(err)
	}
	return string(passwd), nil
}

// registerTLS handles the dynamic registration of certificates for MySQL driver
func (s *AgentServer) registerTLS(name string, sec *agent.Security) error {
	rootCertPool := x509.NewCertPool()
	if ok := rootCertPool.AppendCertsFromPEM(sec.SslCaContent); !ok {
		return errors.New("failed to parse root CA")
	}

	clientCerts := make([]tls.Certificate, 0)
	if len(sec.SslCertContent) > 0 && len(sec.SslKeyContent) > 0 {
		cert, err := tls.X509KeyPair(sec.SslCertContent, sec.SslKeyContent)
		if err != nil {
			return errors.Trace(err)
		}
		clientCerts = append(clientCerts, cert)
	}

	mysql.RegisterTLSConfig(name, &tls.Config{
		RootCAs:      rootCertPool,
		Certificates: clientCerts,
	})
	return nil
}

// convertToColumnValue maps database types to Protobuf-friendly types
func (s *AgentServer) convertToColumnValue(value interface{}) *agent.ColumnValue {
	cv := &agent.ColumnValue{}
	if value == nil {
		return cv
	}

	switch v := value.(type) {
	case int64:
		cv.Value = &agent.ColumnValue_Int64Value{Int64Value: v}
	case float64:
		cv.Value = &agent.ColumnValue_DoubleValue{DoubleValue: v}
	case bool:
		cv.Value = &agent.ColumnValue_BoolValue{BoolValue: v}
	case string:
		cv.Value = &agent.ColumnValue_StringValue{StringValue: v}
	case []uint8:
		cv.Value = &agent.ColumnValue_StringValue{StringValue: string(v)}
	case time.Time:
		cv.Value = &agent.ColumnValue_TimestampValue{TimestampValue: timestamppb.New(v)}
	default:
		// Fallback for types like decimal.Decimal or uint64
		cv.Value = &agent.ColumnValue_StringValue{StringValue: fmt.Sprint(v)}
	}
	return cv
}