package agent

import (
	"bytes"
	"context"
	"crypto/tls"
	"encoding/json"
	stderrors "errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/avast/retry-go/v4"
	"github.com/gin-gonic/gin/binding"
	"github.com/pingcap/errors"
	"github.com/pingcap/log"
	"github.com/tidbcloud/dataflow-service/app/container"
	agentpb "github.com/tidbcloud/dataflow-service/dataflow-agent-ng/idl/pbgen/proto/ng_agent"
	service_config "github.com/tidbcloud/dataflow-service/internal/config"
	"github.com/tidbcloud/dataflow-service/internal/crypto"
	pkglog "github.com/tidbcloud/dataflow-service/pkg/log"
	mgmt "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
	"go.uber.org/zap"
	"google.golang.org/grpc/codes"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/runtime/protoimpl"
	"google.golang.org/protobuf/types/known/emptypb"
)

type ErrorResponse struct {
	Code    int64       `json:"code"`
	Message string      `json:"message"`
	Details interface{} `json:"details"`
}

func (err *ErrorResponse) Error() string {
	return err.Message
}

var errRetryableUpstream = stderrors.New("retryable upstream error")

type retryableUpstreamError struct {
	statusCode int
	response   string
	err        error
}

func (e *retryableUpstreamError) Error() string {
	if e == nil {
		return ""
	}
	if e.err != nil {
		return e.err.Error()
	}
	return "retryable upstream error"
}

func (e *retryableUpstreamError) Unwrap() error {
	if e == nil {
		return nil
	}
	return e.err
}

func (e *retryableUpstreamError) Is(target error) bool {
	return target == errRetryableUpstream
}

func isRetryableUpstreamError(err error) bool {
	if err == nil {
		return false
	}
	if stderrors.Is(err, errRetryableUpstream) {
		return true
	}

	var urlErr *url.Error
	if stderrors.As(err, &urlErr) {
		if urlErr.Timeout() {
			return true
		}
		err = urlErr.Err
	}

	var netErr net.Error
	if stderrors.As(err, &netErr) && (netErr.Timeout()) {
		return true
	}

	errMsg := strings.ToLower(err.Error())
	return strings.Contains(errMsg, "connection reset") ||
		strings.Contains(errMsg, "connection refused") ||
		strings.Contains(errMsg, "server closed idle connection") ||
		strings.Contains(errMsg, "unexpected eof") ||
		strings.Contains(errMsg, "eof")
}

type NGRegionalClient struct {
	BaseUrl  string
	AgentUrl string
}

type TiKVCapacityReservationStatus string

const (
	tikvCapacityReservationBizTypeImport = "TIKV_CAPACITY_RESERVATION_BIZ_TYPE_IMPORT"

	TiKVCapacityReservationStatusProvisioning TiKVCapacityReservationStatus = "TIKV_CAPACITY_RESERVATION_STATUS_PROVISIONING"
	TiKVCapacityReservationStatusReady        TiKVCapacityReservationStatus = "TIKV_CAPACITY_RESERVATION_STATUS_READY"
	TiKVCapacityReservationStatusReleased     TiKVCapacityReservationStatus = "TIKV_CAPACITY_RESERVATION_STATUS_RELEASED"
)

type TiKVSpecUsage struct {
	NodeCount           uint32 `json:"nodeCount,omitempty"`
	TotalCapacityBytes  uint64 `json:"totalCapacityBytes,string,omitempty"`
	TotalAvailableBytes uint64 `json:"totalAvailableBytes,string,omitempty"`
	UsagePercentage     uint32 `json:"usagePercentage,omitempty"`
}

type TiKVCapacityReservation struct {
	SharedPoolID  string                        `json:"sharedPoolId,omitempty"`
	BizType       string                        `json:"bizType,omitempty"`
	ReservationID string                        `json:"reservationId,omitempty"`
	NeedScale     bool                          `json:"needScale,omitempty"`
	Status        TiKVCapacityReservationStatus `json:"status,omitempty"`
	From          *TiKVSpecUsage                `json:"from,omitempty"`
	To            *TiKVSpecUsage                `json:"to,omitempty"`
}

// SharedPool is the main resource for a tidb ng cluster
type SharedPool struct {
	// the shared pool id
	Id string `json:"id,omitempty"`
	// the shared pool component groups
	ComponentGroups []*Component `json:"componentGroups,omitempty"`
	// the shared pool version
	Version string `json:"version,omitempty"`
	// the object storage is used to store the cloud storage engine.
	ObjectStorage *SharedPool_ObjectStorage `json:"objectStorage,omitempty"`
	// the internal object storage is used to tidb global sort and sql audit log.
	InternalStorage *SharedPool_ObjectStorage `json:"internalStorage,omitempty"`
	DataplaneName   string                    `json:"dataplaneName,omitempty"`
	// the SharedPool status
	Phase string `json:"phase,omitempty"`
	// the cluster gateway base url,
	// internal service used to connect to the pd/tikv/tiflash service.
	GatewayBaseUrl string `json:"gatewayBaseUrl,omitempty"`
}

// the object storage configuration.
type SharedPool_ObjectStorage struct {
	Region      string `json:"region,omitempty"`
	Bucket      string `json:"bucket,omitempty"`
	Prefix      string `json:"prefix,omitempty"`
	AccountName string `json:"accountName,omitempty"`
	AccountKey  string `json:"accountKey,omitempty"`
}

