#!/usr/bin/env python3
"""
Benchmark promrs vs prometheus_client (Python) with ~9k metrics.

Benchmarks:
  - record: inc/set/observe on pre-created metric children
  - scrape: generate_latest() encoding all metrics to text
"""
import pyperf

import prometheus_client as pc_py
import promrs as pc_rs

LABEL_COMBINATIONS = 3000
BUCKETS = (0.01, 0.05, 0.1, 0.5, 1.0, 5.0)

# --- Rust setup ---
rs_registry = pc_rs.CollectorRegistry()
rs_counter = pc_rs.Counter("bench_counter", "Benchmark counter", ["id"], registry=rs_registry)
rs_gauge = pc_rs.Gauge("bench_gauge", "Benchmark gauge", ["id"], registry=rs_registry)
rs_histogram = pc_rs.Histogram(
    "bench_histogram", "Benchmark histogram", ["id"], registry=rs_registry, buckets=list(BUCKETS),
)

rs_counter_children = [rs_counter.labels(id=str(i)) for i in range(LABEL_COMBINATIONS)]
rs_gauge_children = [rs_gauge.labels(id=str(i)) for i in range(LABEL_COMBINATIONS)]
rs_histogram_children = [rs_histogram.labels(id=str(i)) for i in range(LABEL_COMBINATIONS)]

for c in rs_counter_children:
    c.inc()
for g in rs_gauge_children:
    g.set(0.0)
for h in rs_histogram_children:
    h.observe(0.0)

# --- Python setup ---
py_registry = pc_py.CollectorRegistry(auto_describe=True)
py_counter = pc_py.Counter("bench_counter", "Benchmark counter", ["id"], registry=py_registry)
py_gauge = pc_py.Gauge("bench_gauge", "Benchmark gauge", ["id"], registry=py_registry)
py_histogram = pc_py.Histogram(
    "bench_histogram", "Benchmark histogram", ["id"], registry=py_registry, buckets=BUCKETS,
)

py_counter_children = [py_counter.labels(id=str(i)) for i in range(LABEL_COMBINATIONS)]
py_gauge_children = [py_gauge.labels(id=str(i)) for i in range(LABEL_COMBINATIONS)]
py_histogram_children = [py_histogram.labels(id=str(i)) for i in range(LABEL_COMBINATIONS)]

for c in py_counter_children:
    c.inc()
for g in py_gauge_children:
    g.set(0.0)
for h in py_histogram_children:
    h.observe(0.0)


# --- Rust benchmarks ---
def rs_record_counter():
    for c in rs_counter_children:
        c.inc()


def rs_record_gauge():
    for g in rs_gauge_children:
        g.set(42.0)


def rs_record_histogram():
    for h in rs_histogram_children:
        h.observe(0.042)


def rs_scrape():
    pc_rs.generate_latest(rs_registry)


# --- Python benchmarks ---
def py_record_counter():
    for c in py_counter_children:
        c.inc()


def py_record_gauge():
    for g in py_gauge_children:
        g.set(42.0)


def py_record_histogram():
    for h in py_histogram_children:
        h.observe(0.042)


def py_scrape():
    pc_py.generate_latest(py_registry)


if __name__ == "__main__":
    runner = pyperf.Runner()
    # Counter
    runner.bench_func("rs_record_counter_9k", rs_record_counter)
    runner.bench_func("py_record_counter_9k", py_record_counter)
    # Gauge
    runner.bench_func("rs_record_gauge_9k", rs_record_gauge)
    runner.bench_func("py_record_gauge_9k", py_record_gauge)
    # Histogram
    runner.bench_func("rs_record_histogram_9k", rs_record_histogram)
    runner.bench_func("py_record_histogram_9k", py_record_histogram)
    # Scrape
    runner.bench_func("rs_scrape_9k", rs_scrape)
    runner.bench_func("py_scrape_9k", py_scrape)
