package jaeger

import (
	"context"
	"fmt"
	share2 "github.com/rpcxio/rpcx-plugins/share"
	"github.com/smallnest/rpcx/share"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/exporters/jaeger"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/sdk/resource"
	"go.opentelemetry.io/otel/sdk/trace"
	"go.opentelemetry.io/otel/semconv/v1.12.0"
	trace2 "go.opentelemetry.io/otel/trace"
)

const (
	CtxSpanName    = "span"
	CtxSpanCtxName = "spanCtx"
)

var TracerProvider *trace.TracerProvider

// InitJaeger @Title 初始化jaeger
func InitJaeger(endpoint, name string) error {
	exporter, err := jaeger.New(jaeger.WithCollectorEndpoint(jaeger.WithEndpoint(endpoint)))
	if err != nil {
		return err
	}
	TracerProvider = trace.NewTracerProvider(
		// Always be sure to batch in production.
		trace.WithBatcher(exporter),
		// Record information about this application in an Resource.
		trace.WithResource(resource.NewWithAttributes(
			semconv.SchemaURL,
			semconv.ServiceNameKey.String(name),
		)),
	)

	otel.SetTextMapPropagator(propagation.NewCompositeTextMapPropagator(propagation.TraceContext{}, propagation.Baggage{}))
	otel.SetTracerProvider(TracerProvider)
	return nil
}

// BeforeHandel @Title 开始Rpc调用
func BeforeHandel(ctx context.Context, baseName, serviceName, serviceMethod string, args interface{}, reply interface{}) (context.Context, error) {
	if span, ok := ctx.Value(CtxSpanName).(trace2.Span); ok {
		span.AddEvent(fmt.Sprintf("调用rpc服务:%s,%s,%s", baseName, serviceName, serviceMethod))
		spanCtx := ctx.Value(CtxSpanCtxName).(context.Context)
		meta := ctx.Value(share.ReqMetaDataKey)
		spanCtx = context.WithValue(spanCtx, share.ReqMetaDataKey, meta)
		share2.Inject(spanCtx, otel.GetTextMapPropagator())
		ctx = spanCtx
	}
	return ctx, nil
}

// AfterHandel @Title 结束Rpc调用
func AfterHandel(ctx context.Context, baseName, serviceName, serviceMethod string, args interface{}, reply interface{}, err error) error {
	if span, ok := ctx.Value(CtxSpanName).(trace2.Span); ok {
		span.AddEvent(fmt.Sprintf("调用rpc服务完成:%s,%s,%s", baseName, serviceName, serviceMethod))
	}
	return err
}