# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Token bucket against fixed window.

100 requests at 0.95 s and 100 more at 1.05 s, against a fixed-window
counter (100 per second), a token bucket (R = 100/s, C = 100), a token bucket
with C = 500, and a sliding-window counter.
"""
import argparse, json

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--limit", type=int, default=100, help="Requests per second allowed.")
p.add_argument("--burst", type=int, default=100, help="Requests in each burst.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
L = args.limit
requests = [0.95] * args.burst + [1.05] * args.burst


def fixed_window():
    counts, ok = {}, 0
    for t in requests:
        w = int(t); counts[w] = counts.get(w, 0) + 1; ok += counts[w] <= L
    return ok

def token_bucket(capacity):
    tokens, last, ok = capacity, 0.0, 0
    for t in requests:
        tokens = min(capacity, tokens + (t - last) * L); last = t
        if tokens >= 1: tokens -= 1; ok += 1
    return ok

def sliding_window():
    """Weighted count: this window's count plus the previous window's, weighted by its overlap."""
    counts, ok = {}, 0
    for t in requests:
        w = int(t); prev = counts.get(w - 1, 0); overlap = 1 - (t - w)
        if counts.get(w, 0) + prev * overlap < L: counts[w] = counts.get(w, 0) + 1; ok += 1
    return ok

rows = [{"limiter": "fixed window (100 per second)", "accepted": fixed_window()},
        {"limiter": "token bucket R=100/s, C=100", "accepted": token_bucket(L)},
        {"limiter": "token bucket R=100/s, C=500", "accepted": token_bucket(5 * L)},
        {"limiter": "sliding window counter", "accepted": sliding_window()}]
for r in rows: print(f"{r['limiter']:<32} accepted {r['accepted']:>3} of {len(requests)} in 0.1 s")

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