package io.prometheus.metrics.core.exemplars;

import static java.util.Objects.requireNonNull;

import io.prometheus.metrics.annotations.StableApi;
import io.prometheus.metrics.core.util.Scheduler;
import io.prometheus.metrics.model.snapshots.Exemplar;
import io.prometheus.metrics.model.snapshots.Exemplars;
import io.prometheus.metrics.model.snapshots.Labels;
import io.prometheus.metrics.tracer.common.SpanContext;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.LongSupplier;
import java.util.function.Supplier;
import javax.annotation.Nullable;

/**
 * The ExemplarSampler selects Spans as exemplars.
 *
 * <p>There are two types of Exemplars: Regular exemplars are sampled implicitly if a supported
 * tracing library is detected. Custom exemplars are provided explicitly in code, for example if a
 * developer wants to make sure an Exemplar is created for a specific code path.
 *
 * <p>Spans will be marked as being an Exemplar by calling {@link
 * SpanContext#markCurrentSpanAsExemplar()}. The tracer implementation should set a Span attribute
 * to mark the current Span as an Exemplar. This attribute can be used by a trace sampling algorithm
 * to make sure traces with Exemplars are sampled.
 *
 * <p>The ExemplarSample is rate-limited, so only a small fraction of Spans will be marked as
 * Exemplars in an application with a large number of requests.
 *
 * <p>See {@link ExemplarSamplerConfig} for configuration options.
 */
@StableApi
public class ExemplarSampler {

  @SuppressWarnings("ReferenceEquality")
  private static boolean sameObject(Object left, Object right) {
    return left == right;
  }

  private final ExemplarSamplerConfig config;
  private final Exemplar[] exemplars;
  private final Exemplar[]
      customExemplars; // Separate from exemplars, because we don't want custom exemplars
  // to be overwritten by automatic exemplar sampling. exemplars.length == customExemplars.length
  private final AtomicBoolean acceptingNewExemplars = new AtomicBoolean(true);
  private final AtomicBoolean acceptingNewCustomExemplars = new AtomicBoolean(true);

  @Nullable
  private final SpanContext
      spanContext; // may be null, in that case SpanContextSupplier.getSpanContext() is used.

  @Nullable private final Supplier<Labels> additionalLabelsSupplier;

  public ExemplarSampler(ExemplarSamplerConfig config) {
    this(config, null, null);
  }

  /**
   * Constructor with an additional {code spanContext} argument. This is useful for testing, but may
   * also be useful in some production scenarios. If {@code spanContext != null} that spanContext is
   * used and {@link io.prometheus.metrics.tracer.initializer.SpanContextSupplier
   * SpanContextSupplier} is not used. If {@code spanContext == null} {@link
   * io.prometheus.metrics.tracer.initializer.SpanContextSupplier#getSpanContext()
   * SpanContextSupplier.getSpanContext()} is called to find a span context.
   */
  public ExemplarSampler(ExemplarSamplerConfig config, @Nullable SpanContext spanContext) {
    this(config, spanContext, null);
  }

  /**
   * Constructor that additionally accepts a supplier of labels to be merged into every
   * automatically-sampled exemplar. The supplier is called each time an exemplar is sampled from a
   * span context, so it can return dynamic values (e.g. a request-scoped identifier). The supplier
   * is only called when a valid, sampled span context is present.
   */
  public ExemplarSampler(
      ExemplarSamplerConfig config,
      @Nullable SpanContext spanContext,
      @Nullable Supplier<Labels> additionalLabelsSupplier) {
    this.config = config;
    this.exemplars = new Exemplar[config.getNumberOfExemplars()];
    this.customExemplars = new Exemplar[exemplars.length];
    this.spanContext = spanContext;
    this.additionalLabelsSupplier = additionalLabelsSupplier;
  }

  public Exemplars collect() {
    // this may run in parallel with observe()
    long now = System.currentTimeMillis();
    List<Exemplar> result = new ArrayList<>(exemplars.length);
    for (int i = 0; i < customExemplars.length; i++) {
      Exemplar exemplar = customExemplars[i];
      if (exemplar != null) {
        if (now - exemplar.getTimestampMillis() > config.getMaxRetentionPeriodMillis()) {
          customExemplars[i] = null;
        } else {
          result.add(exemplar);
        }
      }
    }
    for (int i = 0; i < exemplars.length && result.size() < exemplars.length; i++) {
      Exemplar exemplar = exemplars[i];
      if (exemplar != null) {
        if (now - exemplar.getTimestampMillis() > config.getMaxRetentionPeriodMillis()) {
          exemplars[i] = null;
        } else {
          result.add(exemplar);
        }
      }
    }
    return Exemplars.of(result);
  }

