#!/usr/bin/env python3
"""Hands-on C03 — rate limiting, assembled from six lego blocks."""
import random, time
from collections import deque
from _harness import block, run_all, collect, check, approx

rng = random.Random(3)

def boundary_burst(n_burst=20, n_steady=40, rate=0.5):
    """A burst STRADDLING the window boundary at t=1.0, then steady traffic.

    Straddling is the point: a burst wholly inside one window is handled
    identically by every algorithm here, so it discriminates nothing.
    """
    half = n_burst // 2
    t = [0.95 + i * 0.004 for i in range(half)]          # 10 just before t=1.0
    t += [1.001 + i * 0.004 for i in range(n_burst - half)]  # 10 just after
    return t + [2.0 + i * rate for i in range(n_steady)]

@block(1, "Fixed window", "the obvious algorithm, and the boundary bug that fails interviews")
def b1(s, show):
    class FixedWindow:
        def __init__(self, limit, window): self.limit, self.w = limit, window; self.c = {}
        def allow(self, now):
            k = int(now // self.w)
            self.c = {k: self.c.get(k, 0)}          # only the current window matters
            if self.c[k] < self.limit:
                self.c[k] += 1; return True
            return False
    if show:
        lim = FixedWindow(limit=5, window=1.0)
        # the adversarial pattern: 5 at the END of window 0, 5 at the START of window 1
        times = [0.98, 0.98, 0.99, 0.99, 0.999] + [1.001, 1.001, 1.002, 1.002, 1.003]
        got = [lim.allow(t) for t in times]
        print(f"  limit = 5 per 1.0s window")
        print(f"  requests at t=0.98..0.999 : {sum(got[:5])} allowed")
        print(f"  requests at t=1.001..1.003: {sum(got[5:])} allowed")
        print(f"  ALL {sum(got)} allowed inside a {times[-1]-times[0]:.3f}s span "
              f"-- {sum(got)/ (times[-1]-times[0]):.0f}x the configured rate")
        print("  The counter resets on a wall-clock boundary, so a client that")
        print("  straddles it gets 2x the limit in an arbitrarily short interval.")
        print("  Memory: one integer per client. Correctness: 2x burst. This is the")
        print("  algorithm to name, then reject, in the first minute of the interview.")
    return {"FixedWindow": FixedWindow}

@block(2, "Sliding window log", "exactly correct, and you cannot afford it")
def b2(s, show):
    class SlidingLog:
        def __init__(self, limit, window): self.limit, self.w = limit, window; self.q = deque()
        def allow(self, now):
            while self.q and self.q[0] <= now - self.w: self.q.popleft()
            if len(self.q) < self.limit:
                self.q.append(now); return True
            return False
        def bytes_used(self): return len(self.q) * 8
    if show:
        lim = SlidingLog(limit=5, window=1.0)
        times = [0.98, 0.98, 0.99, 0.99, 0.999] + [1.001, 1.001, 1.002, 1.002, 1.003]
        got = [lim.allow(t) for t in times]
        print(f"  same adversarial pattern: {sum(got)} allowed (fixed window let "
              f"{10}) ")
        print(f"  {'clients':>9}{'limit':>8}{'memory':>12}{'at 1M clients':>16}")
        for limit in (5, 100, 10_000):
            per = limit * 8
            print(f"  {1:>9}{limit:>8}{per:>10} B{per*1_000_000/1e9:>14.1f} GB")
        print("  Exact, because it stores every timestamp. That is also why it is")
        print("  unusable: memory is O(limit) PER CLIENT, so a 10k/min limit across")
        print("  a million clients is 80 GB. Name it as the correctness reference,")
        print("  not as the answer.")
    return {"SlidingLog": SlidingLog}

@block(3, "Token bucket", "the one to actually implement, and why it is lazy")
def b3(s, show):
    class TokenBucket:
        def __init__(self, rate, burst):
            self.rate, self.burst = rate, burst
            self.tokens, self.last = float(burst), 0.0
        def allow(self, now, cost=1.0):
            self.tokens = min(self.burst, self.tokens + (now - self.last) * self.rate)
            self.last = now
            if self.tokens >= cost:
                self.tokens -= cost; return True
            return False
    if show:
        tb = TokenBucket(rate=5.0, burst=5)
        times = [0.98, 0.98, 0.99, 0.99, 0.999] + [1.001, 1.001, 1.002, 1.002, 1.003]
        got = [tb.allow(t) for t in times]
        print(f"  rate = 5/s, burst = 5")
        print(f"  adversarial pattern: {sum(got)} allowed "
              f"(fixed window 10, sliding log 5)")
        print(f"  {'t':>7}{'tokens before':>15}{'allowed':>9}")
        tb2 = TokenBucket(rate=5.0, burst=5)
        for t in (0.0, 0.1, 0.2, 0.4, 1.0, 2.0):
            before = min(tb2.burst, tb2.tokens + (t - tb2.last) * tb2.rate)
            a = tb2.allow(t)
            print(f"  {t:>7.1f}{before:>15.2f}{str(a):>9}")
        print("  No timer, no background thread, no per-request state cleanup: tokens")
        print("  are computed LAZILY from elapsed time on each call. Two floats per")
        print("  client, O(1) time, and burst is an explicit parameter rather than an")
        print("  accident. This is the answer.")
    return {"TokenBucket": TokenBucket}

@block(4, "Sliding window counter", "the memory/accuracy compromise everyone ships")
def b4(s, show):
    class SlidingCounter:
        """Weighted blend of the previous and current fixed windows."""
        def __init__(self, limit, window):
            self.limit, self.w = limit, window
            self.cur_key, self.cur, self.prev = 0, 0, 0
        def allow(self, now):
            k = int(now // self.w)
            if k != self.cur_key:
                self.prev = self.cur if k == self.cur_key + 1 else 0
                self.cur, self.cur_key = 0, k
            frac = 1.0 - (now % self.w) / self.w
            est = self.prev * frac + self.cur
            if est < self.limit:
                self.cur += 1; return True
            return False
    def worst_case(limit, eps):
        """Fill the previous window at its very END, then hammer at 1+eps."""
        sc = SlidingCounter(limit, 1.0)
        for j in range(limit):
            sc.allow(1.0 - 1e-9 * (limit - j))
        admitted = sum(sc.allow(1.0 + eps) for _ in range(limit * 5))
        # The old `limit` requests sit at t~1.0, still inside the trailing
        # window [eps, 1+eps] for any eps < 1. So true occupancy is the sum.
        return admitted, (limit + admitted) / limit

    def over_admission(limit, mult, n=40_000, seed=5):
        """Run ONLY the counter; check each admit against the TRUE trailing count.

        No second limiter, so there is no state-divergence confound: `bad` is
        exactly the count of requests a sliding log would have refused.
        """
        rng = random.Random(seed)
        sc, hist = SlidingCounter(limit, 1.0), deque()
        t, bad, adm = 0.0, 0, 0
        for _ in range(n):
            t += rng.expovariate(limit * mult)
            if sc.allow(t):
                while hist and hist[0] <= t - 1.0: hist.popleft()
                if len(hist) + 1 > limit: bad += 1
                hist.append(t); adm += 1
        return adm, bad, n

    if show:
        times = [0.98, 0.98, 0.99, 0.99, 0.999] + [1.001, 1.001, 1.002, 1.002, 1.003]
        sc = SlidingCounter(limit=5, window=1.0)
        got = [sc.allow(t) for t in times]
        print(f"  adversarial pattern: {sum(got)} allowed (fixed 10, log 5, bucket 5)")
        print(f"  memory: 2 integers per client vs {5*8} B for the log at limit=5")
        print()
        print("  I expected 'bounded error'. Measuring the worst case says otherwise:")
        print(f"  {'gap after boundary':>20}{'admitted':>10}{'true/limit':>12}")
        for eps in (0.1, 0.3, 0.5, 0.9, 0.99):
            adm, ratio = worst_case(100, eps)
            print(f"  {eps:>19.2f}s{adm:>10}{ratio:>11.2f}x")
        print("  The worst case tends to 2x -- the SAME bound as the fixed window")
        print("  this algorithm exists to fix. It does not remove the 2x; it makes")
        print("  the 2x require a specific arrival pattern instead of any burst.")
        print()
        print("  And the published '0.003% wrongly allowed' does not survive either.")
        print("  Measured, limit=100, Poisson arrivals, 40k requests each:")
        print(f"  {'offered load':>14}{'admitted':>10}{'over-limit':>12}{'% of all':>10}")
        for mult in (0.5, 0.9, 1.0, 1.5, 3.0):
            adm, bad, n = over_admission(100, mult)
            print(f"  {mult:>13.1f}x{adm:>10}{bad:>12}{bad/n*100:>9.2f}%")
        print("  Zero error while traffic is under the limit; 15-23% once it is at")
        print("  or above it. Cloudflare's figure is real and is measured in the")
        print("  regime where the limiter is not limiting. In the regime a limiter")
        print("  exists for, the error is four orders of magnitude larger.")
    return {"SlidingCounter": SlidingCounter}

@block(5, "Two servers", "every single-node algorithm is wrong the moment you scale out")
def b5(s, show):
    if show:
        print("  Run the token bucket independently on N servers, limit 5/s each")
        print(f"  {'servers':>9}{'per-server limit':>18}{'effective limit':>17}")
        for n in (1, 2, 4, 16):
            print(f"  {n:>9}{5:>18}{5*n:>17}")
        print("  Sharding the LIMIT instead (5/n per server) is worse: a client whose")
        print("  requests land unevenly gets throttled far below its quota.")
        print()
        print("  Three real options, and the trade each makes:")
        print(f"  {'design':<26}{'accuracy':>10}{'latency':>10}  {'blast radius':<20}")
        for name, acc, lat, blast in (
                ("central store (Redis)", "exact", "+1 RTT", "hard dependency"),
                ("local + async sync", "approx", "0", "drift on partition"),
                ("consistent-hash owner", "exact", "+1 RTT", "one shard per key")):
            print(f"  {name:<26}{acc:>10}{lat:>10}  {blast:<20}")
        print("  The follow-up is always 'what if Redis is down'. The answer that")
        print("  scores is fail-OPEN with a local fallback limiter, because a rate")
        print("  limiter that fails closed converts a cache outage into a full outage.")
    return {}

@block(6, "Atomicity", "check-then-set across a network is a race, not an implementation detail")
def b6(s, show):
    class RedisLike:
        def __init__(self): self.d = {}
        def get(self, k): return self.d.get(k, 0)
        def set(self, k, v): self.d[k] = v
        def incr(self, k):                      # atomic
            self.d[k] = self.d.get(k, 0) + 1; return self.d[k]
    def racy(store, key, limit, n_workers):
        allowed = 0
        for _ in range(n_workers):
            v = store.get(key)                  # every worker reads the same value
            if v < limit:
                allowed += 1
        for _ in range(allowed): store.incr(key)
        return allowed
    def atomic(store, key, limit, n_workers):
        allowed = 0
        for _ in range(n_workers):
            if store.incr(key) <= limit: allowed += 1
        return allowed
    if show:
        print(f"  limit = 5, {10} concurrent workers hitting the same key")
        r1 = RedisLike(); r2 = RedisLike()
        print(f"  GET-then-SET (read all, then write): {racy(r1, 'k', 5, 10):>2} allowed  "
              f"<- WRONG")
        print(f"  INCR and compare (single round trip): {atomic(r2, 'k', 5, 10):>2} allowed  "
              f"<- correct")
        print("  The racy version is what you write first. It is correct under no")
        print("  concurrency and wrong under exactly the load a rate limiter exists")
        print("  for. The fix is one atomic operation -- INCR, or a Lua script for")
        print("  the token bucket, since 'read tokens, compute, write tokens' is")
        print("  three round trips and two races.")
    return {"RedisLike": RedisLike}

def assembly(s):
    print("\nSix blocks = a production limiter. One traffic pattern, four algorithms.\n")
    times = boundary_burst()
    algos = [
        ("fixed window",     s["FixedWindow"](5, 1.0)),
        ("sliding log",      s["SlidingLog"](5, 1.0)),
        ("token bucket",     s["TokenBucket"](5.0, 5)),
        ("sliding counter",  s["SlidingCounter"](5, 1.0)),
    ]
    print(f"  {'algorithm':<20}{'allowed':>9}{'burst allowed':>15}"
          f"{'state/client':>14}{'exact':>7}")
    for name, lim in algos:
        got = [lim.allow(t) for t in times]
        burst = sum(got[:20])
        state = {"fixed window": "1 int", "sliding log": "N floats",
                 "token bucket": "2 floats", "sliding counter": "2 ints"}[name]
        exact = "yes" if name == "sliding log" else "no"
        print(f"  {name:<20}{sum(got):>9}{burst:>15}{state:>14}{exact:>7}")
    print("\n  60 requests: a 20-request burst STRADDLING the window boundary at")
    print("  t=1.0, then 40 spread over 20s at the configured rate. Straddling is")
    print("  the whole point -- a burst wholly inside one window is handled")
    print("  identically by all four, so it discriminates nothing. Put the burst on")
    print("  the boundary and the fixed window's 2x failure appears immediately.")
    print("\n  What to say, in order: fixed window is O(1) state and allows 2x at the")
    print("  boundary; sliding log is exact and O(limit) memory per client; token")
    print("  bucket is O(1) state, lazy, and makes burst an explicit parameter;")
    print("  sliding counter is O(1) state and -- per block 4 -- has the SAME 2x")
    print("  worst case as the fixed window, just harder to trigger. Then: none of")
    print("  them survive two servers without a shared store, and the shared store")
    print("  needs ONE atomic operation, and it must fail open.")
    print("\n  Note the burst column: the counter allowed 6, one MORE than the token")
    print("  bucket's 5, on a pattern chosen to embarrass the fixed window. That one")
    print("  request is the whole difference between 'bounded error' as a slogan and")
    print("  as a measurement -- and it is why the token bucket is the answer.")
    print("\n  Built: fixed window -> sliding log -> token bucket -> sliding counter")
    print("  -> distribution -> atomicity.")
    print("  Not built, and worth an extra 10 minutes if the interview goes there:")
    print("  hierarchical limits (per-user AND per-org), cost-weighted requests")
    print("  (an LLM call is not one unit), and the 429 + Retry-After contract.")

def parts():
    """Every mechanism this page builds, ready to import.

        >>> from c03_rate_limiter import parts
        >>> p = parts()
        >>> sorted(p)                      # doctest: +ELLIPSIS
        [...]
    """
    return collect()


def verify():
    """Re-derive every headline claim on this page from scratch."""
    # Rebuild the algorithms independently of the blocks, so a bug in a block
    # cannot make its own claim pass.
    class FW:
        def __init__(s, lim, w): s.lim, s.w, s.c = lim, w, {}
        def allow(s, t):
            k = int(t // s.w); s.c = {k: s.c.get(k, 0)}
            if s.c[k] < s.lim: s.c[k] += 1; return True
            return False
    class SL:
        def __init__(s, lim, w): s.lim, s.w, s.q = lim, w, deque()
        def allow(s, t):
            while s.q and s.q[0] <= t - s.w: s.q.popleft()
            if len(s.q) < s.lim: s.q.append(t); return True
            return False
    class TB:
        def __init__(s, r, b): s.r, s.b, s.tok, s.last = r, b, float(b), 0.0
        def allow(s, t, cost=1.0):
            s.tok = min(s.b, s.tok + (t - s.last) * s.r); s.last = t
            if s.tok >= cost: s.tok -= cost; return True
            return False
    class SC:
        def __init__(s, lim, w): s.lim, s.w, s.k, s.cur, s.prev = lim, w, 0, 0, 0
        def allow(s, t):
            k = int(t // s.w)
            if k != s.k:
                s.prev = s.cur if k == s.k + 1 else 0
                s.cur, s.k = 0, k
            if s.prev * (1.0 - (t % s.w) / s.w) + s.cur < s.lim:
                s.cur += 1; return True
            return False

    adv = [0.98, 0.98, 0.99, 0.99, 0.999, 1.001, 1.001, 1.002, 1.002, 1.003]

    # B1 -- the fixed window admits 2x the limit across a boundary.
    got_fw = sum(FW(5, 1.0).allow(t) for t in adv)
    check("B1  fixed window admits 2x the limit at a boundary",
          got_fw == 10, f"admitted {got_fw} against a limit of 5")

    # B2 -- the sliding log is exact on the same pattern.
    l = SL(5, 1.0); got_sl = sum(l.allow(t) for t in adv)
    check("B2  sliding log is exact on the same pattern",
          got_sl == 5, f"admitted {got_sl}, the correct answer")

    # B2 -- and its state is O(limit) per client, not O(1).
    big = SL(10_000, 60.0)
    for i in range(10_000): big.allow(i * 1e-4)
    check("B2  sliding log state is O(limit) per client",
          len(big.q) == 10_000, f"{len(big.q)*8:,} B for one client at limit=10k")

    # B3 -- the token bucket matches the log's exact answer from two floats.
    t = TB(5.0, 5); got_tb = sum(t.allow(x) for x in adv)
    check("B3  token bucket matches the log's answer",
          got_tb == got_sl, f"admitted {got_tb}, same as the sliding log")

    # B4 -- the sliding counter's worst case tends to 2x, NOT to a small bound.
    worst = 0.0
    for eps in (0.5, 0.9, 0.99):
        sc = SC(100, 1.0)
        for j in range(100): sc.allow(1.0 - 1e-9 * (100 - j))
        adm = sum(sc.allow(1.0 + eps) for _ in range(500))
        worst = max(worst, (100 + adm) / 100)
    check("B4  sliding counter's worst case approaches 2x, like the fixed window",
          1.95 <= worst < 2.0, f"measured {worst:.2f}x at a 0.99s gap")

    # B4 -- and its error is ~0 below the limit but large at/above it.
    def over(mult, n=40_000, seed=5):
        rng = random.Random(seed); sc, hist = SC(100, 1.0), deque()
        tt, bad = 0.0, 0
        for _ in range(n):
            tt += rng.expovariate(100 * mult)
            if sc.allow(tt):
                while hist and hist[0] <= tt - 1.0: hist.popleft()
                if len(hist) + 1 > 100: bad += 1
                hist.append(tt)
        return bad / n
    under, overld = over(0.5), over(1.5)
    check("B4  counter error is zero under the limit",
          under == 0.0, f"{under*100:.2f}% at 0.5x offered load")
    check("B4  ...and 15-25% at or above it",
          0.15 <= overld <= 0.25, f"{overld*100:.2f}% at 1.5x offered load")

    # B6 -- check-then-act admits `workers`, atomic INCR admits `limit`.
    store = {}
    reads = [store.get("k", 0) for _ in range(10)]
    racy = sum(1 for v in reads if v < 5)
    atomic = 0
    store["k"] = 0
    for _ in range(10):
        store["k"] += 1
        if store["k"] <= 5: atomic += 1
    check("B6  GET-then-SET admits one per concurrent worker",
          racy == 10, f"{racy} admitted against a limit of 5")
    check("B6  atomic INCR admits exactly the limit",
          atomic == 5, f"{atomic} admitted")

    # Assembly -- the four algorithms rank as the page claims on the burst.
    times = boundary_burst()
    burst = {name: sum([a.allow(x) for x in times][:20]) for name, a in
             (("fixed", FW(5, 1.0)), ("log", SL(5, 1.0)),
              ("bucket", TB(5.0, 5)), ("counter", SC(5, 1.0)))}
    check("ASM fixed window is the worst on a boundary-straddling burst",
          burst["fixed"] > burst["counter"] >= burst["log"] == burst["bucket"],
          f"fixed {burst['fixed']}, counter {burst['counter']}, "
          f"log {burst['log']}, bucket {burst['bucket']}")


if __name__ == "__main__":
    run_all(assembly, "HANDS-ON C03 — Rate limiting, block by block", verify=verify)