// shared pool component
type Component struct {
	// id means the component group id
	Id string `json:"id,omitempty"`
	// the component type
	ComponentType string `json:"componentType,omitempty"`
	// the component node quota
	NodeQuota *NodeQuota `json:"nodeQuota,omitempty"`
	// include storage and tikv additional storage
	Storages []*Component_Storage `json:"storages,omitempty"`
	// the component configuration
	// include the image repo and image tag
	Configuration *Component_Configuration `json:"configuration,omitempty"`
	// base version of the component
	// tiproxy version is not same as the shared pool version
	Version string `json:"version,omitempty"`
	Phase   string `json:"phase,omitempty"`
}

// the storage configuration
type Component_Storage struct {
	// the storage name
	Name string `json:"name,omitempty"`
	// the storage type
	// aws s3/io2 or gcp gcs
	// or azure blob
	Type string `json:"type,omitempty"`
	// the storage size
	SizeGi uint32 `json:"sizeGi,omitempty"`
	// the storage iops
	Iops uint32 `json:"iops,omitempty"`
	// the storage throughput
	ThroughputMbs uint32 `json:"throughputMbs,omitempty"`
	// the storage mount path
	MountPath string `json:"mountPath,omitempty"`
}

// the component configuration
type Component_Configuration struct {
	ConfigurationFileContent string `json:"configurationFileContent,omitempty"`
	ImageRepo                string `json:"imageRepo,omitempty"`
	ImageTag                 string `json:"imageTag,omitempty"`
}

type TiDBGroup struct {
	// This is the unique id for the tidb group.
	Id string `json:"id,omitempty"`
	// user cluster id
	ClusterId string `json:"clusterId,omitempty"`
	// user level id
	TenantId string `json:"tenantId,omitempty"`
	// shared pool id
	SharedPoolId string `json:"sharedPoolId,omitempty"`
	// the component node quota
	NodeQuota *NodeQuota `json:"nodeQuota,omitempty"`
	// cluster gateway base url,
	// internal service used to connect to the tidb group.
	GatewayBaseUrl string `json:"gatewayBaseUrl,omitempty"`
	// only display the tidb group status,
	// not included the pd/tikv/tiflash status.
	Phase string `json:"phase,omitempty"`
	// the tidb group version
	Version       string `json:"version,omitempty"`
	DataplaneName string `json:"dataplaneName,omitempty"`
}

// the node quota is used to describe the node spec
type NodeQuota struct {

	// e.g. 8C16G
	NodeSpec *NodeSpec `json:"nodeSpec,omitempty"`
	// node count
	NodeCount uint32 `json:"nodeCount,omitempty"`
	// generate by request unit
	// the max/min node count is used to limit the node count
	MaxNodeCount uint32 `json:"maxNodeCount,omitempty"`
	MinNodeCount uint32 `json:"minNodeCount,omitempty"`
}

// the node spec is used to describe the node spec
type NodeSpec struct {

	// the node spec is used to describe the node spec
	Cpu      uint32 `json:"cpu,omitempty"`
	MemoryGi uint32 `json:"memoryGi,omitempty"`
	// the cloud provider will reserve some resources for the node
	// so available_cpu and available_memory_mi is not equal to the node spec
	// not returned in the api
	AvailableCpu      uint32 `json:"availableCpu,omitempty"`
	AvailableMemoryMi uint32 `json:"availableMemoryMi,omitempty"`
	WithLocalDisk     bool   `json:"withLocalDisk,omitempty"`
	LocalDiskSizeGi   uint32 `json:"localDiskSizeGi,omitempty"`
}

type ScaleSharedPoolRequest struct {
	SharedPoolId string                                 `json:"sharedPoolId,omitempty"`
	Component    *ScaleSharedPoolRequest_ScaleComponent `json:"component,omitempty"`
	Force        bool                                   `json:"force,omitempty"`
}

type ScaleSharedPoolRequest_ScaleComponent struct {
	state         protoimpl.MessageState
	sizeCache     protoimpl.SizeCache
	unknownFields protoimpl.UnknownFields

	// id means the component group id
	Id string `json:"id,omitempty"`
	// the component type
	ComponentType string `json:"componentType,omitempty"`
	// the component node quota
	NodeQuota *NodeQuota `json:"nodeQuota,omitempty"`
	// include storage and tikv additional storage
	Storages []*Component_Storage `json:"storages,omitempty"`
}

// describe the dm cluster
type DMCluster struct {
	Id              string     `json:"id,omitempty"`
	SharedPoolId    string     `json:"sharedPoolId,omitempty"`
	MasterNodeQuota *NodeQuota `json:"masterNodeQuota,omitempty"`
	WorkerNodeQuota *NodeQuota `json:"workerNodeQuota,omitempty"`
	Phase           string     `json:"phase,omitempty"`
	Backend         string     `json:"backend,omitempty"`
	DataplaneName   string     `json:"dataplaneName,omitempty"`
}

type CreateDMClusterRequest struct {
	DmCluster *DMCluster `json:"dmCluster,omitempty"`
}

type ScaleDMClusterRequest struct {
	SharedPoolId    string     `json:"sharedPoolId,omitempty"`
	DmClusterId     string     `json:"dmClusterId,omitempty"`
	MasterNodeQuota *NodeQuota `json:"masterNodeQuota,omitempty"`
	WorkerNodeQuota *NodeQuota `json:"workerNodeQuota,omitempty"`
	Force           bool       `json:"force,omitempty"`
}

var PEMChangedError = errors.New("pem changed")