  public void reset() {
    for (int i = 0; i < exemplars.length; i++) {
      exemplars[i] = null;
      customExemplars[i] = null;
    }
  }

  public void observe(double value) {
    if (!acceptingNewExemplars.get()) {
      return; // This is the hot path in a high-throughput application and should be as efficient as
      // possible.
    }
    rateLimitedObserve(acceptingNewExemplars, value, () -> doObserve(value));
  }

  public void observeWithExemplar(double value, Labels labels) {
    if (!acceptingNewCustomExemplars.get()) {
      return; // This is the hot path in a high-throughput application and should be as efficient as
      // possible.
    }
    rateLimitedObserve(
        acceptingNewCustomExemplars, value, () -> doObserveWithExemplar(value, labels));
  }

  private long doObserve(double value) {
    if (exemplars.length == 1) {
      return doObserveSingleExemplar(value);
    } else {
      double[] classicUpperBounds = config.getHistogramClassicUpperBounds();
      if (classicUpperBounds != null) {
        return doObserveWithUpperBounds(value, classicUpperBounds);
      } else {
        return doObserveWithoutUpperBounds(value);
      }
    }
  }

  private long doObserveSingleExemplar(double value) {
    long now = System.currentTimeMillis();
    Exemplar current = exemplars[0];
    if (current == null
        || now - current.getTimestampMillis() > config.getMinRetentionPeriodMillis()) {
      return updateExemplar(0, value, now);
    }
    return 0;
  }

  private long doObserveSingleExemplar(double amount, Labels labels) {
    long now = System.currentTimeMillis();
    Exemplar current = customExemplars[0];
    if (current == null
        || now - current.getTimestampMillis() > config.getMinRetentionPeriodMillis()) {
      return updateCustomExemplar(0, amount, labels, now);
    }
    return 0;
  }

  private long doObserveWithUpperBounds(double value, double[] classicUpperBounds) {
    long now = System.currentTimeMillis();
    for (int i = 0; i < classicUpperBounds.length; i++) {
      if (value <= classicUpperBounds[i]) {
        Exemplar previous = exemplars[i];
        if (previous == null
            || now - previous.getTimestampMillis() > config.getMinRetentionPeriodMillis()) {
          return updateExemplar(i, value, now);
        } else {
          return 0;
        }
      }
    }
    return 0; // will never happen, as upperBounds contains +Inf
  }

  private long doObserveWithoutUpperBounds(double value) {
    final long now = System.currentTimeMillis();
    Exemplar smallest = null;
    int smallestIndex = -1;
    Exemplar largest = null;
    int largestIndex = -1;
    int nullIndex = -1;
    for (int i = exemplars.length - 1; i >= 0; i--) {
      Exemplar exemplar = exemplars[i];
      if (exemplar == null) {
        nullIndex = i;
      } else if (now - exemplar.getTimestampMillis() > config.getMaxRetentionPeriodMillis()) {
        exemplars[i] = null;
        nullIndex = i;
      } else {
        if (smallest == null || exemplar.getValue() < smallest.getValue()) {
          smallest = exemplar;
          smallestIndex = i;
        }
        if (largest == null || exemplar.getValue() > largest.getValue()) {
          largest = exemplar;
          largestIndex = i;
        }
      }
    }
    if (nullIndex >= 0) {
      return updateExemplar(nullIndex, value, now);
    }
    if (now - requireNonNull(smallest).getTimestampMillis() > config.getMinRetentionPeriodMillis()
        && value < smallest.getValue()) {
      return updateExemplar(smallestIndex, value, now);
    }
    if (now - requireNonNull(largest).getTimestampMillis() > config.getMinRetentionPeriodMillis()
        && value > largest.getValue()) {
      return updateExemplar(largestIndex, value, now);
    }
    long oldestTimestamp = 0;
    int oldestIndex = -1;
    for (int i = 0; i < exemplars.length; i++) {
      Exemplar exemplar = exemplars[i];
      if (exemplar != null && !sameObject(exemplar, smallest) && !sameObject(exemplar, largest)) {
        if (oldestTimestamp == 0 || exemplar.getTimestampMillis() < oldestTimestamp) {
          oldestTimestamp = exemplar.getTimestampMillis();
          oldestIndex = i;
        }
      }
    }
    if (oldestIndex != -1 && now - oldestTimestamp > config.getMinRetentionPeriodMillis()) {
      return updateExemplar(oldestIndex, value, now);
    }
    return 0;
  }

