// Package agent add by blacknc <blacknc@163.com> at 2024/09/04 16:57:06
package agent

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"path/filepath"
	"strconv"
	"time"

	"github.com/avast/retry-go/v4"
	"github.com/pingcap/errors"
	"github.com/tidbcloud/dataflow-service/app/container"
	"github.com/tidbcloud/dataflow-service/idl/pbgen/proto/agent"
	"github.com/tidbcloud/dataflow-service/internal/config"
	"github.com/tidbcloud/dataflow-service/internal/constants"
	"github.com/tidbcloud/dataflow-service/internal/grpc/grpcli"
	infraapi "github.com/tidbcloud/dataflow-service/pkg/infra-api"
	regionalapi "github.com/tidbcloud/dataflow-service/pkg/regional-api"
	"github.com/tidbcloud/pkgv2/log"
	"go.uber.org/zap"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/types/known/emptypb"
)

var ErrUnSupportedOperation = errors.New("This operation is not currently supported")

func NewRegionalAgentClient(ctx context.Context, clusterID uint64, provider, region string) (AgentClient, error) {
	cluster, err := grpcli.GetAnyClusterByID(ctx, clusterID)
	if err != nil {
		log.WithContextL(ctx).Error("failed to get cluster by id", zap.Uint64("cluster", clusterID), zap.Error(err))
		return nil, err
	}

	regionalEndpoint, err2 := BuildRegionalAPIEndpoint(ctx, provider, region, cluster.TenantId)
	if err2 != nil {
		log.WithContextL(ctx).Error("failed to build regional api endpoint", zap.Error(err2))
		return nil, err2
	}
	return &regionalAgentClientImpl{
		ctx:         ctx,
		clusterID:   clusterID,
		tenantID:    cluster.TenantId,
		provider:    provider,
		region:      region,
		regionalAPI: regionalapi.NewTCIDMHTTPClient(regionalEndpoint),
	}, nil
}

// BuildRegionalAPIEndpoint
// Example: https://regional-api.eastus.dev.azure.tidbcloud.com
func BuildRegionalAPIEndpoint(ctx context.Context, provider, physicalRegion string, tenantID uint64) (string, error) {
	if config.Global().DebugConfig.Enable && config.Global().DebugConfig.TCIDMRegionalServer != "" {
		return config.Global().DebugConfig.TCIDMRegionalServer, nil
	}
	env := config.Global().Env
	if env == "nightly" {
		env = "dev"
	}

	cloudProvider := convertToCloudProvider(provider)
	logger := log.WithContextL(ctx).With(
		zap.String("provider", provider),
		zap.String("cloudProvider", cloudProvider),
		zap.String("physical_region", physicalRegion),
		zap.Uint64("tenantID", tenantID),
	)

	var svcAddr string
	err := retry.Do(func() (err error) {
		logicalRegion, err := container.Global().RegionGrpcClient.GetRegionByTenantID(ctx, tenantID, cloudProvider, physicalRegion)
		if err != nil {
			logger.Error("failed to get region grpc client", zap.Error(err))
			return err
		}
		svcAddr = logicalRegion.GetRegion().GetServiceAddress()
		logger.Info("get regional api endpoint from region grpc client and override region", zap.String("svcAddr", svcAddr))
		return nil
	}, retry.Attempts(3), retry.Delay(100*time.Millisecond), retry.DelayType(retry.FixedDelay))
	if err != nil {
		logger.Error("failed to get regional api endpoint from region grpc client", zap.Error(err))
		return "", err
	}
	return svcAddr, nil
}

func convertToCloudProvider(provider string) string {
	switch provider {
	case infraapi.ProviderTypeEKS:
		return constants.CloudProviderAWS
	case infraapi.ProviderTypeGKE:
		return constants.CloudProviderGCP
	case infraapi.ProviderTypeAKS:
		return constants.CloudProviderAzure
	default:
		return provider
	}
}

type regionalAgentClientImpl struct {
	ctx          context.Context
	clusterID    uint64
	tenantID     uint64
	provider     string
	region       string
	proxyTimeout time.Duration
	regionalAPI  regionalapi.TCIDMHTTPClient
}

func (a *regionalAgentClientImpl) SetHttpTimeout(t time.Duration) {
	a.proxyTimeout = t
}

