// Package amqptrace provides W3C TraceContext propagation through AMQP messages
// used by the platform's event bus. It wraps platform/rabbit with OTel context
// injection/extraction so producers and consumers share the same trace.
//
// Usage (producer):
//
//	import "git.k6n.net/mats/platform/amqptrace"
//
//	ev := event.Event{Type: event.OrderCreated, ...}
//	if err := amqptrace.Publish(ctx, conn, "order", ev); err != nil { ... }
//
// Usage (consumer):
//
//	amqptrace.Listen(ch, "order", "order.created", func(ctx context.Context, d amqp.Delivery) error {
//	    // ctx carries the parent span from the publisher.
//	    _, span := otel.Tracer("my-service").Start(ctx, "handle order")
//	    defer span.End()
//	    ...
//	})
package amqptrace

import (
	"context"
	"encoding/json"
	"time"

	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/trace"

	"git.k6n.net/mats/platform/event"
	"git.k6n.net/mats/platform/rabbit"
	amqp "github.com/rabbitmq/amqp091-go"
)

// ---------------------------------------------------------------------------
// Carrier — adapts amqp.Table to the OTel TextMapCarrier
// ---------------------------------------------------------------------------

type amqpHeadersCarrier struct{ headers amqp.Table }

func (c amqpHeadersCarrier) Get(key string) string {
	if v, ok := c.headers[key]; ok {
		if s, ok := v.(string); ok {
			return s
		}
	}
	return ""
}

func (c amqpHeadersCarrier) Set(key, value string) {
	c.headers[key] = value
}

func (c amqpHeadersCarrier) Keys() []string {
	keys := make([]string, 0, len(c.headers))
	for k := range c.headers {
		keys = append(keys, k)
	}
	return keys
}

// ---------------------------------------------------------------------------
// Producer
// ---------------------------------------------------------------------------

// Publish marshals the event and publishes it to exchange with the W3C trace
// context from ctx injected into AMQP headers. Consumers using Listen or
// ListenAndHandle will automatically create child spans.
func Publish(ctx context.Context, c *amqp.Connection, exchange string, ev event.Event) error {
	body, err := json.Marshal(ev)
	if err != nil {
		return err
	}
	headers := make(amqp.Table)
	otel.GetTextMapPropagator().Inject(ctx, amqpHeadersCarrier{headers})

	ch, err := c.Channel()
	if err != nil {
		return err
	}
	defer ch.Close()
	return ch.Publish(exchange, string(ev.Type), true, false, amqp.Publishing{
		ContentType:  "application/json",
		Body:         body,
		DeliveryMode: amqp.Persistent,
		Timestamp:    time.Now(),
		Headers:      headers,
	})
}

// PublishEvent is a convenience wrapper around Publish that takes a
// context.Context in its signature, matching the naming convention of
// platform/rabbit.PublishEvent. It is equivalent to calling Publish.
func PublishEvent(ctx context.Context, c *amqp.Connection, exchange string, ev event.Event) error {
	return Publish(ctx, c, exchange, ev)
}

// ---------------------------------------------------------------------------
// Consumer
// ---------------------------------------------------------------------------

// tracedHandler wraps a delivery handler to extract the W3C trace context
// from AMQP headers and start a child span named after the routing key.
// The handler receives a context that carries the parent trace.
func tracedHandler(pattern string, handler func(context.Context, amqp.Delivery) error) func(amqp.Delivery) error {
	return func(d amqp.Delivery) error {
		ctx := otel.GetTextMapPropagator().Extract(context.Background(), amqpHeadersCarrier{d.Headers})

		tracer := otel.Tracer("amqptrace.consumer")
		ctx, span := tracer.Start(ctx, d.RoutingKey,
			trace.WithAttributes(
				attribute.String("messaging.system", "rabbitmq"),
				attribute.String("messaging.destination", d.Exchange),
				attribute.String("messaging.destination_kind", "topic"),
				attribute.String("messaging.rabbitmq.routing_key", d.RoutingKey),
				attribute.Int("messaging.rabbitmq.delivery_tag", int(d.DeliveryTag)),
			),
		)
		defer span.End()

		err := handler(ctx, d)
		if err != nil {
			span.RecordError(err)
			span.SetStatus(codes.Error, err.Error())
		}
		return err
	}
}

// Listen binds an exclusive queue with a routing-key pattern, wraps the
// handler to extract W3C trace context from AMQP headers, and starts a
// child span for each message. The handler type accepts a context.Context
// (with the parent span) and an amqp.Delivery.
//
// Delegates to rabbit.ListenToPattern for topology setup (exchange declare,
// queue bind, consume goroutine) — wraps only the trace-propagation layer.
func Listen(ch *amqp.Channel, exchange, pattern string, handler func(context.Context, amqp.Delivery) error) error {
	wrapped := tracedHandler(pattern, handler)
	return rabbit.ListenToPattern(ch, exchange, pattern, wrapped)
}