// Copyright 2020 getensh.com. All rights reserved.
// Use of this source code is governed by getensh.com.

/**
 * @Author: mac
 * @Description: rpcx opentracing serverInterceptor
 * @Date: 2020/3/25 10:25
 */

package trace

import (
	"context"
	"github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"
	"github.com/opentracing/opentracing-go/log"
	"github.com/smallnest/rpcx/protocol"
	"github.com/smallnest/rpcx/server"
	"github.com/smallnest/rpcx/share"
	"github.com/uber/jaeger-client-go"
	"net"
)

type RpcxOpenTracingServerPlugin struct{}

func (p RpcxOpenTracingServerPlugin) Register(name string, rcvr interface{}, metadata string) error {
	span := opentracing.StartSpan(
		"rpcx.Register")
	defer span.Finish()

	span.LogFields(log.String("register_service", name))

	return nil
}

func (p RpcxOpenTracingServerPlugin) RegisterFunction(serviceName, fname string, fn interface{}, metadata string) error {
	span := opentracing.StartSpan(
		"rpcx.RegisterFunction")
	defer span.Finish()

	span.LogFields(log.String("register_function", serviceName+"."+fname))
	return nil
}

func (p RpcxOpenTracingServerPlugin) PostConnAccept(conn net.Conn) (net.Conn, bool) {
	span1 := opentracing.StartSpan(
		"rpcx.AcceptConn")
	defer span1.Finish()

	span1.LogFields(log.String("remote_addr", conn.RemoteAddr().String()))
	return conn, true
}

func (p RpcxOpenTracingServerPlugin) PreHandleRequest(ctx context.Context, r *protocol.Message) error {
	wireContext, err := share.GetSpanContextFromContext(ctx)
	if err != nil || wireContext == nil {
		return err
	}
	span := opentracing.StartSpan(
		"rpcx.service."+r.ServicePath+"."+r.ServiceMethod,
		ext.RPCServerOption(wireContext))

	clientConn := ctx.Value(server.RemoteConnContextKey).(net.Conn)
	span.LogFields(
		log.String("network", clientConn.RemoteAddr().Network()),
		log.String("remote_addr", clientConn.RemoteAddr().String()))

	if rpcxContext, ok := ctx.(*share.Context); ok {
		rpcxContext.SetValue(share.OpentracingSpanServerKey, span)
	}

	// 设置 request ID,可用于日志记录, 使用ctx.Value("RequestId")可以取出
	if sc, ok := span.Context().(jaeger.SpanContext); ok {
		traceId := sc.TraceID().String()
		ctx = context.WithValue(ctx, "RequestId", traceId)
	}
	return nil
}

func (p RpcxOpenTracingServerPlugin) PostWriteResponse(ctx context.Context, req *protocol.Message, res *protocol.Message, err error) error {
	if rpcxContext, ok := ctx.(*share.Context); ok {
		span := rpcxContext.Value(share.OpentracingSpanServerKey)
		if span != nil {
			sp := span.(opentracing.Span)
			sp.Finish()
		}
	}
	return nil
}