// proxyPOST used to send POST request to tcidm regional server
func (a *regionalAgentClientImpl) proxyPOST(ctx context.Context, path string, req any) ([]byte, error) {
	headers := map[string]string{}
	if a.proxyTimeout.Milliseconds() > 0 {
		headers["X-Timeout"] = strconv.FormatInt(a.proxyTimeout.Milliseconds(), 10)
	}
	return a.regionalAPI.ProxyPOST(ctx, path, req, headers)
}

// proxyToDataflowAgent used to send request to dataflow-agent which will proxy by tcidm regional server
func (a *regionalAgentClientImpl) proxyToDataflowAgent(ctx context.Context, method, path string, body any) ([]byte, error) {

	path = filepath.Join("proxy/dataflow-agent", path)
	extraHeaders := map[string]string{
		constants.ClusterIDHeader: strconv.FormatUint(a.clusterID, 10),
	}
	agentClient, err := a.proxyImpl()
	if err != nil {
		return nil, err
	}
	return agentClient.FetchResponseWithHeader(ctx, method, path, body, extraHeaders, nil)
}

// FetchResponseWithHeader used to send request for other client such as DMOpenapiClient、LightningClient
func (a *regionalAgentClientImpl) FetchResponseWithHeader(ctx context.Context,
	method, path string,
	req any,
	extraHeaders map[string]string,
	errorUnwrapper func(logger *zap.Logger, bodyBytes []byte) error,
) ([]byte, error) {

	logger := log.WithContextL(ctx).With(
		zap.String("method", method),
		zap.String("path", path),
		zap.Any("header", extraHeaders),
	)
	var (
		bodyBytes []byte
		err       error
	)
	switch method {
	case http.MethodGet:
		bodyBytes, err = a.regionalAPI.ProxyGET(ctx, path, extraHeaders)
	case http.MethodPost:
		bodyBytes, err = a.regionalAPI.ProxyPOST(ctx, path, req, extraHeaders)
	case http.MethodDelete:
		bodyBytes, err = a.regionalAPI.ProxyDELETE(ctx, path, extraHeaders)
	case http.MethodPut:
		bodyBytes, err = a.regionalAPI.ProxyPUT(ctx, path, req, extraHeaders)
	default:
		err = errors.New("[regionalAgentClientImpl.FetchResponseWithHeader]unsupported request method:" + method)
	}
	if err != nil {
		logger.Error("failed to get response", zap.Error(err))
		if errorUnwrapper != nil {
			if err := errorUnwrapper(logger, []byte(err.Error())); err != nil {
				return nil, err
			}
		}
		return nil, err
	}
	return bodyBytes, nil
}

