package metrics

import (
	"context"
	"strconv"
	"time"

	"github.com/ThreeDotsLabs/watermill/message"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/tidbcloud/tiinf/transport/tieventbus/metadata"
)

var (
	publisherLabelKeys = []string{
		LabelKeyEventType,
		LabelKeySource,
		LabelKeyTarget,
		LabelKeySuccess,
	}
)

// PublisherPrometheusMetricsDecorator decorates a publisher to capture Prometheus metrics.
type PublisherPrometheusMetricsDecorator struct {
	pub                  message.Publisher
	publisherName        string
	publishTimeSeconds   *prometheus.HistogramVec
	eventsPublishedTotal *prometheus.CounterVec
}

// Publish updates the relevant publisher metrics and calls the wrapped publisher's Publish.
func (m PublisherPrometheusMetricsDecorator) Publish(topic string, messages ...*message.Message) (err error) {
	if len(messages) == 0 {
		return m.pub.Publish(topic)
	}

	ctxArr := make([]context.Context, len(messages))
	for i, msg := range messages {
		ctxArr[i] = msg.Context()
		msg.SetContext(setPublishObservedToCtx(msg.Context()))
		if _, ok := msg.Metadata[metadata.CreatedAtKey]; !ok {
			msg.Metadata.Set(metadata.CreatedAtKey, strconv.FormatInt(time.Now().UTC().UnixNano(), 10))
		}
	}
	start := time.Now()
	defer func() {
		duration := time.Since(start).Seconds()
		for i, msg := range messages {
			if publishAlreadyObserved(ctxArr[i]) {
				// decorator idempotency when applied decorator multiple times
				continue
			}
			ctx := msg.Context()
			labels := labelsFromCtx(ctx, publisherLabelKeys...)
			if err != nil {
				labels[LabelKeySuccess] = "false"
			} else {
				labels[LabelKeySuccess] = "true"
			}
			m.eventsPublishedTotal.With(labels).Inc()
			m.publishTimeSeconds.With(labels).Observe(duration)
		}
	}()

	return m.pub.Publish(topic, messages...)
}

// Close decreases the total publisher count, closes the Prometheus HTTP server and calls wrapped Close.
func (m PublisherPrometheusMetricsDecorator) Close() error {
	return m.pub.Close()
}