package grpccorrelation

import (
	"context"

	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
	"gitlab.com/gitlab-org/labkit/correlation"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

func extractFromContext(ctx context.Context, propagateIncomingCorrelationID bool) context.Context {
	generateCorrelationID := true
	md, ok := metadata.FromIncomingContext(ctx)
	if ok {
		if propagateIncomingCorrelationID {
			// Extract correlation_id
			values := md.Get(metadataCorrelatorKey)
			if len(values) > 0 && values[0] != "" {
				ctx = correlation.ContextWithCorrelation(ctx, values[0])
				generateCorrelationID = false
			}
		}

		// Extract client name
		clientNames := md.Get(metadataClientNameKey)
		if len(clientNames) > 0 {
			ctx = correlation.ContextWithClientName(ctx, clientNames[0])
		}
	}
	if generateCorrelationID {
		ctx = correlation.ContextWithCorrelation(ctx, correlation.SafeRandomID())
	}
	return ctx
}

// UnaryServerCorrelationInterceptor propagates Correlation-IDs from incoming upstream services.
func UnaryServerCorrelationInterceptor(opts ...ServerCorrelationInterceptorOption) grpc.UnaryServerInterceptor {
	config := applyServerCorrelationInterceptorOptions(opts)
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
		ctx = extractFromContext(ctx, config.propagateIncomingCorrelationID)
		return handler(ctx, req)
	}
}

// StreamServerCorrelationInterceptor propagates Correlation-IDs from incoming upstream services.
func StreamServerCorrelationInterceptor(opts ...ServerCorrelationInterceptorOption) grpc.StreamServerInterceptor {
	config := applyServerCorrelationInterceptorOptions(opts)
	return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
		wrapped := grpc_middleware.WrapServerStream(ss)
		wrapped.WrappedContext = extractFromContext(ss.Context(), config.propagateIncomingCorrelationID)

		return handler(srv, wrapped)
	}
}
