//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://gitee.com/micro-tools/wf/extend/utils.

package grpctracing

import (
	"context"
	"gitee.com/micro-tools/wf/extend/utils/krpc/internal"
	"gitee.com/micro-tools/wf/extend/utils/krpc/internal/grpcctx"
	"gitee.com/micro-tools/wf/extend/utils/krpc/internal/grpcutils"
	"gitee.com/micro-tools/wf/net/gtrace"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/baggage"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/trace"
	"google.golang.org/grpc"
	grpcCodes "google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
	"sync"
)

var mu sync.Mutex

// UnaryServerInterceptor returns a grpc.UnaryServerInterceptor suitable
// for use in a grpc.NewServer call.
func UnaryServerInterceptor(
	ctx context.Context,
	req interface{},
	info *grpc.UnaryServerInfo,
	handler grpc.UnaryHandler) (interface{}, error) {
	tracer := otel.GetTracerProvider().Tracer(
		tracingInstrumentGrpcServer,
		trace.WithInstrumentationVersion(internal.VERSION),
	)
	requestMetadata, _ := metadata.FromIncomingContext(ctx)
	metadataCopy := requestMetadata.Copy()

	entries, spanCtx := Extract(ctx, metadataCopy)
	ctx = baggage.ContextWithValues(ctx, entries...)
	ctx = trace.ContextWithRemoteSpanContext(ctx, spanCtx)
	name, attr := spanInfo(info.FullMethod, peerFromCtx(ctx))
	ctx, span := tracer.Start(
		ctx,
		name,
		trace.WithSpanKind(trace.SpanKindServer),
		trace.WithAttributes(attr...),
	)

	defer func() {
		span.End()
	}()

	span.SetAttributes(gtrace.CommonLabels()...)

	span.AddEvent(tracingEventGrpcRequest, trace.WithAttributes(
		attribute.Any(tracingEventGrpcRequestBaggage, gtrace.GetBaggageMap(ctx)),
		attribute.Any(tracingEventGrpcMetadataIncoming, grpcctx.Ctx.IncomingMap(ctx)),
		attribute.String(
			tracingEventGrpcRequestMessage,
			grpcutils.MarshalMessageToJsonStringForTracing(
				req, "Request", tracingMaxContentLogSize,
			),
		),
	))

	res, err := handler(ctx, req)

	span.AddEvent(tracingEventGrpcResponse, trace.WithAttributes(
		attribute.String(
			tracingEventGrpcResponseMessage,
			grpcutils.MarshalMessageToJsonStringForTracing(
				res, "Response", tracingMaxContentLogSize,
			),
		),
	))

	if err != nil {
		s, _ := status.FromError(err)
		span.SetStatus(codes.Error, s.Message())
		span.SetAttributes(statusCodeAttr(s.Code()))
	} else {
		span.SetAttributes(statusCodeAttr(grpcCodes.OK))
	}

	return res, err
}

// StreamServerInterceptor returns a grpc.StreamServerInterceptor suitable
// for use in a grpc.NewServer call.
func StreamServerInterceptor(
	srv interface{},
	ss grpc.ServerStream,
	info *grpc.StreamServerInfo,
	handler grpc.StreamHandler) error {
	tracer := otel.GetTracerProvider().Tracer(
		tracingInstrumentGrpcServer,
		trace.WithInstrumentationVersion(internal.VERSION),
	)

	ctx := ss.Context()
	requestMetadata, _ := metadata.FromIncomingContext(ctx)
	metadataCopy := requestMetadata.Copy()
	entries, spanCtx := Extract(ctx, metadataCopy)
	ctx = baggage.ContextWithValues(ctx, entries...)
	ctx = trace.ContextWithRemoteSpanContext(ctx, spanCtx)
	name, attr := spanInfo(info.FullMethod, peerFromCtx(ctx))
	ctx, span := tracer.Start(
		ctx,
		name,
		trace.WithSpanKind(trace.SpanKindServer),
		trace.WithAttributes(attr...),
	)
	defer span.End()

	span.SetAttributes(gtrace.CommonLabels()...)

	err := handler(srv, wrapServerStream(ctx, ss))

	if err != nil {
		s, _ := status.FromError(err)
		span.SetStatus(codes.Error, s.Message())
		span.SetAttributes(statusCodeAttr(s.Code()))
	} else {
		span.SetAttributes(statusCodeAttr(grpcCodes.OK))
	}

	return err
}