package tracing

import (
	"context"
	"net/http"
	"strings"
	"sync"

	"github.com/gin-gonic/gin"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/trace"
)

const TraceIDHeader = "X-Trace-ID"

var instrumentDefaultHTTPTransportOnce sync.Once

func Configure() {
	otel.SetTextMapPropagator(propagation.NewCompositeTextMapPropagator(
		traceIDPropagator{},
		propagation.TraceContext{},
		propagation.Baggage{},
	))
	instrumentDefaultHTTPTransportOnce.Do(func() {
		http.DefaultTransport = propagationTransport{base: http.DefaultTransport}
	})
}

func Middleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		ctx := otel.GetTextMapPropagator().Extract(
			c.Request.Context(),
			propagation.HeaderCarrier(c.Request.Header),
		)
		c.Request = c.Request.WithContext(ctx)
		if spanContext := trace.SpanContextFromContext(ctx); spanContext.HasTraceID() {
			c.Header(TraceIDHeader, spanContext.TraceID().String())
		}
		c.Next()
	}
}

func TraceIDFromContext(ctx context.Context) string {
	spanContext := trace.SpanContextFromContext(ctx)
	if !spanContext.HasTraceID() {
		return ""
	}
	return spanContext.TraceID().String()
}

type propagationTransport struct {
	base http.RoundTripper
}

func (t propagationTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	req = req.Clone(req.Context())
	req.Header = req.Header.Clone()
	otel.GetTextMapPropagator().Inject(req.Context(), propagation.HeaderCarrier(req.Header))
	return t.base.RoundTrip(req)
}

type traceIDPropagator struct{}

func (traceIDPropagator) Inject(ctx context.Context, carrier propagation.TextMapCarrier) {
	spanContext := trace.SpanContextFromContext(ctx)
	if spanContext.HasTraceID() {
		carrier.Set(TraceIDHeader, spanContext.TraceID().String())
	}
}

func (traceIDPropagator) Extract(ctx context.Context, carrier propagation.TextMapCarrier) context.Context {
	traceID, err := trace.TraceIDFromHex(strings.TrimSpace(carrier.Get(TraceIDHeader)))
	if err != nil || !traceID.IsValid() {
		return ctx
	}
	spanContext := trace.NewSpanContext(trace.SpanContextConfig{
		TraceID:    traceID,
		SpanID:     trace.SpanID{1},
		TraceFlags: trace.FlagsSampled,
		Remote:     true,
	})
	return trace.ContextWithRemoteSpanContext(ctx, spanContext)
}

func (traceIDPropagator) Fields() []string {
	return []string{TraceIDHeader}
}