  // Returns the timestamp of the newly added Exemplar (which is System.currentTimeMillis())
  // or 0 if no Exemplar was added.
  private long doObserveWithExemplar(double amount, Labels labels) {
    if (customExemplars.length == 1) {
      return doObserveSingleExemplar(amount, labels);
    } else {
      double[] classicUpperBounds = config.getHistogramClassicUpperBounds();
      if (classicUpperBounds != null) {
        return doObserveWithExemplarWithUpperBounds(amount, labels, classicUpperBounds);
      } else {
        return doObserveWithExemplarWithoutUpperBounds(amount, labels);
      }
    }
  }

  private long doObserveWithExemplarWithUpperBounds(
      double value, Labels labels, double[] classicUpperBounds) {
    long now = System.currentTimeMillis();
    for (int i = 0; i < classicUpperBounds.length; i++) {
      if (value <= classicUpperBounds[i]) {
        Exemplar previous = customExemplars[i];
        if (previous == null
            || now - previous.getTimestampMillis() > config.getMinRetentionPeriodMillis()) {
          return updateCustomExemplar(i, value, labels, now);
        } else {
          return 0;
        }
      }
    }
    return 0; // will never happen, as upperBounds contains +Inf
  }

  private long doObserveWithExemplarWithoutUpperBounds(double amount, Labels labels) {
    final long now = System.currentTimeMillis();
    int nullPos = -1;
    int oldestPos = -1;
    Exemplar oldest = null;
    for (int i = customExemplars.length - 1; i >= 0; i--) {
      Exemplar exemplar = customExemplars[i];
      if (exemplar == null) {
        nullPos = i;
      } else if (now - exemplar.getTimestampMillis() > config.getMaxRetentionPeriodMillis()) {
        customExemplars[i] = null;
        nullPos = i;
      } else {
        if (oldest == null || exemplar.getTimestampMillis() < oldest.getTimestampMillis()) {
          oldest = exemplar;
          oldestPos = i;
        }
      }
    }
    if (nullPos != -1) {
      return updateCustomExemplar(nullPos, amount, labels, now);
    } else if (now - requireNonNull(oldest).getTimestampMillis()
        > config.getMinRetentionPeriodMillis()) {
      return updateCustomExemplar(oldestPos, amount, labels, now);
    } else {
      return 0;
    }
  }

  /**
   * Observing requires a system call to {@link System#currentTimeMillis()}, and it requires
   * iterating over the existing exemplars to check if one of the existing exemplars can be
   * replaced.
   *
   * <p>To avoid performance issues, we rate limit observing exemplars to {@link
   * ExemplarSamplerConfig#getSampleIntervalMillis()} milliseconds.
   */
  @SuppressWarnings("FutureReturnValueIgnored")
  private void rateLimitedObserve(AtomicBoolean accepting, double value, LongSupplier observeFunc) {
    if (Double.isNaN(value)) {
      return;
    }
    if (!accepting.compareAndSet(true, false)) {
      return;
    }
    // observeFunc returns the current timestamp or 0 if no Exemplar was added.
    long now = observeFunc.getAsLong();
    long sleepTime =
        now == 0 ? config.getSampleIntervalMillis() : durationUntilNextExemplarExpires(now);
    Scheduler.schedule(
        () -> accepting.compareAndSet(false, true), sleepTime, TimeUnit.MILLISECONDS);
  }

  private long durationUntilNextExemplarExpires(long now) {
    long oldestTimestamp = now;
    for (Exemplar exemplar : exemplars) {
      if (exemplar == null) {
        return config.getSampleIntervalMillis();
      } else if (exemplar.getTimestampMillis() < oldestTimestamp) {
        oldestTimestamp = exemplar.getTimestampMillis();
      }
    }
    long oldestAge = now - oldestTimestamp;
    if (oldestAge < config.getMinRetentionPeriodMillis()) {
      return config.getMinRetentionPeriodMillis() - oldestAge;
    }
    return config.getSampleIntervalMillis();
  }

