package bus

import (
	"context"

	"github.com/nats-io/nats.go"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/trace"
)

// natsHeaderCarrier 让 OTel 传播器能读写 NATS 消息头（W3C traceparent 等），
// 从而把链路上下文跨 NATS 总线带到下游服务——这是分布式追踪能连成一棵树的关键。
type natsHeaderCarrier nats.Header

func (c natsHeaderCarrier) Get(key string) string { return nats.Header(c).Get(key) }
func (c natsHeaderCarrier) Set(key, value string) { nats.Header(c).Set(key, value) }
func (c natsHeaderCarrier) Keys() []string {
	out := make([]string, 0, len(c))
	for k := range c {
		out = append(out, k)
	}
	return out
}

// injectTrace 把 ctx 里的链路上下文写进 NATS 消息头（发布前调用）。
func injectTrace(ctx context.Context, h nats.Header) {
	otel.GetTextMapPropagator().Inject(ctx, natsHeaderCarrier(h))
}

// extractTrace 从 NATS 消息头还原上游链路上下文，挂到 ctx 上（消费时调用）。
func extractTrace(ctx context.Context, h nats.Header) context.Context {
	if h == nil {
		return ctx
	}
	return otel.GetTextMapPropagator().Extract(ctx, natsHeaderCarrier(h))
}

// tracer 是 bus 自身的 span 发射器（与业务同 scope，方便 Jaeger 归组）。
func tracer() trace.Tracer {
	return otel.Tracer("github.com/sundynix/sundynix-agentix")
}