func NewNGRegionalClientByClusterID(ctx context.Context, clusterID uint64) (*NGRegionalClient, error) {
	ctx, cancelFunc := context.WithTimeout(ctx, time.Second*10)
	defer cancelFunc()
	listResp, err := container.Global().MgmtSvcGrpcClient.ListClusters(ctx, &mgmt.ListClustersRequest{Filter: &mgmt.ClusterFilter{ClusterIds: []uint64{clusterID}}})
	if err != nil {
		log.Error("list mgmt clusters error", zap.Error(err), zap.Uint64("clusterID", clusterID))
		return nil, errors.Annotate(err, "list clusters by ID error")
	}
	if len(listResp.Clusters) == 0 {
		return nil, errors.Annotate(err, "db not exist")
	}

	if listResp.Clusters[0].GetNextGenCluster() == nil || listResp.Clusters[0].GetNextGenCluster().GetSharedPoolId() == "" {
		return nil, errors.New("cluster is not a nextgen cluster or shared pool id is empty")
	}

	sharedPoolId, perr := strconv.ParseUint(listResp.Clusters[0].GetNextGenCluster().GetSharedPoolId(), 10, 64)
	if perr != nil {
		return nil, errors.Annotate(perr, "parse shared pool id error")
	}

	return NewNGRegionalClientByResourcePoolID(ctx, sharedPoolId)
}

func NewNGRegionalClientByResourcePoolID(ctx context.Context, resourcePoolID uint64) (*NGRegionalClient, error) {
	ctx, cancelFunc := context.WithTimeout(ctx, time.Second*10)
	defer cancelFunc()

	// check resource pool existence
	_, err := container.Global().MgmtSvcGrpcClient.GetResourcePoolByID(ctx, resourcePoolID)
	if err != nil {
		log.Error("get mgmt resource pool error", zap.Error(err), zap.Uint64("resourcePoolID", resourcePoolID))
		return nil, errors.Annotate(err, "get resource pool by ID error")
	}

	baseURL := fmt.Sprintf("%s/proxy/nextgen/resourcePools/%d", service_config.Global().GrpcConfig.MgmtSvcProxyAddr, resourcePoolID)
	client := &NGRegionalClient{
		BaseUrl: baseURL,
	}

	client.AgentUrl = baseURL + fmt.Sprintf("/dataplane/proxy/sharedPools/%d/dataflow-agent", resourcePoolID)
	return client, nil
}

func (c *NGRegionalClient) GetTiDBGroup(ctx context.Context, clusterID uint64) (*TiDBGroup, error) {
	url := fmt.Sprintf("%s/api/premium/v1/clusters/%d/tidbGroups?clusterId=%d", c.BaseUrl, clusterID, clusterID)
	respBytes, err := getResponseWithRetry(ctx, "GET", url, nil)
	if err != nil {
		log.Error("failed to get tidb group", zap.String("resp", string(respBytes)), zap.Error(err))
		return nil, err
	}
	var groupsResp struct {
		TidbGroups []*TiDBGroup `json:"tidbGroups"`
		TotalCount uint64       `json:"totalCount"`
	}
	if err := json.Unmarshal(respBytes, &groupsResp); err != nil {
		log.Error("unmarshal tidb group failed", zap.String("resp", string(respBytes)), zap.Error(err))
		return nil, errors.Trace(err)
	}
	if groupsResp.TotalCount == 0 {
		log.Error("get tidb group failed", zap.String("resp", string(respBytes)), zap.Error(err))
		return nil, errors.Trace(err)
	}
	log.Info("Get TiDB Group", zap.Uint64("clusterID", clusterID), zap.Any("group", groupsResp.TidbGroups[0]))
	return groupsResp.TidbGroups[0], nil
}

func (c *NGRegionalClient) ScaleSharedPool(ctx context.Context, sharedPoolId uint64, req *ScaleSharedPoolRequest) error {
	url := fmt.Sprintf("%s/api/premium/v1/sharedPools/%d:scale", c.BaseUrl, sharedPoolId)
	payload, err := json.Marshal(req)
	if err != nil {
		return errors.Trace(err)
	}
	respBytes, err := getResponseWithRetry(ctx, "POST", url, payload)
	if err != nil {
		log.Error("failed to scale tidb group dxf", zap.Uint64("sharedPoolId", sharedPoolId), zap.String("payload", string(payload)), zap.String("respBytes", string(respBytes)), zap.Error(err))
		return err
	}
	log.Info("scale tidb group dxf", zap.Uint64("sharedPoolId", sharedPoolId), zap.String("payload", string(payload)), zap.String("respBytes", string(respBytes)))
	return nil
}

func (c *NGRegionalClient) GetSharePool(ctx context.Context, sharedPoolId uint64) (*SharedPool, error) {
	url := fmt.Sprintf("%s/api/premium/v1/sharedPools/%d", c.BaseUrl, sharedPoolId)
	respBytes, err := getResponseWithRetry(ctx, "GET", url, nil)
	if err != nil {
		log.Error("failed to get share pool", zap.Uint64("sharedPoolId", sharedPoolId), zap.String("respBytes", string(respBytes)), zap.Error(err))
		return nil, errors.Trace(err)
	}
	var sharedPools SharedPool
	if err := json.Unmarshal(respBytes, &sharedPools); err != nil {
		log.Error("unmarshal share pool failed", zap.Uint64("sharedPoolId", sharedPoolId), zap.String("resp", string(respBytes)), zap.Error(err))
		return nil, errors.Trace(err)
	}
	log.Info("success to get share pool", zap.Uint64("sharedPoolId", sharedPoolId), zap.String("respBytes", string(respBytes)))
	return &sharedPools, nil
}

