package grpctrace

import (
	"regexp"
	"runtime/pprof"

	"code.justin.tv/common/chitin/internal/trace"
	"code.justin.tv/release/trace/common"
	"code.justin.tv/release/trace/events"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

var (
	txidKey = "chitin-txid"

	clientProfile *pprof.Profile
)

func init() {
	pkgname := "code.justin.tv/common/chitin/grpctrace"

	// The pprof tool expects profiles to match the following regexp:
	//
	//     `\A(\w+) profile: total \d+\n\z`
	//
	// This means that the profile name can only contain the characters [0-9A-
	// Za-z_], contrary to the runtime/pprof package docs, which say:
	//
	//     The convention is to use a 'import/path.' prefix to create separate
	//     name spaces for each package.
	//
	// See https://golang.org/issue/13195
	cmdPprofName := regexp.MustCompile("[^0-9A-Za-z_]")
	clientProfile = pprof.NewProfile(cmdPprofName.ReplaceAllString(pkgname, "_") + "_client")
}

func Invoke(ctx context.Context, method string, args, reply interface{}, cc *grpc.ClientConn, opts ...grpc.CallOption) error {
	ctx = trace.NewSpan(ctx)
	md, _ := metadata.FromContext(ctx)
	md = md.Copy()
	trace.AugmentGRPCMetadata(ctx, md)
	ctx = metadata.NewContext(ctx, md)

	pprofKey := new(byte) // This needs to have a unique address, so it must not be size zero
	clientProfile.Add(pprofKey, 1)

	// TODO: find a way to get the peer
	ev := &events.Event{
		Kind: common.Kind_REQUEST_HEAD_PREPARED,
		Extra: &events.Extra{
			Grpc: &events.ExtraGRPC{
				Method: method,
			},
		},
	}
	trace.SendEvent(ctx, ev)

	defer func() {
		// TODO: report whether grpc.Invoke returned an error
		ev := &events.Event{
			Kind: common.Kind_RESPONSE_BODY_RECEIVED,
		}
		trace.SendEvent(ctx, ev)

		clientProfile.Remove(pprofKey)
	}()

	return grpc.Invoke(ctx, method, args, reply, cc, opts...)
}

func NewClientStream(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) {
	ctx = trace.NewSpan(ctx)
	md, _ := metadata.FromContext(ctx)
	md = md.Copy()
	trace.AugmentGRPCMetadata(ctx, md)
	ctx = metadata.NewContext(ctx, md)

	pprofKey := new(byte) // This needs to have a unique address, so it must not be size zero
	clientProfile.Add(pprofKey, 1)

	ev := &events.Event{
		Kind: common.Kind_REQUEST_HEAD_PREPARED,
		Extra: &events.Extra{
			Grpc: &events.ExtraGRPC{
				Method: method,
			},
		},
	}
	trace.SendEvent(ctx, ev)

	cs, err := grpc.NewClientStream(ctx, desc, cc, method, opts...)

	ourCS := &clientStream{
		stream: stream{
			ctx:    ctx,
			stream: cs,
			recvErr: func(err error) {
				ev := &events.Event{
					Kind: common.Kind_RESPONSE_BODY_RECEIVED,
				}
				trace.SendEvent(ctx, ev)

				clientProfile.Remove(pprofKey)
			},
		},
		client: cs,
	}

	return ourCS, err
}
