// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package trace

import (
	"strings"

	"golang.org/x/net/context"
	"google.golang.org/api/option"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

const grpcMetadataKey = "x-cloud-trace-context"

// GRPCClientInterceptor returns a grpc.UnaryClientInterceptor that traces all outgoing requests from a gRPC client.
// The calling context should already have a *trace.Span; a child span will be
// created for the outgoing gRPC call. If the calling context doesn't have a span,
// the call will not be traced.
//
// The functionality in gRPC that this feature relies on is currently experimental.
func GRPCClientInterceptor() grpc.UnaryClientInterceptor {
	return grpc.UnaryClientInterceptor(grpcUnaryInterceptor)
}

func grpcUnaryInterceptor(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
	// TODO: also intercept streams.
	span := FromContext(ctx).NewChild(method)
	defer span.Finish()

	if span != nil {
		header := spanHeader(span.trace.traceID, span.span.ParentSpanId, span.trace.globalOptions)
		md, ok := metadata.FromOutgoingContext(ctx)
		if !ok {
			md = metadata.Pairs(grpcMetadataKey, header)
		} else {
			md = md.Copy() // metadata is immutable, copy.
			md[grpcMetadataKey] = []string{header}
		}
		ctx = metadata.NewOutgoingContext(ctx, md)
	}

	err := invoker(ctx, method, req, reply, cc, opts...)
	if err != nil {
		// TODO: standardize gRPC label names?
		span.SetLabel("error", err.Error())
	}
	return err
}

// GRPCServerInterceptor returns a grpc.UnaryServerInterceptor that enables the tracing of the incoming
// gRPC calls. Incoming call's context can be used to extract the span on servers that enabled this option:
//
//	span := trace.FromContext(ctx)
//
// The functionality in gRPC that this feature relies on is currently experimental.
func GRPCServerInterceptor(tc *Client) grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
		md, _ := metadata.FromIncomingContext(ctx)
		if header, ok := md[grpcMetadataKey]; ok {
			span := tc.SpanFromHeader("", strings.Join(header, ""))
			defer span.Finish()
			ctx = NewContext(ctx, span)
		}
		return handler(ctx, req)
	}
}

// EnableGRPCTracing automatically traces all outgoing gRPC calls from cloud.google.com/go clients.
//
// The functionality in gRPC that this relies on is currently experimental.
//
// Deprecated: Use option.WithGRPCDialOption(grpc.WithUnaryInterceptor(GRPCClientInterceptor())) instead.
var EnableGRPCTracing option.ClientOption = option.WithGRPCDialOption(grpc.WithUnaryInterceptor(GRPCClientInterceptor()))