func (c *NGRegionalClient) GetSharePoolByClusterID(ctx context.Context, clusterID uint64) (*SharedPool, error) {
	tidbGroup, err := c.GetTiDBGroup(ctx, clusterID)
	if err != nil {
		log.Error("failed to get tidb group", zap.Uint64("clusterID", clusterID), zap.Error(err))
		return nil, errors.Trace(err)
	}

	sharedPoolId, err := strconv.ParseUint(tidbGroup.SharedPoolId, 10, 64)
	if err != nil {
		log.Error("failed to parse shared pool id", zap.Uint64("clusterID", clusterID), zap.String("sharedPoolId", tidbGroup.SharedPoolId), zap.Error(err))
		return nil, errors.Trace(err)
	}

	sharedPool, err := c.GetSharePool(ctx, sharedPoolId)
	if err != nil {
		log.Error("failed to get shared pool", zap.Uint64("clusterID", clusterID), zap.Uint64("sharedPoolId", sharedPoolId), zap.Error(err))
		return nil, errors.Trace(err)
	}
	return sharedPool, nil
}

func (c *NGRegionalClient) ExecuteSQL(ctx context.Context, req *agentpb.ExecSQLReq) (*agentpb.ExecSQLResp, error) {
	url := fmt.Sprintf("%s/api/v1/cluster/execsql", c.AgentUrl)
	respBytes, err := getResponseProtoJson(ctx, http.MethodPost, url, req)
	if err != nil {
		log.Error("execute sql failed", zap.Error(err))
		return nil, err
	}
	var resp agentpb.ExecSQLResp
	// use protojson to unmarshal the JSON, so that the enum value could be correctly parsed
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func postProtoRequest(ctx context.Context, url string, req proto.Message, resp proto.Message) error {
	respBytes, err := getResponseProtoJson(ctx, http.MethodPost, url, req)
	if err != nil {
		return err
	}
	if err := protojson.Unmarshal(respBytes, resp); err != nil {
		return errors.Trace(err)
	}
	return nil
}

func (c *NGRegionalClient) postProtoRequestWithCluster(
	ctx context.Context,
	url string,
	clusterID uint64,
	username string,
	password string,
	buildReq func(cluster *agentpb.Cluster) proto.Message,
	resp proto.Message,
) error {
	cluster, err := c.buildEncryptedCluster(ctx, clusterID, username, password)
	if err != nil {
		return err
	}
	return postProtoRequest(ctx, url, buildReq(cluster), resp)
}

func (c *NGRegionalClient) CreateImportIntoJobTask(
	ctx context.Context,
	taskID string,
	clusterID uint64,
	username string,
	password string,
	sql string,
	setupSQLs []string,
) (*agentpb.CreateImportIntoJobTaskResp, error) {
	url := fmt.Sprintf("%s/api/v1/import/import-into-job-tasks", c.AgentUrl)
	var resp agentpb.CreateImportIntoJobTaskResp
	if err := c.postProtoRequestWithCluster(ctx, url, clusterID, username, password, func(cluster *agentpb.Cluster) proto.Message {
		return &agentpb.CreateImportIntoJobTaskReq{
			TaskId:    taskID,
			ClusterId: clusterID,
			Cluster:   cluster,
			Sql:       sql,
			SetupSqls: setupSQLs,
		}
	}, &resp); err != nil {
		return nil, err
	}
	return &resp, nil
}

func (c *NGRegionalClient) BatchCreateImportIntoJobTasks(
	ctx context.Context,
	clusterID uint64,
	username string,
	password string,
	tasks []*agentpb.BatchCreateImportIntoJobTasksReq_Task,
) (*agentpb.BatchCreateImportIntoJobTasksResp, error) {
	url := fmt.Sprintf("%s/api/v1/import/import-into-job-tasks:batchCreate", c.AgentUrl)
	var resp agentpb.BatchCreateImportIntoJobTasksResp
	if err := c.postProtoRequestWithCluster(ctx, url, clusterID, username, password, func(cluster *agentpb.Cluster) proto.Message {
		return &agentpb.BatchCreateImportIntoJobTasksReq{
			ClusterId: clusterID,
			Cluster:   cluster,
			Tasks:     tasks,
		}
	}, &resp); err != nil {
		return nil, err
	}
	return &resp, nil
}

func (c *NGRegionalClient) GetImportIntoJobTask(ctx context.Context, taskID string) (*agentpb.GetImportIntoJobTaskResp, error) {
	url := fmt.Sprintf("%s/api/v1/import/import-into-job-tasks/%s", c.AgentUrl, taskID)
	respBytes, err := getResponseWithRetry(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	var resp agentpb.GetImportIntoJobTaskResp
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (c *NGRegionalClient) BatchGetImportIntoJobTasks(ctx context.Context, taskIDs []string) (*agentpb.BatchGetImportIntoJobTasksResp, error) {
	req := &agentpb.BatchGetImportIntoJobTasksReq{TaskIds: taskIDs}
	url := fmt.Sprintf("%s/api/v1/import/import-into-job-tasks:batchGet", c.AgentUrl)
	var resp agentpb.BatchGetImportIntoJobTasksResp
	if err := postProtoRequest(ctx, url, req, &resp); err != nil {
		return nil, err
	}
	return &resp, nil
}

func (c *NGRegionalClient) BatchCancelImportIntoJobs(
	ctx context.Context,
	clusterID uint64,
	username string,
	password string,
	jobIDs []uint64,
) (*agentpb.BatchCancelImportIntoJobsResp, error) {
	url := fmt.Sprintf("%s/api/v1/import/import-into-jobs:batchCancel", c.AgentUrl)
	var resp agentpb.BatchCancelImportIntoJobsResp
	if err := c.postProtoRequestWithCluster(ctx, url, clusterID, username, password, func(cluster *agentpb.Cluster) proto.Message {
		return &agentpb.BatchCancelImportIntoJobsReq{
			ClusterId: clusterID,
			Cluster:   cluster,
			JobIds:    jobIDs,
		}
	}, &resp); err != nil {
		return nil, err
	}
	return &resp, nil
}

func (c *NGRegionalClient) GetImportFiles(ctx context.Context, clusterId uint64, sourceUri string, maxScanFiles int32, tableMetas ...*agentpb.GetImportFilesReq_TableMeta) (*agentpb.GetImportFilesResp, error) {
	req := &agentpb.GetImportFilesReq{
		SourceUri: sourceUri,
		ClusterId: clusterId,
		TableMeta: tableMetas,
	}
	if maxScanFiles > 0 {
		req.MaxScanFiles = maxScanFiles
	}
	url := fmt.Sprintf("%s/api/v1/import/get-import-files", c.AgentUrl)

	respBytes, err := getResponseProtoJson(ctx, http.MethodPost, url, req)
	if err != nil {
		return nil, err
	}
	var resp agentpb.GetImportFilesResp
	// use protojson to unmarshal the JSON, so that the enum value could be correctly parsed
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (c *NGRegionalClient) CreateImportFilesTask(
	ctx context.Context,
	taskID string,
	clusterID uint64,
	sourceURI string,
	maxScanFiles int32,
	tableMetas ...*agentpb.GetImportFilesReq_TableMeta,
) (*agentpb.CreateImportFilesTaskResp, error) {
	req := &agentpb.CreateImportFilesTaskReq{
		TaskId:       taskID,
		ClusterId:    clusterID,
		SourceUri:    sourceURI,
		MaxScanFiles: maxScanFiles,
		TableMeta:    tableMetas,
	}
	url := fmt.Sprintf("%s/api/v1/import/files-tasks", c.AgentUrl)
	respBytes, err := getResponseProtoJson(ctx, http.MethodPost, url, req)
	if err != nil {
		return nil, err
	}
	var resp agentpb.CreateImportFilesTaskResp
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (c *NGRegionalClient) GetImportFilesTask(ctx context.Context, taskID string) (*agentpb.GetImportFilesTaskResp, error) {
	url := fmt.Sprintf("%s/api/v1/import/files-tasks/%s", c.AgentUrl, taskID)
	respBytes, err := getResponseWithRetry(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	var resp agentpb.GetImportFilesTaskResp
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (c *NGRegionalClient) CreateSchemasTask(
	ctx context.Context,
	taskID string,
	clusterID uint64,
	sourceURI string,
	username string,
	password string,
	schemas ...*agentpb.CreateSchemasReq_Schema,
) (*agentpb.CreateSchemasTaskResp, error) {
	cluster, err := c.buildEncryptedCluster(ctx, clusterID, username, password)
	if err != nil {
		return nil, err
	}
	req := &agentpb.CreateSchemasTaskReq{
		TaskId:    taskID,
		ClusterId: clusterID,
		SourceUri: sourceURI,
		Schemas:   schemas,
		Cluster:   cluster,
	}
	url := fmt.Sprintf("%s/api/v1/import/create-schemas-tasks", c.AgentUrl)
	respBytes, err := getResponseProtoJson(ctx, http.MethodPost, url, req)
	if err != nil {
		return nil, err
	}
	var resp agentpb.CreateSchemasTaskResp
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (c *NGRegionalClient) GetSchemasTask(ctx context.Context, taskID string) (*agentpb.GetSchemasTaskResp, error) {
	url := fmt.Sprintf("%s/api/v1/import/create-schemas-tasks/%s", c.AgentUrl, taskID)
	respBytes, err := getResponseWithRetry(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	var resp agentpb.GetSchemasTaskResp
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (c *NGRegionalClient) CreateImportDataSizeEstimateTask(
	ctx context.Context,
	taskID string,
	clusterID uint64,
	sourceURI string,
	maxScanFiles int32,
	csvConfig *agentpb.CsvConfig,
	tableMetas ...*agentpb.CreateImportDataSizeEstimateTaskReq_TableMeta,
) (*agentpb.CreateImportDataSizeEstimateTaskResp, error) {
	req := &agentpb.CreateImportDataSizeEstimateTaskReq{
		TaskId:       taskID,
		ClusterId:    clusterID,
		SourceUri:    sourceURI,
		MaxScanFiles: maxScanFiles,
		TableMeta:    tableMetas,
		CsvConfig:    csvConfig,
	}
	url := fmt.Sprintf("%s/api/v1/import/estimate-data-size-tasks", c.AgentUrl)
	respBytes, err := getResponseProtoJson(ctx, http.MethodPost, url, req)
	if err != nil {
		return nil, err
	}
	var resp agentpb.CreateImportDataSizeEstimateTaskResp
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (c *NGRegionalClient) GetImportDataSizeEstimateTask(ctx context.Context, taskID string) (*agentpb.GetImportDataSizeEstimateTaskResp, error) {
	url := fmt.Sprintf("%s/api/v1/import/estimate-data-size-tasks/%s", c.AgentUrl, taskID)
	respBytes, err := getResponseWithRetry(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	var resp agentpb.GetImportDataSizeEstimateTaskResp
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (c *NGRegionalClient) CreateSchemas(ctx context.Context, clusterId uint64, sourceUri string, username string, password string, schemas ...*agentpb.CreateSchemasReq_Schema) (*emptypb.Empty, error) {
	err := retry.Do(func() error {
		cluster, err := c.buildEncryptedCluster(ctx, clusterId, username, password)
		if err != nil {
			return err
		}
		req := &agentpb.CreateSchemasReq{
			User:      username,
			ClusterId: clusterId,
			SourceUri: sourceUri,
			Schemas:   schemas,
			Cluster:   cluster,
		}
		url := fmt.Sprintf("%s/api/v1/import/create-schemas", c.AgentUrl)
		resp, err := getResponseProtoJson(ctx, http.MethodPost, url, req)
		if err != nil {
			return err
		}
		parsedResp := &agentpb.CreateSchemasResp{}
		if err := protojson.Unmarshal(resp, parsedResp); err != nil {
			return errors.Trace(err)
		}
		if err := createSchemasErrorFromResp(parsedResp); err != nil {
			return err
		}
		return nil
	}, retry.RetryIf(func(err error) bool {
		return errors.ErrorEqual(err, PEMChangedError)
	}), retry.Attempts(service_config.Global().ImportConfig.MaxRetryTimesOnExecImportingQuerySQL), retry.LastErrorOnly(true))
	return &emptypb.Empty{}, err
}

func (c *NGRegionalClient) buildEncryptedCluster(ctx context.Context, clusterID uint64, username, password string) (*agentpb.Cluster, error) {
	pemResp, err := c.GetPublicKey(ctx, &agentpb.GetPublicKeyRequest{})
	if err != nil {
		return nil, err
	}
	data, err := crypto.EncryptWithPEM(pemResp.GetPem(), []byte(password))
	if err != nil {
		return nil, err
	}
	return &agentpb.Cluster{
		ClusterId:         clusterID,
		User:              username,
		PasswordEncrypted: true,
		EncryptedPassword: &agentpb.EncryptedData{
			Data: data,
			Pem:  pemResp.GetPem(),
		},
	}, nil
}

func createSchemasErrorFromResp(resp *agentpb.CreateSchemasResp) error {
	if resp == nil {
		return nil
	}
	if resp.GetPemChanged() {
		return PEMChangedError
	}
	if resp.GetErrorMsg() != "" {
		return &ErrorResponse{Message: resp.GetErrorMsg()}
	}
	return nil
}

func (c *NGRegionalClient) CreateTiKVCapacityReservation(
	ctx context.Context,
	sharedPoolID, reservationID string,
	requiredDataBytes uint64,
) (*TiKVCapacityReservation, error) {
	payload, err := json.Marshal(map[string]any{
		"sharedPoolId":      sharedPoolID,
		"bizType":           tikvCapacityReservationBizTypeImport,
		"reservationId":     reservationID,
		"requiredDataBytes": requiredDataBytes,
	})
	if err != nil {
		return nil, errors.Trace(err)
	}

	url := fmt.Sprintf("%s/api/premium/v1/sharedPools/%s/tikvCapacityReservations", c.BaseUrl, sharedPoolID)
	respBytes, err := getResponseWithRetry(ctx, http.MethodPost, url, payload)
	if err != nil {
		return nil, err
	}
	return parseTiKVCapacityReservation(respBytes)
}

func (c *NGRegionalClient) GetTiKVCapacityReservation(ctx context.Context, sharedPoolID, reservationID string) (*TiKVCapacityReservation, error) {
	url := fmt.Sprintf("%s/api/premium/v1/sharedPools/%s/tikvCapacityReservations/%s/%s", c.BaseUrl, sharedPoolID, tikvCapacityReservationBizTypeImport, reservationID)
	respBytes, err := getResponseWithRetry(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	return parseTiKVCapacityReservation(respBytes)
}

func (c *NGRegionalClient) DeleteTiKVCapacityReservation(ctx context.Context, sharedPoolID, reservationID string) error {
	url := fmt.Sprintf("%s/api/premium/v1/sharedPools/%s/tikvCapacityReservations/%s/%s", c.BaseUrl, sharedPoolID, tikvCapacityReservationBizTypeImport, reservationID)
	_, err := getResponseWithRetry(ctx, http.MethodDelete, url, nil)
	return err
}

func parseTiKVCapacityReservation(respBytes []byte) (*TiKVCapacityReservation, error) {
	var resp TiKVCapacityReservation
	if err := json.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func getResponseProtoJsonWithRetry(ctx context.Context, method, url string, req proto.Message) ([]byte, error) {
	payload, err := protojson.Marshal(req)
	if err != nil {
		return nil, err
	}
	return getResponseWithRetry(ctx, method, url, payload)
}

func getResponseProtoJson(ctx context.Context, method, url string, req proto.Message) ([]byte, error) {
	payload, err := protojson.Marshal(req)
	if err != nil {
		return nil, err
	}
	return getResponse(ctx, method, url, payload)
}

const (
	ngRegionalRetryAttempts = 3
	ngRegionalRetryDelay    = 200 * time.Millisecond
)

var sharedHttpClient = &http.Client{
	Transport: otelhttp.NewTransport(&http.Transport{
		TLSClientConfig:     &tls.Config{InsecureSkipVerify: true},
		MaxIdleConns:        500,
		IdleConnTimeout:     90 * time.Second,
		MaxIdleConnsPerHost: 100,
	}),
}

func getResponseWithRetry(ctx context.Context, method, url string, payload []byte) ([]byte, error) {
	var respBytes []byte
	err := retry.Do(
		func() error {
			var err error
			respBytes, err = getResponse(ctx, method, url, payload)
			return err
		},
		retry.Attempts(ngRegionalRetryAttempts),
		retry.Delay(ngRegionalRetryDelay),
		retry.DelayType(retry.FixedDelay),
		retry.Context(ctx),
		retry.LastErrorOnly(true),
		retry.OnRetry(func(n uint, err error) {
			pkglog.GetLogger(ctx).Warn("retry ng regional request",
				zap.Uint("attempt", n+1),
				zap.String("method", method),
				zap.String("url", url),
				zap.Error(err),
			)
		}),
		retry.RetryIf(func(err error) bool {
			return isRetryableUpstreamError(err)
		}),
	)
	if err != nil {
		return nil, err
	}
	return respBytes, nil
}

func getResponse(ctx context.Context, method, url string, payload []byte) ([]byte, error) {
	logger := pkglog.GetLogger(ctx).With(zap.String("method", method), zap.String("url", url))
	httpReq, err := http.NewRequestWithContext(ctx, method, url, bytes.NewReader(payload))
	if err != nil {
		return nil, errors.Trace(err)
	}
	httpReq.Header.Set("Content-Type", binding.MIMEJSON)
	httpReq.Header.Set("Authorization", "Bearer "+service_config.Global().GrpcConfig.MgmtSvcProxyToken)

	httpClient := sharedHttpClient
	httpResp, err := httpClient.Do(httpReq)
	if err != nil {
		traceErr := errors.Trace(err)
		if isRetryableUpstreamError(traceErr) {
			return nil, &retryableUpstreamError{err: traceErr}
		}
		return nil, traceErr
	}
	defer httpResp.Body.Close() // nolint: errcheck

	respBytes, err := io.ReadAll(httpResp.Body)
	if err != nil {
		return nil, errors.Trace(err)
	}

	if httpResp.StatusCode != http.StatusOK {
		if httpResp.StatusCode < 500 {
			var errResp = ErrorResponse{}
			if parseErr := json.Unmarshal(respBytes, &errResp); parseErr == nil {
				logger.Info("response failed but been catch", zap.String("response", string(respBytes)))
				return nil, &errResp
			}
			msg := strings.TrimSpace(string(respBytes))
			if msg == "" {
				msg = http.StatusText(httpResp.StatusCode)
			}
			logger.Info("response failed with plain text body", zap.String("response", string(respBytes)))
			return nil, &ErrorResponse{
				Code:    int64(httpStatusToGRPCCode(httpResp.StatusCode)),
				Message: msg,
			}
		} else {
			if httpResp.StatusCode == http.StatusBadGateway ||
				httpResp.StatusCode == http.StatusServiceUnavailable ||
				httpResp.StatusCode == http.StatusGatewayTimeout {
				logger.Warn("retryable upstream response failed",
					zap.Int("status_code", httpResp.StatusCode),
					zap.String("response", string(respBytes)),
				)
				return nil, &retryableUpstreamError{
					statusCode: httpResp.StatusCode,
					response:   string(respBytes),
					err:        errors.New("request failed"),
				}
			}
		}
		logger.Error("response failed", zap.String("response", string(respBytes)))
		return nil, errors.New("request failed")
	}

	return respBytes, nil
}

func httpStatusToGRPCCode(statusCode int) codes.Code {
	switch statusCode {
	case http.StatusBadRequest:
		return codes.InvalidArgument
	case http.StatusUnauthorized:
		return codes.Unauthenticated
	case http.StatusForbidden:
		return codes.PermissionDenied
	case http.StatusNotFound:
		return codes.NotFound
	case http.StatusConflict:
		return codes.AlreadyExists
	case http.StatusTooManyRequests:
		return codes.ResourceExhausted
	case http.StatusNotImplemented:
		return codes.Unimplemented
	default:
		return codes.Unknown
	}
}

func (c *NGRegionalClient) CheckS3Access(ctx context.Context, req *agentpb.CheckS3AccessReq) (*agentpb.CheckS3AccessResp, error) {
	url := fmt.Sprintf("%s/api/v1/s3/check_access", c.AgentUrl)

	respBytes, err := getResponseProtoJson(ctx, http.MethodPost, url, req)
	if err != nil {
		return nil, err
	}
	var resp agentpb.CheckS3AccessResp
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (c *NGRegionalClient) GetPublicKey(ctx context.Context, req *agentpb.GetPublicKeyRequest) (*agentpb.GetPublicKeyResponse, error) {
	url := fmt.Sprintf("%s/api/v1/crypto/pem", c.AgentUrl)

	respBytes, err := getResponseProtoJson(ctx, http.MethodPost, url, req)
	if err != nil {
		return nil, err
	}
	var resp agentpb.GetPublicKeyResponse
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (c *NGRegionalClient) CheckAliOSSAccess(ctx context.Context, req *agentpb.CheckAliOSSAccessReq) (*agentpb.CheckAliOSSAccessResp, error) {
	url := fmt.Sprintf("%s/api/v1/alioss/check_access", c.AgentUrl)

	respBytes, err := getResponseProtoJson(ctx, http.MethodPost, url, req)
	if err != nil {
		return nil, err
	}
	var resp agentpb.CheckAliOSSAccessResp
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (c *NGRegionalClient) CheckAzblobAccess(ctx context.Context, req *agentpb.CheckAzblobAccessReq) (*agentpb.CheckAzblobAccessResp, error) {
	url := fmt.Sprintf("%s/api/v1/azblob/check_access", c.AgentUrl)

	respBytes, err := getResponseProtoJson(ctx, http.MethodPost, url, req)
	if err != nil {
		return nil, err
	}
	var resp agentpb.CheckAzblobAccessResp
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (c *NGRegionalClient) ExecSQLOnDMSource(ctx context.Context, req *agentpb.ExecSQLOnDMSourceRequest) (
	*agentpb.ExecSQLOnDMSourceResponse, error,
) {
	url := fmt.Sprintf("%s/api/v1/dm/source/exec_sql", c.AgentUrl)

	respBytes, err := getResponseProtoJson(ctx, http.MethodPost, url, req)
	if err != nil {
		return nil, err
	}
	var resp agentpb.ExecSQLOnDMSourceResponse
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (c *NGRegionalClient) DMOpenAPIProxy(ctx context.Context, req *agentpb.DMOpenAPIProxyRequest) (*agentpb.DMOpenAPIProxyResponse, error) {
	url := fmt.Sprintf("%s/api/v1/dm/openapi/proxy", c.AgentUrl)

	respBytes, err := getResponseProtoJson(ctx, http.MethodPost, url, req)
	if err != nil {
		return nil, err
	}
	var resp agentpb.DMOpenAPIProxyResponse
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	log.Info("get dm response", zap.String("response", string(respBytes)))
	return &resp, nil
}

func (c *NGRegionalClient) PrecheckDMJob(ctx context.Context, req *agentpb.PrecheckDMJobReq) (*agentpb.PrecheckDMJobResp, error) {
	url := fmt.Sprintf("%s/api/v1/dm/jobprecheck", c.AgentUrl)
	respBytes, err := getResponseProtoJsonWithRetry(ctx, http.MethodPost, url, req)
	if err != nil {
		return nil, err
	}
	var resp agentpb.PrecheckDMJobResp
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

var (
	DMClusterNotFoundErr = errors.New("dm cluster not found")
)

func (c *NGRegionalClient) GetDMClusterByID(ctx context.Context, sharedPoolId, dmClusterID uint64) (*DMCluster, error) {
	url := fmt.Sprintf("%s/api/premium/v1/sharedPools/%d/dmClusters/%d", c.BaseUrl, sharedPoolId, dmClusterID)
	respBytes, err := getResponseWithRetry(ctx, "GET", url, nil)
	if err != nil {
		// check if dm cluster not found
		var errResp *ErrorResponse
		if stderrors.As(err, &errResp) && errResp.Code == 5 { // grpc not found code
			log.Info("dm cluster not found", zap.Uint64("dmClusterID", dmClusterID))
			return nil, DMClusterNotFoundErr
		}
		log.Error("failed to get dm cluster", zap.Uint64("dmClusterID", dmClusterID), zap.Error(err))
		return nil, errors.Trace(err)
	}
	var dmCluster DMCluster
	if err := json.Unmarshal(respBytes, &dmCluster); err != nil {
		log.Error("unmarshal dm cluster failed", zap.Uint64("dmClusterID", dmClusterID), zap.String("resp", string(respBytes)), zap.Error(err))
		return nil, errors.Trace(err)
	}
	log.Info("success to get dm cluster", zap.Uint64("dmClusterID", dmClusterID), zap.String("respBytes", string(respBytes)))
	return &dmCluster, nil
}

func (c *NGRegionalClient) ScaleDMCluster(ctx context.Context, sharedPoolId, dmClusterID uint64, req *ScaleDMClusterRequest) error {
	url := fmt.Sprintf("%s/api/premium/v1/sharedPools/%d/dmClusters/%d:scale", c.BaseUrl, sharedPoolId, dmClusterID)
	payload, err := json.Marshal(req)
	if err != nil {
		return errors.Trace(err)
	}
	respBytes, err := getResponseWithRetry(ctx, "POST", url, payload)
	if err != nil {
		log.Error("failed to scale dm cluster", zap.Uint64("dmClusterID", dmClusterID), zap.String("payload", string(payload)), zap.String("respBytes", string(respBytes)), zap.Error(err))
		return err
	}
	log.Info("scale dm cluster", zap.Uint64("dmClusterID", dmClusterID), zap.String("payload", string(payload)), zap.String("respBytes", string(respBytes)))
	return nil
}

func (c *NGRegionalClient) CreateDMCluster(ctx context.Context, sharedPoolId uint64, req *CreateDMClusterRequest) (*DMCluster, error) {
	url := fmt.Sprintf("%s/api/premium/v1/sharedPools/%d/dmClusters", c.BaseUrl, sharedPoolId)
	payload, err := json.Marshal(req)
	if err != nil {
		return nil, errors.Trace(err)
	}
	respBytes, err := getResponseWithRetry(ctx, "POST", url, payload)
	if err != nil {
		log.Error("failed to create dm cluster", zap.String("payload", string(payload)), zap.String("respBytes", string(respBytes)), zap.Error(err))
		return nil, err
	}
	var dmCluster DMCluster
	if err := json.Unmarshal(respBytes, &dmCluster); err != nil {
		log.Error("unmarshal dm cluster failed", zap.String("resp", string(respBytes)), zap.Error(err))
		return nil, errors.Trace(err)
	}
	log.Info("success to create dm cluster", zap.String("respBytes", string(respBytes)))
	return &dmCluster, nil
}