package io.prometheus.metrics.benchmarks;

import io.prometheus.metrics.config.EscapingScheme;
import io.prometheus.metrics.expositionformats.ExpositionFormatWriter;
import io.prometheus.metrics.expositionformats.OpenMetricsTextFormatWriter;
import io.prometheus.metrics.expositionformats.PrometheusTextFormatWriter;
import io.prometheus.metrics.model.snapshots.GaugeSnapshot;
import io.prometheus.metrics.model.snapshots.GaugeSnapshot.GaugeDataPointSnapshot;
import io.prometheus.metrics.model.snapshots.Labels;
import io.prometheus.metrics.model.snapshots.MetricSnapshot;
import io.prometheus.metrics.model.snapshots.MetricSnapshots;
import io.prometheus.metrics.model.snapshots.SummarySnapshot;
import io.prometheus.metrics.model.snapshots.SummarySnapshot.SummaryDataPointSnapshot;
import io.prometheus.metrics.model.snapshots.Unit;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import org.openjdk.jmh.annotations.Benchmark;
import org.openjdk.jmh.annotations.Scope;
import org.openjdk.jmh.annotations.State;

/**
 * Results on a machine with dedicated Ubuntu 24.04 LTS, AMD Ryzen™ 9 7900 × 24, 96.0 GiB RAM:
 *
 * <pre>
 * Benchmark                                             Mode  Cnt       Score       Error  Units
 * TextFormatUtilBenchmark.openMetricsWriteToByteArray  thrpt   25  826847.708 ± 10941.611  ops/s
 * TextFormatUtilBenchmark.openMetricsWriteToNull       thrpt   25  847756.101 ±  5299.128  ops/s
 * TextFormatUtilBenchmark.prometheusWriteToByteArray   thrpt   25  874804.601 ±  9730.060  ops/s
 * TextFormatUtilBenchmark.prometheusWriteToNull        thrpt   25  910782.719 ± 17617.167  ops/s
 * </pre>
 */
public class TextFormatUtilBenchmark {

  private static final MetricSnapshots SNAPSHOTS;

  static {
    MetricSnapshot gaugeSnapshot =
        GaugeSnapshot.builder()
            .name("gauge_snapshot_name")
            .dataPoint(
                GaugeDataPointSnapshot.builder()
                    .labels(Labels.of("name", "value"))
                    .scrapeTimestampMillis(1000L)
                    .value(123.45d)
                    .build())
            .build();

    MetricSnapshot summaryDataPointSnapshot =
        SummarySnapshot.builder()
            .name("summary_snapshot_name_bytes")
            .dataPoint(
                SummaryDataPointSnapshot.builder()
                    .count(5)
                    .labels(Labels.of("name", "value"))
                    .sum(123456d)
                    .build())
            .unit(Unit.BYTES)
            .build();

    SNAPSHOTS = MetricSnapshots.of(gaugeSnapshot, summaryDataPointSnapshot);
  }

  private static final ExpositionFormatWriter OPEN_METRICS_TEXT_FORMAT_WRITER =
      OpenMetricsTextFormatWriter.create();
  private static final ExpositionFormatWriter PROMETHEUS_TEXT_FORMAT_WRITER =
      PrometheusTextFormatWriter.create();

  @State(Scope.Benchmark)
  public static class WriterState {

    final ByteArrayOutputStream byteArrayOutputStream;

    public WriterState() {
      this.byteArrayOutputStream = new ByteArrayOutputStream();
    }
  }

  @Benchmark
  public OutputStream openMetricsWriteToByteArray(WriterState writerState) throws IOException {
    // avoid growing the array
    ByteArrayOutputStream byteArrayOutputStream = writerState.byteArrayOutputStream;
    byteArrayOutputStream.reset();
    OPEN_METRICS_TEXT_FORMAT_WRITER.write(
        byteArrayOutputStream, SNAPSHOTS, EscapingScheme.ALLOW_UTF8);
    return byteArrayOutputStream;
  }

  @Benchmark
  public OutputStream openMetricsWriteToNull() throws IOException {
    OutputStream nullOutputStream = NullOutputStream.INSTANCE;
    OPEN_METRICS_TEXT_FORMAT_WRITER.write(nullOutputStream, SNAPSHOTS, EscapingScheme.ALLOW_UTF8);
    return nullOutputStream;
  }

  @Benchmark
  public OutputStream prometheusWriteToByteArray(WriterState writerState) throws IOException {
    // avoid growing the array
    ByteArrayOutputStream byteArrayOutputStream = writerState.byteArrayOutputStream;
    byteArrayOutputStream.reset();
    PROMETHEUS_TEXT_FORMAT_WRITER.write(
        byteArrayOutputStream, SNAPSHOTS, EscapingScheme.ALLOW_UTF8);
    return byteArrayOutputStream;
  }

  @Benchmark
  public OutputStream prometheusWriteToNull() throws IOException {
    OutputStream nullOutputStream = NullOutputStream.INSTANCE;
    PROMETHEUS_TEXT_FORMAT_WRITER.write(nullOutputStream, SNAPSHOTS, EscapingScheme.ALLOW_UTF8);
    return nullOutputStream;
  }

  static final class NullOutputStream extends OutputStream {

    static final OutputStream INSTANCE = new NullOutputStream();

    private NullOutputStream() {
      super();
    }

    @Override
    public void write(int b) {}

    @Override
    public void write(byte[] b) {}

    @Override
    public void write(byte[] b, int off, int len) {}

    @Override
    public void flush() {}

    @Override
    public void close() {}
  }
}