func (a *regionalAgentClientImpl) CheckS3Access(ctx context.Context, req *agent.CheckS3AccessReq) (*agent.CheckS3AccessResp, error) {

	respBytes, err := a.proxyPOST(ctx, "/s3/check_access", CheckS3AccessParam{
		CheckS3Access: req,
		ClusterID:     a.clusterID,
	})
	if err != nil {
		return nil, err
	}
	var resp agent.CheckS3AccessResp
	if err := json.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (a *regionalAgentClientImpl) CheckGcsAccess(ctx context.Context, req *agent.CheckGcsAccessReq) (*agent.CheckGcsAccessResp, error) {

	respBytes, err := a.proxyPOST(ctx, "/gcs/check_access", CheckGcsAccessParam{
		CheckGcsAccess: req,
		ClusterID:      a.clusterID,
	})
	if err != nil {
		return nil, err
	}
	var resp agent.CheckGcsAccessResp
	if err := json.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, err
}

func (a *regionalAgentClientImpl) CheckAzblobAccess(ctx context.Context, req *agent.CheckAzblobAccessReq) (*agent.CheckAzblobAccessResp, error) {

	respBytes, err := a.proxyPOST(ctx, "/azblob/check_access", CheckAzblobAccessParam{
		CheckAzblobAccess: req,
		ClusterID:         a.clusterID,
	})
	if err != nil {
		return nil, err
	}
	var resp agent.CheckAzblobAccessResp
	if err := json.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, err
}

func (a *regionalAgentClientImpl) PrecheckImportJob(ctx context.Context, req *agent.PrecheckImportJobReq) (*agent.PrecheckImportJobResp, error) {
	path := fmt.Sprintf("/api/v1/orgs/%s/projects/%s/clusters/%s/imports:precheck", req.OrgId, req.ProjectId, req.ClusterId)
	// Use protojson to serialize the request to properly handle protobuf oneof fields (e.g., ImportSource.source)
	reqBytes, err := protojson.Marshal(req)
	if err != nil {
		return nil, errors.Trace(err)
	}
	respBytes, err := a.proxyToDataflowAgent(ctx, http.MethodPost, path, bytes.NewReader(reqBytes))
	if err != nil {
		return nil, err
	}
	var resp agent.PrecheckImportJobResp
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (a *regionalAgentClientImpl) LightningPrecheck(ctx context.Context, req *agent.LightningPrecheckReq) (*agent.LightningPrecheckResp, error) {
	respBytes, err := a.proxyPOST(ctx, "/import/precheck", ImportPrecheckParam{
		LightningPrecheck: req,
		ClusterID:         a.clusterID,
	})
	if err != nil {
		return nil, err
	}
	var resp agent.LightningPrecheckResp
	if err := json.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (a *regionalAgentClientImpl) ExecSQLOnCluster(ctx context.Context, req *agent.ExecSQLReq) (*agent.ExecSQLResp, error) {

	respBytes, err := a.proxyPOST(ctx, "/cluster/execsql", ExecSQLOnClusterParam{
		ExecSQL: req,
	})
	if err != nil {
		return nil, err
	}
	var resp agent.ExecSQLResp
	if err := json.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (a *regionalAgentClientImpl) CheckDBConnectivity(ctx context.Context, req *agent.CheckDBConnectivityReq) (*agent.CheckDBConnectivityResp, error) {

	respBytes, err := a.proxyToDataflowAgent(ctx, http.MethodPost, "/api/v1/db-connectivity", req)
	if err != nil {
		return nil, err
	}
	var resp agent.CheckDBConnectivityResp
	if err := json.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (a *regionalAgentClientImpl) GetClusterClientSecret(ctx context.Context, req *agent.GetClusterClientSecretReq) (*agent.GetClusterClientSecretResp, error) {

	respBytes, err := a.proxyToDataflowAgent(ctx, http.MethodPost, "/api/v1/cluster-client-secret", req)
	if err != nil {
		return nil, err
	}
	var resp agent.GetClusterClientSecretResp
	if err := json.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (a *regionalAgentClientImpl) GetTiDBServerSecret(ctx context.Context, req *agent.GetTiDBServerSecretReq) (*agent.GetTiDBServerSecretResp, error) {

	respBytes, err := a.proxyToDataflowAgent(ctx, http.MethodPost, "/api/v1/tidb-server-secret", req)
	if err != nil {
		return nil, err
	}
	var resp agent.GetTiDBServerSecretResp
	if err := json.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (a *regionalAgentClientImpl) proxyImpl() (*regionalAgentClientImpl, error) {
	agentClient := *a
	regionalEndpoint, err := BuildRegionalAPIEndpoint(a.ctx, a.provider, a.region, a.tenantID)
	if err != nil {
		log.WithContextL(a.ctx).Error("failed to build regional api endpoint", zap.Error(err))
		return nil, err
	}
	agentClient.regionalAPI = regionalapi.NewTCIDMHTTPClient(regionalEndpoint).
		SetResponseHandler(regionalapi.RawResponseHandler)
	return &agentClient, nil
}

func (a *regionalAgentClientImpl) GetEngineClient(engineAddr string) *EngineClient {
	return nil
}

func (a *regionalAgentClientImpl) GetDMOpenapiClient(clusterID uint64, dmAddr string) *DMOpenapiClient {
	agentClient := *a
	agentClient.proxyTimeout = defaultTimeout
	regionalEndpoint, err := BuildRegionalAPIEndpoint(a.ctx, a.provider, a.region, a.tenantID)
	if err != nil {
		// BuildRegionalAPIEndpoint will retry 3 times, so we can panic here, current request will be failed
		log.WithContextL(a.ctx).Panic("failed to build regional api endpoint", zap.Error(err))
	}

	agentClient.regionalAPI = regionalapi.NewTCIDMHTTPClient(regionalEndpoint).
		SetResponseHandler(regionalapi.RawResponseHandler)
	return &DMOpenapiClient{
		AgentClient: &agentClient,
		dmAddr:      dmAddr,
		clusterID:   clusterID,
	}
}

func (a *regionalAgentClientImpl) GetLightningClient(clusterID, importID uint64) LightningAPI {
	agentClient := *a
	agentClient.proxyTimeout = defaultTimeout
	ns := fmt.Sprintf("tidb%d", clusterID)
	serviceName := fmt.Sprintf("import-%d", importID)
	return &regionalLightningClient{
		AgentClient: &agentClient,
		ns:          ns,
		serviceName: serviceName,
		clusterID:   clusterID,
	}
}

func (a *regionalAgentClientImpl) PrecheckDMJob(ctx context.Context, req *agent.PrecheckDMJobReq) (*agent.PrecheckDMJobResp, error) {

	respBytes, err := a.proxyToDataflowAgent(ctx, http.MethodPost, "/api/v1/dm/jobprecheck", req)
	if err != nil {
		return nil, err
	}
	var resp agent.PrecheckDMJobResp
	if err := json.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (a *regionalAgentClientImpl) GetSourceTables(ctx context.Context, req *agent.CheckDBConnectivityReq) (*agent.GetSourceTablesResp, error) {

	respBytes, err := a.proxyToDataflowAgent(ctx, http.MethodPost, "/api/v1/source-tables", req)
	if err != nil {
		return nil, err
	}
	var resp agent.GetSourceTablesResp
	if err := json.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (a *regionalAgentClientImpl) GetSourceBinlogStatus(ctx context.Context, req *agent.CheckDBConnectivityReq) (*agent.GetSourceBinlogStatusResp, error) {

	respBytes, err := a.proxyToDataflowAgent(ctx, http.MethodPost, "/api/v1/source/binlog-status", req)
	if err != nil {
		return nil, err
	}
	var resp agent.GetSourceBinlogStatusResp
	if err := json.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (a *regionalAgentClientImpl) GetClusterTables(ctx context.Context, req *agent.Cluster) (*agent.GetClusterTablesResp, error) {
	respBytes, err := a.proxyPOST(ctx, "/cluster-tables", GetClusterTablesParam{
		Cluster: req,
	})
	if err != nil {
		return nil, err
	}

	var resp agent.GetClusterTablesResp
	if err := json.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (a *regionalAgentClientImpl) GetUserFileContent(ctx context.Context, req *agent.GetUserFileReq) (*agent.GetUserFileResp, error) {

	respBytes, err := a.proxyToDataflowAgent(ctx, http.MethodPost, "/api/v1/user-file", req)
	if err != nil {
		return nil, err
	}
	var resp agent.GetUserFileResp
	if err := json.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (a *regionalAgentClientImpl) NsLookup(ctx context.Context, req *agent.NsLookupReq) (*agent.NsLookupResp, error) {

	respBytes, err := a.proxyToDataflowAgent(ctx, http.MethodPost, "/api/v1/nslookup", req)
	if err != nil {
		return nil, err
	}
	var resp agent.NsLookupResp
	if err := json.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (a *regionalAgentClientImpl) LoadDataToCluster(ctx context.Context, req *agent.LoadDataReq) (*agent.LoadDataResp, error) {

	respBytes, err := a.proxyToDataflowAgent(ctx, http.MethodPost, "/api/v1/cluster/load-data", req)
	if err != nil {
		return nil, err
	}
	var resp agent.LoadDataResp
	if err := json.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (a *regionalAgentClientImpl) PreviewCsvFile(ctx context.Context, req *agent.PreviewCsvFileReq) (*agent.PreviewCsvFileResp, error) {
	respBytes, err := a.proxyToDataflowAgent(ctx, http.MethodPost, "/api/v1/import/preview-csv-file", req)
	if err != nil {
		return nil, err
	}
	var resp agent.PreviewCsvFileResp
	if err := json.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (a *regionalAgentClientImpl) SearchImportFiles(ctx context.Context, req *agent.SearchImportFilesReq) (*agent.SearchImportFilesResp, error) {
	respBytes, err := a.proxyPOST(ctx, "/import/search-import-files", SearchImportFilesParam{
		SearchImportFiles: req,
		ClusterID:         a.clusterID,
	})
	if err != nil {
		return nil, err
	}
	var resp agent.SearchImportFilesResp
	if err := json.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (a *regionalAgentClientImpl) PrecheckImportFile(ctx context.Context, req *agent.PrecheckImportFileReq) (*emptypb.Empty, error) {

	_, err := a.proxyPOST(ctx, "/import/precheck-import-file", PrecheckImportFileParam{
		PrecheckImportFile: req,
		ClusterID:          a.clusterID,
	})
	if err != nil {
		return nil, err
	}

	return &emptypb.Empty{}, nil
}

func (a *regionalAgentClientImpl) DropClusterTables(ctx context.Context, req *agent.DropClusterTablesReq) (*emptypb.Empty, error) {

	_, err := a.proxyToDataflowAgent(ctx, http.MethodPost, "/api/v1/import/drop-cluster-tables", req)
	if err != nil {
		return nil, err
	}
	return &emptypb.Empty{}, nil
}