package grpctrace

import (
	"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"
	"google.golang.org/grpc/transport"
)

func RegisterService(s *grpc.Server, sd *grpc.ServiceDesc, ss interface{}) {
	desc := &grpc.ServiceDesc{
		ServiceName: sd.ServiceName,
		HandlerType: sd.HandlerType,
		Methods:     make([]grpc.MethodDesc, len(sd.Methods)),
		Streams:     make([]grpc.StreamDesc, len(sd.Streams)),
	}

	for i := range sd.Methods {
		w := &wrapper{
			MethodHandler: sd.Methods[i].Handler,
		}
		desc.Methods[i] = grpc.MethodDesc{
			MethodName: sd.Methods[i].MethodName,
			Handler:    w.method,
		}
	}
	for i := range sd.Streams {
		w := &wrapper{
			StreamHandler: sd.Streams[i].Handler,
		}
		desc.Streams[i] = grpc.StreamDesc{
			StreamName:    sd.Streams[i].StreamName,
			Handler:       w.stream,
			ServerStreams: sd.Streams[i].ServerStreams,
			ClientStreams: sd.Streams[i].ClientStreams,
		}
	}

	s.RegisterService(desc, ss)
}

type wrapper struct {
	MethodHandler func(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error)
	StreamHandler func(srv interface{}, stream grpc.ServerStream) error
}

func (w *wrapper) method(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
	md, _ := metadata.FromContext(ctx)
	ctx = contextFromGRPCMetadata(ctx, md)
	md = md.Copy()
	augmentGRPCMetadata(ctx, md)

	if stream, ok := transport.StreamFromContext(ctx); ok {
		ev := &events.Event{
			Kind: common.Kind_REQUEST_BODY_RECEIVED,
			Extra: &events.Extra{
				Grpc: &events.ExtraGRPC{
					Method: stream.Method(),
				},
			},
		}
		trace.SendEvent(ctx, ev)
		defer func() {
			ev := &events.Event{
				Kind: common.Kind_RESPONSE_BODY_SENT,
			}
			trace.SendEvent(ctx, ev)
		}()
	}

	return w.MethodHandler(srv, ctx, dec, interceptor)
}

func (w *wrapper) stream(srv interface{}, ss grpc.ServerStream) error {
	ctx := ss.Context()
	md, _ := metadata.FromContext(ctx)
	ctx = contextFromGRPCMetadata(ctx, md)
	md = md.Copy()
	augmentGRPCMetadata(ctx, md)

	if stream, ok := transport.StreamFromContext(ctx); ok {
		ev := &events.Event{
			Kind: common.Kind_REQUEST_HEAD_RECEIVED,
			Extra: &events.Extra{
				Grpc: &events.ExtraGRPC{
					Method: stream.Method(),
				},
			},
		}
		trace.SendEvent(ctx, ev)
		defer func() {
			ev := &events.Event{
				Kind: common.Kind_RESPONSE_BODY_SENT,
			}
			trace.SendEvent(ctx, ev)
		}()
	}

	ourSS := &serverStream{
		stream: stream{
			ctx:    ctx,
			stream: ss,
		},
		ServerStream: ss,
	}

	return w.StreamHandler(srv, ourSS)
}
