"""A load balancer that routes on the backends' own response times.

Four policies over one fleet and one arrival trace, on a virtual clock of
integer ticks. There is no wall clock, no thread and no sleep anywhere: the
only source of randomness is a `random.Random(seed)` threaded explicitly
through arrivals, service draws and the power-of-d sample.

The toy exists for one question. A latency-aware balancer sends traffic to
whichever backend answers fastest. What happens when the fastest answer in
the fleet is an error?

Retries and health checks are deliberately absent; see the commentary's §7
and §8. `../circuit-breaker/` already owns retry storms.
"""

import math
import random
from collections import deque


# --------------------------------------------------------------- backends

class Backend:
    """One server: FIFO, `slots` requests in service at once, work-conserving.

    `latency` is how long a call takes, in ticks. `healthy=False` means every
    call *fails* after exactly that long -- a broken backend that is, to
    anything measuring response times, indistinguishable from a fast one.
    That equivalence is the whole toy, and `dispatch` is where it lives: the
    duration does not depend on `healthy`, only the verdict does.
    """

    def __init__(self, name, latency, healthy=True, slots=1, cv=0.0):
        self.name = name
        self.latency = latency
        self.healthy = healthy
        self.slots = slots
        self.cv = cv                    # coefficient of variation, 0 = constant
        self.free_at = [0] * slots      # per-slot tick the slot next falls idle
        self.outstanding = deque()      # completion ticks, ascending (FIFO)
        self.sent = 0                   # requests routed here
        self.ok = 0
        self.err = 0
        self.ewma = float(latency)      # the balancer's belief about this backend

    def draw(self, rng):
        """Sample a service time. Constant unless `cv` is set (see §6.5)."""
        if self.cv <= 0:
            return self.latency
        sigma = math.sqrt(math.log(1 + self.cv * self.cv))
        mu = math.log(self.latency) - sigma * sigma / 2
        return max(1, int(round(rng.lognormvariate(mu, sigma))))

    def retire(self, now):
        """Forget requests that have completed by `now`. FIFO, so the deque
        is sorted and we can stop at the first one still running."""
        while self.outstanding and self.outstanding[0] <= now:
            self.outstanding.popleft()

    def inflight(self):
        return len(self.outstanding)

    def dispatch(self, now, rng):
        """Accept one request at tick `now`. Returns (latency, ok).

        The request waits for the earliest-free slot, so a saturated backend
        makes new arrivals queue and its observed latency climbs.
        """
        i = min(range(self.slots), key=lambda k: self.free_at[k])
        start = max(now, self.free_at[i])
        done = start + self.draw(rng)
        self.free_at[i] = done
        self.outstanding.append(done)
        self.sent += 1
        if self.healthy:
            self.ok += 1
        else:
            self.err += 1
        return done - now, self.healthy


# --------------------------------------------------------------- policies

class RoundRobin:
    """The dumb one. Reads no signal at all, so no signal can mislead it."""

    name = "round-robin"

    def __init__(self):
        self.i = -1

    def pick(self, fleet, now, rng):
        self.i = (self.i + 1) % len(fleet)
        return fleet[self.i]

    def observe(self, b, latency, ok):
        pass


class ScorePolicy:
    """Pick the lowest-scoring backend out of `d` sampled at random.

    `d=None` scores the whole fleet -- the classic "least X" policy. Any
    integer `d` makes it power-of-d-choices over the same score, which §6.3
    shows is a dial on the blast radius rather than a fix.
    """

    def __init__(self, d=None):
        self.d = d
        self.name = self.base + ("" if d is None else "(d=%d)" % d)
        self.last = ()                  # the candidate set of the last pick

    def score(self, b):
        raise NotImplementedError

    def pick(self, fleet, now, rng):
        cands = fleet if self.d is None else rng.sample(fleet, min(self.d, len(fleet)))
        self.last = cands               # so the demo can count won-vs-sampled
        return min(cands, key=self.score)

    def observe(self, b, latency, ok):
        pass


class LeastConn(ScorePolicy):
    """Fewest requests in flight. Ties go to the earliest in the list."""

    base = "least-conn"

    def score(self, b):
        return b.inflight()


class PeakEwma(ScorePolicy):
    """Finagle's peak-EWMA: an exponentially weighted mean of observed
    latency, scaled by the queue the next request would join. This is the
    policy the toy is about, and `observe` is why -- it learns from the
    duration of a call and never from its outcome.
    """

    base = "peak-ewma"

    def __init__(self, d=None, alpha=0.2):
        ScorePolicy.__init__(self, d)
        self.alpha = alpha

    def score(self, b):
        return b.ewma * (b.inflight() + 1)

    def observe(self, b, latency, ok):
        b.ewma = (1 - self.alpha) * b.ewma + self.alpha * latency


class PeakEwmaOverSuccess(PeakEwma):
    """The fix (§7.3): divide the score by the backend's observed success
    rate, Laplace-smoothed so a backend with no history is not divided by
    zero. One term, and it is the only term in the toy that reads `ok`.
    """

    base = "ewma/success"

    def score(self, b):
        rate = (b.ok + 1) / (b.ok + b.err + 2)
        return b.ewma * (b.inflight() + 1) / rate


# ------------------------------------------------------------------- run

def poisson(rng, lam):
    """Knuth's method. `lam` is well under 1 here, so this loops ~once."""
    limit, k, p = math.exp(-lam), 0, 1.0
    while True:
        p *= rng.random()
        if p <= limit:
            return k
        k += 1


def simulate(policy, fleet, ticks=20000, lam=0.32, seed=1):
    """Run one policy against one fleet. Every tick: retire what finished,
    admit this tick's arrivals, route each one, feed the outcome back."""
    rng = random.Random(seed)
    samples = []                        # (latency, ok) in arrival order
    for now in range(ticks):
        for b in fleet:
            b.retire(now)
        for _ in range(poisson(rng, lam)):
            b = policy.pick(fleet, now, rng)
            latency, ok = b.dispatch(now, rng)
            policy.observe(b, latency, ok)
            samples.append((latency, ok))
    total = len(samples)
    good = sum(1 for _, ok in samples if ok)
    return {
        "policy": policy.name,
        "total": total,
        "good": good,
        "sent": [b.sent for b in fleet],
        "broken_share": (sum(b.sent for b in fleet if not b.healthy) / total
                         if total else 0.0),
        "all": [lat for lat, _ in samples],
        "ok": [lat for lat, ok in samples if ok],
    }


def pct(xs, q):
    """The q-quantile by nearest rank. No interpolation: every number this
    toy prints is a latency some request actually experienced."""
    if not xs:
        return float("nan")
    s = sorted(xs)
    return s[min(len(s) - 1, int(q * len(s)))]


def row(r):
    return ("  %-16s share=%6.2f%%  goodput=%5d/%d  p99(all)=%3.0f  p99(ok)=%3.0f"
            % (r["policy"], 100 * r["broken_share"], r["good"], r["total"],
               pct(r["all"], .99), pct(r["ok"], .99)))