package grpc

import (
	"bytes"
	"context"
	"fmt"
	"time"

	"github.com/tidbcloud/cluster-service-ng/pkg/utils"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/proto"
)

// marshalStatusToJSON marshals gRPC status to JSON without HTML escaping.
func marshalStatusToJSON(s *status.Status) string {
	marshaller := protojson.MarshalOptions{
		UseProtoNames:   true,
		EmitUnpopulated: true,
	}
	data, err := marshaller.Marshal(s.Proto())
	if err != nil {
		return fmt.Sprintf(`{"error": %q}`, err.Error())
	}
	// Unescape HTML characters (<, >, &) for readable logs
	data = bytes.ReplaceAll(data, []byte(`\u003c`), []byte(`<`))
	data = bytes.ReplaceAll(data, []byte(`\u003e`), []byte(`>`))
	data = bytes.ReplaceAll(data, []byte(`\u0026`), []byte(`&`))
	return string(data)
}

func ServerPayloadLogInterceptor(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
	startTime := time.Now()

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

	resp, err := handler(ctx, req)

	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 := resp.(proto.Message)
		respJson := "null"
		if !ok {
			tilog.Warn(ctx, "response is not proto.Message", zap.String("resp_type", fmt.Sprintf("%T", resp)))
		} else {
			respJson = marshaller.Format(respBp)
		}
		// Remove sensitive information
		reqJson = utils.RedactSecrets(reqJson)
		respJson = utils.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 := marshalStatusToJSON(errStatus)
			errJson = utils.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.
			tilog.WithContextL(ctx).WithOptions(zap.WithCaller(false), zap.AddStacktrace(zap.FatalLevel)).Error(
				"<--- grpc server",
				zap.String("method", info.FullMethod),
				zap.String("request", reqJson),
				zap.String("response", respJson),
				zap.Duration("duration", d),
				zap.String("error", errJson),
			)

			return
		}

		tilog.WithContextL(ctx).WithOptions(zap.WithCaller(false)).Info(
			"<--- grpc server",
			zap.String("method", info.FullMethod),
			zap.String("request", reqJson),
			zap.String("response", respJson),
			zap.Duration("duration", d),
		)
	}()

	return resp, err
}

// CallingTimeoutInterceptor returns a new unary client interceptor that sets a timeout on the request context.
func CallingTimeoutInterceptor(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 {
		timedCtx, cancel := context.WithTimeout(ctx, timeout)
		defer cancel()
		return invoker(timedCtx, method, req, reply, cc, opts...)
	}
}

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)
		respJson := "null"
		if !ok {
			tilog.Warn(ctx, "response is not proto.Message", zap.String("resp_type", fmt.Sprintf("%T", reply)))
		} else {
			respJson = marshaller.Format(respBp)
		}
		// Remove sensitive information
		reqJson = utils.RedactSecrets(reqJson)
		respJson = utils.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 := marshalStatusToJSON(errStatus)
			errJson = utils.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.
			tilog.WithContextL(ctx).WithOptions(zap.WithCaller(false), zap.AddStacktrace(zap.FatalLevel)).Error(
				"---> gRPC client",
				zap.String("url", cc.Target()),
				zap.String("method", method),
				zap.String("request", reqJson),
				zap.String("response", respJson),
				zap.Duration("duration", d),
				zap.String("error", errJson),
			)
			return
		}

		tilog.WithContextL(ctx).WithOptions(zap.WithCaller(false)).Info(
			"---> gRPC client",
			zap.String("url", cc.Target()),
			zap.String("method", method),
			zap.String("request", reqJson),
			zap.String("response", respJson),
			zap.Duration("duration", d),
		)
	}()

	return err
}