"""
Practice: retries, measured. Lab 13's script 01 (01_retries.py) and resilience.py: which failures to
retry, exponential backoff with jitter, what coin-flip math predicts, how stacked retries multiply,
a retry loop that honors Retry-After, and the test server's rate limit (a token bucket) that let
only 12 of a 113-call wave in.

Fill in the TODOs, then run the tests:
    pytest test_retry_practice.py        (locally)
or press "Run tests" on the chapter page.
"""
import math

RETRYABLE_STATUS = {408, 409, 429, 500, 502, 503, 504}


class Failure(Exception):
    """A failed call. `status` is an HTTP status code (429, 503, 400...) or "timeout" or
    "connection". `retry_after` is the server's "try again in N seconds", if it sent one."""

    def __init__(self, status, retry_after=None):
        super().__init__(f"failed: {status}")
        self.status, self.retry_after = status, retry_after


class GaveUp(Exception):
    """retry() used every attempt. `attempts` is how many calls it made."""

    def __init__(self, attempts):
        super().__init__(f"gave up after {attempts} attempt(s)")
        self.attempts = attempts


def is_retryable(status):
    """True for failures that may go away on their own: a timeout, a dropped connection, or a
    status in RETRYABLE_STATUS (rate limits and server errors). False for everything else: a 400
    (the request itself is wrong), a 401 or 403 (bad key, no permission)."""
    return status in ("timeout", "connection") or status in RETRYABLE_STATUS


def backoff(n, base=0.5, cap=8.0, jitter="none", u=0.5):
    """The wait before retry n (n = 0 for the first retry). Exponential: base × 2**n, at most cap.
    jitter "none": exactly that. "full": u × that (anywhere from 0 to it). "equal": half of it,
    plus u × the other half. `u` is the random number, between 0 and 1, so tests can fix it."""
    exp = min(cap, base * 2 ** n)
    if jitter == "none":
        return exp
    if jitter == "full":
        return u * exp
    if jitter == "equal":
        return exp / 2 + u * exp / 2
    raise ValueError(jitter)


def chance_through(p, tries):
    """Each call fails at random with chance p, independently (a coin flip). The chance that at
    least one of `tries` calls gets through."""
    return 1 - p ** tries


def wilson(k, n, z=1.96):
    """The 95% Wilson interval for k successes out of n, as (low, high)."""
    if n == 0:
        return (0.0, 1.0)
    phat = k / n
    denom = 1 + z * z / n
    centre = (phat + z * z / (2 * n)) / denom
    half = z * math.sqrt(phat * (1 - phat) / n + z * z / (4 * n * n)) / denom
    return (centre - half, centre + half)


def calls_to_dead_service(tries_per_layer):
    """Every layer (your code, the SDK, a gateway) retries the layer below it. Against a service
    that always fails, how many calls does one user message send? `tries_per_layer` lists each
    layer's total tries (1 = that layer doesn't retry)."""
    return math.prod(tries_per_layer)


def retry(call, attempts=5, wait=lambda n: backoff(n), sleep=lambda s: None):
    """Call call() until it returns, at most `attempts` times, and return its result.
    - A Failure that isn't retryable (is_retryable(e.status) is False) is raised at once.
    - Otherwise, before the next try, sleep for wait(n) (n = 0 before the first retry), but never
      less than the failure's retry_after, if it has one.
    - Don't sleep after the last attempt: raise GaveUp(attempts) instead."""
    for i in range(attempts):
        try:
            return call()
        except Failure as e:
            if not is_retryable(e.status):
                raise
            if i == attempts - 1:
                break
            s = wait(i)
            if e.retry_after is not None:
                s = max(s, e.retry_after)
            sleep(s)
    raise GaveUp(attempts)


class TokenBucket:
    """The test server's rate limit: `rate` tokens a second, at most `burst` saved up, starting
    full. take(now) is True (let in) if a token is there, and uses it; False means turned away
    (a 429). `now` is the time in seconds; it never goes backwards."""

    def __init__(self, rate, burst, start=0.0):
        self.rate, self.burst = rate, burst
        self.tokens, self.last = float(burst), start

    def take(self, now):
        self.tokens = min(self.burst, self.tokens + (now - self.last) * self.rate)
        self.last = now
        if self.tokens >= 1:
            self.tokens -= 1
            return True
        return False