package agent

import (
	"context"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	base "tcidm/commons/http"
	"tcidm/metrics"
	"time"

	"github.com/pingcap/errors"
	dataflowgrpc "github.com/tidbcloud/dataflow-service/idl/pbgen/proto/agent"
	"github.com/tidbcloud/tiinf/errs"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"google.golang.org/genproto/googleapis/rpc/status"
	"google.golang.org/grpc/codes"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/types/known/emptypb"
)

// @Time    : 2024/8/27 15:06
// @Author  : lvhongmeng

//const defaultTimeout = 27 * time.Second

type AgentClient struct {
	BaseURL string

	Client base.HttpRequester
}

func (a *AgentClient) SetHttpTimeout(t time.Duration) {
	a.Client.SetHttpTimeout(t)
}

func (a *AgentClient) getSchema() string {
	return "http"
}

func (a *AgentClient) getResponse(ctx context.Context, method, url string, req any) ([]byte, error) {
	return a.getResponseWithHeader(ctx, method, url, req, nil, defaultErrorUnwrapper)
}

func (a *AgentClient) getResponseWithHeader(ctx context.Context,
	method, reqUrl string,
	req any,
	extraHeaders map[string]string,
	errorUnwrapper func(logger *zap.Logger, bodyBytes []byte) error,
) ([]byte, error) {
	logger := tilog.WithContextL(ctx)

	var err error
	startTime := time.Now()
	defer func() {
		cost := time.Since(startTime).Milliseconds()
		ok := strconv.FormatBool(err == nil)
		u, _ := url.Parse(reqUrl)
		metrics.DownstreamRequestDuration.WithLabelValues("dataflow-agent", u.Path, ok).Observe(float64(cost))
	}()

	statusCode, bodyBytes, err := a.Client.RequestWithHeaders(ctx, method, reqUrl, req, extraHeaders)
	if err != nil {
		logger.Error("failed to get response", zap.Error(err))
		return nil, err
	}

	var printedBodyBytes string
	if len(printedBodyBytes) > 200 {
		printedBodyBytes = string(bodyBytes[:200]) + "..."
	} else {
		printedBodyBytes = string(bodyBytes)
	}
	logger.Info("received response",
		zap.String("method", method),
		zap.String("url", reqUrl),
		zap.Int("status-code", statusCode),
		zap.String("resp", printedBodyBytes))

	if statusCode != http.StatusOK {
		if err = errorUnwrapper(logger, bodyBytes); err != nil {
			return nil, err
		}
	}
	return bodyBytes, nil
}

var grpcErrCode2TiErr = map[codes.Code]errs.Error{
	codes.Canceled:           errs.ErrInternalTimeout,
	codes.Unknown:            errs.ErrInternalUnknown,
	codes.InvalidArgument:    errs.ErrBadRequest,
	codes.DeadlineExceeded:   errs.ErrInternalTimeout,
	codes.NotFound:           errs.ErrNotFound,
	codes.AlreadyExists:      errs.ErrRecordExist,
	codes.PermissionDenied:   errs.ErrForbidden,
	codes.ResourceExhausted:  errs.ErrTooManyRequests,
	codes.FailedPrecondition: errs.ErrInternalGRPC,
	codes.Aborted:            errs.ErrInternalGRPC,
	codes.OutOfRange:         errs.ErrBadRequest,
	codes.Unimplemented:      errs.ErrInternalGRPC,
	codes.Internal:           errs.ErrInternalUnknown,
	codes.Unavailable:        errs.ErrInternalUnknown,
	codes.DataLoss:           errs.ErrInternalUnknown,
	codes.Unauthenticated:    errs.ErrUnauthorized,
}

func defaultErrorUnwrapper(logger *zap.Logger, bodyBytes []byte) error {
	var errResp status.Status
	if err := protojson.Unmarshal(bodyBytes, &errResp); err == nil {
		if tierr, ok := grpcErrCode2TiErr[codes.Code(errResp.Code)]; ok {
			message := errResp.Message
			if codes.Code(errResp.Code) == codes.DeadlineExceeded {
				message = context.DeadlineExceeded.Error() + ":" + errResp.Message
			}
			return errs.NewError(tierr, message, nil)
		} else {
			return errors.New(errResp.Message)
		}
	} else {
		logger.Error("failed to unmarshal response body as grpc Status", zap.Error(err))
	}
	return nil
}