  private long updateCustomExemplar(int index, double value, Labels labels, long now) {
    if (!labels.contains(Exemplar.TRACE_ID) && !labels.contains(Exemplar.SPAN_ID)) {
      labels = mergeLabels(labels, sampleTraceContextLabels());
    }
    customExemplars[index] =
        Exemplar.builder().value(value).labels(labels).timestampMillis(now).build();
    return now;
  }

  private long updateExemplar(int index, double value, long now) {
    Labels traceLabels = doSampleExemplar();
    if (!traceLabels.isEmpty()) {
      exemplars[index] =
          Exemplar.builder().value(value).labels(traceLabels).timestampMillis(now).build();
      return now;
    } else {
      return 0;
    }
  }

  private Labels doSampleExemplar() {
    Labels labels = sampleTraceContextLabels();
    if (labels.isEmpty()) {
      return labels;
    }
    // Per-metric supplier first (more specific), then the global supplier. On a name
    // collision the earlier (more specific) value is kept; the reserved trace_id/span_id
    // labels always win over both.
    labels = mergeAdditionalLabels(labels, additionalLabelsSupplier);
    labels = mergeAdditionalLabels(labels, ExemplarLabelsSupplier.getExemplarLabelsSupplier());
    return labels;
  }

  private Labels sampleTraceContextLabels() {
    // Using the qualified name so that Micrometer can exclude the dependency on
    // prometheus-metrics-tracer-initializer
    // as they provide their own implementation of SpanContextSupplier.
    // If we had an import statement for SpanContextSupplier the dependency would be needed in any
    // case.
    SpanContext spanContext =
        this.spanContext != null
            ? this.spanContext
            : io.prometheus.metrics.tracer.initializer.SpanContextSupplier.getSpanContext();
    try {
      if (spanContext != null) {
        if (spanContext.isCurrentSpanSampled()) {
          String spanId = spanContext.getCurrentSpanId();
          String traceId = spanContext.getCurrentTraceId();
          if (spanId != null && traceId != null) {
            spanContext.markCurrentSpanAsExemplar();
            return Labels.of(Exemplar.TRACE_ID, traceId, Exemplar.SPAN_ID, spanId);
          }
        }
      }
    } catch (NoClassDefFoundError ignored) {
      // ignore
    }
    return Labels.EMPTY;
  }

  /**
   * Merge labels from {@code supplier} into {@code base}, dropping any label whose name already
   * exists in {@code base}. Never throws: a {@code null} supplier, a {@code null}/empty result, a
   * colliding label name, or an exception thrown by the supplier all result in {@code base} being
   * returned unchanged (minus the offending labels). A misbehaving supplier must never break metric
   * collection.
   */
  private static Labels mergeAdditionalLabels(Labels base, @Nullable Supplier<Labels> supplier) {
    if (supplier == null) {
      return base;
    }
    Labels extra;
    try {
      extra = supplier.get();
    } catch (Throwable ignored) {
      // A misbehaving supplier (any RuntimeException or Error) must never break metric collection.
      return base;
    }
    if (extra == null || extra.isEmpty()) {
      return base;
    }
    return mergeLabels(base, extra);
  }

  /**
   * Merge {@code extra} into {@code base}, dropping any label whose name already exists in {@code
   * base}.
   */
  private static Labels mergeLabels(Labels base, Labels extra) {
    if (extra.isEmpty()) {
      return base;
    }
    // Count name collisions with base in a single pass so we can merge exactly once below: base
    // (trace_id/span_id and any more-specific supplier) always wins, so colliding labels are
    // dropped. extra is itself a valid Labels (no internal duplicates), so the surviving labels
    // never collide with each other and merge() cannot throw on a duplicate name.
    int size = extra.size();
    int collisions = 0;
    for (int i = 0; i < size; i++) {
      if (base.contains(extra.getName(i))) {
        collisions++;
      }
    }
    if (collisions == 0) {
      return base.merge(extra);
    }
    if (collisions == size) {
      return base;
    }
    int kept = size - collisions;
    String[] names = new String[kept];
    String[] values = new String[kept];
    int j = 0;
    for (int i = 0; i < size; i++) {
      String name = extra.getName(i);
      if (!base.contains(name)) {
        names[j] = name;
        values[j] = extra.getValue(i);
        j++;
      }
    }
    return base.merge(names, values);
  }
}