# /// script
# requires-python = ">=3.10"
# dependencies = ["prometheus_client"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Count your series.

A metrics endpoint built with prometheus_client: a request counter labelled by
method, status and route, a latency histogram labelled by route, and a gauge.
Counts the series per metric against the product of label cardinalities, then
adds one label with 10,000 values (a user id) to the request counter and
measures the scrape size and the memory before and after.
"""
import argparse, json, random, tracemalloc
from prometheus_client import CollectorRegistry, Counter, Gauge, Histogram, generate_latest

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--users", type=int, default=10_000, help="Values of the added label.")
p.add_argument("--requests", type=int, default=200_000, help="Requests to record.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

METHODS, STATUSES, ROUTES = ["GET", "POST"], ["200", "404", "500"], [f"/r{i}" for i in range(20)]
def build(with_user):
    tracemalloc.start(); reg = CollectorRegistry()
    labels = ["method", "status", "route"] + (["user"] if with_user else [])
    req = Counter("http_requests", "Requests.", labels, registry=reg)
    lat = Histogram("http_latency_seconds", "Latency.", ["route"], registry=reg)
    Gauge("queue_length", "Queue length.", registry=reg).set(3)
    random.seed(0)
    for i in range(args.requests):
        m, s, r = random.choice(METHODS), random.choice(STATUSES), random.choice(ROUTES)
        req.labels(m, s, r, *( [f"u{random.randrange(args.users)}"] if with_user else [])).inc()
        lat.labels(r).observe(random.random())
    body = generate_latest(reg); mem = tracemalloc.get_traced_memory()[0]; tracemalloc.stop()
    series = {}
    for line in body.decode().splitlines():
        if line and not line.startswith("#"):
            name = line.split("{")[0].split(" ")[0]
            for suf in ("_total", "_created", "_bucket", "_count", "_sum"):
                if name.endswith(suf): name = name[: -len(suf)]; break
            series[name] = series.get(name, 0) + 1
    return series, len(body), mem

before, size0, mem0 = build(False)
after, size1, mem1 = build(True)
expected = {"http_requests": len(METHODS) * len(STATUSES) * len(ROUTES), "http_latency_seconds": len(ROUTES), "queue_length": 1}
rows = [{"metric": m, "label_product": expected[m], "series_before": before.get(m, 0), "series_after": after.get(m, 0)} for m in expected]
summary = {"size_before_kb": size0 / 1024, "size_after_kb": size1 / 1024, "mem_before_mb": mem0 / 2**20, "mem_after_mb": mem1 / 2**20,
           "size_x": size1 / size0, "mem_x": mem1 / mem0, "series_before": sum(before.values()), "series_after": sum(after.values()),
           "series_x": sum(after.values()) / sum(before.values())}
for r in rows: print(f"{r['metric']:<22} label product {r['label_product']:>4}  exposed lines before {r['series_before']:>6}  after adding user {r['series_after']:>7}")
print(f"scrape {summary['size_before_kb']:.0f} KB -> {summary['size_after_kb']:.0f} KB ({summary['size_x']:.0f}x); "
      f"memory {summary['mem_before_mb']:.1f} MB -> {summary['mem_after_mb']:.1f} MB ({summary['mem_x']:.0f}x)")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"summary": summary, "tables": {"metrics": rows}}, f)