func GetAgentClient(ctx context.Context, dataflowAgentUrl string) (*AgentClient, error) {
	var client base.HttpRequester

	client = base.NewAgentHttpRequester()

	return &AgentClient{BaseURL: dataflowAgentUrl, Client: client}, nil
}

func (a *AgentClient) LightningPrecheck(ctx context.Context, req *dataflowgrpc.LightningPrecheckReq) (*dataflowgrpc.LightningPrecheckResp, error) {
	url := fmt.Sprintf("%s://%s/api/v1/import/lightning-precheck", a.getSchema(), a.BaseURL)
	respBytes, err := a.getResponse(ctx, http.MethodPost, url, req)
	if err != nil {
		return nil, err
	}
	var resp dataflowgrpc.LightningPrecheckResp
	// 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 (a *AgentClient) SearchImportFiles(ctx context.Context, req *dataflowgrpc.SearchImportFilesReq) (*dataflowgrpc.SearchImportFilesResp, error) {
	url := fmt.Sprintf("%s://%s/api/v1/import/search-import-files", a.getSchema(), a.BaseURL)

	respBytes, err := a.getResponse(ctx, http.MethodPost, url, req)
	if err != nil {
		return nil, err
	}
	var resp dataflowgrpc.SearchImportFilesResp
	// 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 (a *AgentClient) PrecheckImportFile(ctx context.Context, req *dataflowgrpc.PrecheckImportFileReq) (*emptypb.Empty, error) {
	url := fmt.Sprintf("%s://%s/api/v1/import/precheck-import-file", a.getSchema(), a.BaseURL)

	_, err := a.getResponse(ctx, http.MethodPost, url, req)
	if err != nil {
		return nil, err
	}

	return &emptypb.Empty{}, nil
}

func (a *AgentClient) CheckS3Access(ctx context.Context, req *dataflowgrpc.CheckS3AccessReq) (*dataflowgrpc.CheckS3AccessResp, error) {
	url := fmt.Sprintf("%s://%s/api/v1/s3/check_access", a.getSchema(), a.BaseURL)

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

func (a *AgentClient) CheckGcsAccess(ctx context.Context, req *dataflowgrpc.CheckGcsAccessReq) (*dataflowgrpc.CheckGcsAccessResp, error) {
	url := fmt.Sprintf("%s://%s/api/v1/gcs/check_access", a.getSchema(), a.BaseURL)

	respBytes, err := a.getResponse(ctx, http.MethodPost, url, req)
	if err != nil {
		return nil, err
	}
	var resp dataflowgrpc.CheckGcsAccessResp
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, err
}

func (a *AgentClient) CheckAzblobAccess(ctx context.Context, req *dataflowgrpc.CheckAzblobAccessReq) (*dataflowgrpc.CheckAzblobAccessResp, error) {
	url := fmt.Sprintf("%s://%s/api/v1/azblob/check_access", a.getSchema(), a.BaseURL)

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

func (a *AgentClient) GetClusterTables(ctx context.Context, req *dataflowgrpc.Cluster) (*dataflowgrpc.GetClusterTablesResp, error) {
	url := fmt.Sprintf("%s://%s/api/v1/cluster-tables", a.getSchema(), a.BaseURL)
	respBytes, err := a.getResponse(ctx, http.MethodPost, url, req)
	if err != nil {
		return nil, err
	}
	var resp dataflowgrpc.GetClusterTablesResp
	// 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 (a *AgentClient) ExecSQLOnCluster(ctx context.Context, req *dataflowgrpc.ExecSQLReq) (*dataflowgrpc.ExecSQLResp, error) {
	url := fmt.Sprintf("%s://%s/api/v1/cluster/execsql", a.getSchema(), a.BaseURL)

	respBytes, err := a.getResponse(ctx, http.MethodPost, url, req)
	if err != nil {
		return nil, err
	}
	var resp dataflowgrpc.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
}