package tracing

import (
	"context"

	"go.opentelemetry.io/otel"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

type annotationTextMapCarrier map[string]string

func (c annotationTextMapCarrier) Get(key string) string {
	return c[key]
}
func (c annotationTextMapCarrier) Set(key, value string) {
	c[key] = value
}
func (c annotationTextMapCarrier) Keys() []string {
	keys := make([]string, 0, len(c))
	for k := range c {
		keys = append(keys, k)
	}
	return keys
}

func getAnnotations(obj metav1.Object) map[string]string {
	annotations := obj.GetAnnotations()
	if annotations == nil {
		annotations = make(map[string]string)
	}

	return annotations
}

func ExtractTraceContext(ctx context.Context, obj metav1.Object) context.Context {
	return otel.GetTextMapPropagator().Extract(ctx, annotationTextMapCarrier(getAnnotations(obj)))
}

func InjectTraceContext(ctx context.Context, obj metav1.Object) {
	annotations := getAnnotations(obj)
	otel.GetTextMapPropagator().Inject(ctx, annotationTextMapCarrier(annotations))
	obj.SetAnnotations(annotations)
}