package api

import (
	"context"
	"go.opentelemetry.io/otel/propagation"
)

const traceParentHeader = "traceparent"
const traceStateHeader = "tracestate"

// TraceParentAndStateFromContext returns W3C conform trace parent and state from context
func TraceParentAndStateFromContext(ctx context.Context) (string, string) {
	mapCarrier := propagation.MapCarrier{}
	propagator := propagation.TraceContext{}
	propagator.Inject(ctx, mapCarrier)

	// Get trace parent from context w3c conform format
	// Format: <version>-<trace-id>-<parent-id>-<trace-flags>
	// Example:	"00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
	traceParent := mapCarrier[traceParentHeader]
	traceState := mapCarrier[traceStateHeader]

	return traceParent, traceState
}

// AddTraceParentAndStateToContext adds trace and state related information to the given context.
func AddTraceParentAndStateToContext(ctx context.Context, traceParent string, traceState string) context.Context {
	mapCarrier := propagation.MapCarrier{}
	mapCarrier[traceParentHeader] = traceParent
	mapCarrier[traceStateHeader] = traceState

	propagator := propagation.TraceContext{}
	return propagator.Extract(ctx, mapCarrier)
}