package grpc

import (
	"context"
	"fmt"
	"time"

	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	grpcmetadata "google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/proto"

	"github.com/tidbcloud/tidb-management-service/utilities/grpc/metadata"
)

func ClientPayloadLogInterceptor(
	ctx context.Context,
	method string,
	req, reply any,
	cc *grpc.ClientConn,
	invoker grpc.UnaryInvoker,
	opts ...grpc.CallOption,
) error {
	startTime := time.Now()

	reqSnapshot := proto.Clone(req.(proto.Message))

	err := invoker(ctx, method, req, reply, cc, opts...)

	defer func() {
		d := time.Since(startTime)
		marshaller := protojson.MarshalOptions{
			UseProtoNames:   true,
			EmitUnpopulated: true,
		}
		// Process the request
		reqJson := marshaller.Format(reqSnapshot)
		// Process the response
		respBp, ok := reply.(proto.Message)
		if !ok {
			tilog.Warn(
				ctx,
				"response is not proto.Message",
				zap.String("resp_type", fmt.Sprintf("%T", reply)),
				zap.Any("response", reply),
			)
		}

		// Extract base response error
		if baseResp := getBaseRespFromResp(reply); baseResp.GetErrCode() != 0 {
			err = status.Error(codes.Code(baseResp.GetErrCode()), baseResp.GetErrMsg())
		}

		respJson := marshaller.Format(respBp)
		// Remove sensitive information
		reqJson = RedactSecrets(reqJson)
		respJson = RedactSecrets(respJson)

		if err != nil {
			// Process the error
			errStatus, ok := status.FromError(err)
			if !ok {
				tilog.Warn(ctx, "error is not grpc status", zap.Error(err))
				errStatus = status.New(codes.Unknown, err.Error())
			}
			errJson := marshaller.Format(errStatus.Proto())
			errJson = RedactSecrets(errJson)

			// Disable stack trace here.
			// The Go error stack and the call stack of the current goroutine:
			//  - Capturing the call stack directly in the gRPC log interceptor does not yield the business call stack, as the business function has already returned by that point.
			//  - The error should include the stack where the error was returned.
			GetPayloadLogger(ctx).Error(
				"---> gRPC client",
				zap.String("url", cc.Target()),
				zap.String("method", method),
				zap.String("request", reqJson),
				zap.String("response", respJson),
				zap.Duration("cost", d),
				zap.String("error", errJson),
			)
			return
		}

		GetPayloadLogger(ctx).Info(
			"---> gRPC client",
			zap.String("url", cc.Target()),
			zap.String("method", method),
			zap.String("request", reqJson),
			zap.String("response", respJson),
			zap.Duration("cost", d),
		)
	}()

	return err
}

func TimeoutInterceptor(timeout time.Duration) grpc.UnaryClientInterceptor {
	return func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
		if timeout <= 0 {
			return invoker(ctx, method, req, reply, cc, opts...)
		}
		ctx, cancel := context.WithTimeout(ctx, timeout)
		defer cancel()
		return invoker(ctx, method, req, reply, cc, opts...)
	}
}

func AppendRequestMetaInterceptor(
	ctx context.Context,
	method string,
	req, reply any,
	cc *grpc.ClientConn,
	invoker grpc.UnaryInvoker,
	opts ...grpc.CallOption,
) error {
	// Append request meta here

	// check outgoing meta
	md, ok := grpcmetadata.FromOutgoingContext(ctx)
	// not set request meta
	// orgid or projectid must be one
	if !ok ||
		(md[metadata.OrgIDHeaderKey] == nil || md[metadata.DefaultProjectIDHeaderKey] == nil || md[metadata.ClusterIDHeaderKey] == nil) {
		reqMeta, err := metadata.ExtractRequestMeta(ctx)
		if err != nil {
			// only log error
			// not break the request
			tilog.Warn(ctx, "extract original request from context failed", zap.Error(err))
		} else {
			ctx = metadata.AppendToOutgoingMeta(ctx, reqMeta)
		}
	}

	return invoker(ctx, method, req, reply, cc, opts...)
}