"""A lease-based distributed lock, and the resource it fails to protect.

N independent lock servers, majority acquire, TTL leases -- Redlock's shape.
No threads, no sockets, no wall clock, no RNG: time is an integer tick and a
run is a pure function of an explicit event list, so one interleaving can be
staged on purpose.

The point is not that the lock is buggy. Every lease this file grants is
audited, and `overlapping_grants` asserts that no two of them are ever valid
at the same instant. The resource is corrupted anyway, because a client can
be paused between "I hold the lock" and "I write", and no lock implementation
can close that gap. Martin Kleppmann's argument; the fix is `Resource.fencing`.
"""


class LockServer:
    """One lock server. Holds at most one lease at a time, and remembers every
    lease it ever issued so the run can be audited afterwards."""

    def __init__(self, sid):
        self.id = sid
        self.holder = None
        self.expires = 0
        self.token = 0          # monotonic; the fencing token comes from here
        self.grants = []        # audit log: every lease, as [start, end)

    def acquire(self, client, now, ttl):
        """Grant a lease, unless a live one is outstanding. The expiry test is
        the only thing that releases a lease from a client that never came
        back -- there is nobody to ask whether A is still alive."""
        if self.holder is not None and now < self.expires:
            return None
        self.token += 1
        self.holder = client
        self.expires = now + ttl
        self.grants.append(dict(server=self.id, client=client, token=self.token,
                                start=now, end=self.expires))
        return self.token

    def release(self, client, now):
        """Only the current holder can release, and only while its lease is
        live. A client whose lease already expired releases nothing -- which is
        why A's release at the end of the demo is a no-op."""
        if self.holder != client or now >= self.expires:
            return False
        self.holder = None
        self.expires = now
        for g in reversed(self.grants):
            if g["client"] == client and g["start"] <= now < g["end"]:
                g["end"] = now      # truncate the audit record to reality
                break
        return True

    def holder_at(self, t):
        """Which grant, if any, was valid at instant t. Reads the audit log
        rather than the live state, so it answers about the past."""
        for g in self.grants:
            if g["start"] <= t < g["end"]:
                return g
        return None


class LockService:
    """N independent servers; a client needs a majority to hold the lock.
    N=1 is a single lock server, N=5 is Redlock's recommended deployment."""

    def __init__(self, n):
        self.servers = [LockServer(i) for i in range(1, n + 1)]
        self.quorum = n // 2 + 1

    def acquire(self, client, now, ttl):
        """Ask everyone, keep it if a majority said yes, otherwise hand back
        the partial grants immediately. Returns the highest token seen, which
        is the fencing token."""
        toks = [s.acquire(client, now, ttl) for s in self.servers]
        got = [t for t in toks if t is not None]
        if len(got) >= self.quorum:
            return max(got)
        for s, t in zip(self.servers, toks):
            if t is not None:
                s.release(client, now)
        return None

    def release(self, client, now):
        for s in self.servers:
            s.release(client, now)

    def held_by(self, client, t):
        """Did `client` hold a quorum of valid leases at instant t?"""
        n = sum(1 for s in self.servers
                if (s.holder_at(t) or {}).get("client") == client)
        return n >= self.quorum


class Resource:
    """The storage service the clients are locking *around*. It is the thing
    that must not be corrupted, and it is a separate system from the lock --
    which is the whole problem. With `fencing`, it refuses any write whose
    token is not strictly higher than the highest it has already applied."""

    def __init__(self, fencing=False):
        self.fencing = fencing
        self.value = ""
        self.max_token = 0
        self.log = []

    def read(self, client, now):
        return self.value

    def write(self, client, now, value, token):
        # `<`, not `<=`: tokens are unique per grant, so a stale writer is
        # always strictly lower, and the loose test lets one holder write
        # twice under a single lease. See commentary 7.3.
        ok = not (self.fencing and token < self.max_token)
        if ok:
            self.max_token = max(self.max_token, token)
            self.value = value
        self.log.append(dict(t=now, client=client, value=value,
                             token=token, ok=ok))
        return ok


class Schedule:
    """The driver. Integer ticks, an explicit event list, ties broken by
    insertion order. Events may schedule later events, so the list is kept
    sorted as it is consumed."""

    def __init__(self):
        self.events = []
        self.trace = []
        self.seq = 0

    def at(self, t, fn):
        self.seq += 1
        self.events.append((t, self.seq, fn))

    def say(self, t, s):
        self.trace.append(f"  t={t:>3}  {s}")

    def run(self):
        done = 0
        while done < len(self.events):
            self.events.sort()
            t, _, fn = self.events[done]
            done += 1
            fn(t)


def scenario(ttl, pause_from, pause_to, b_arrives, n=5, fencing=False):
    """Two clients, one resource, one lease-based lock.

    A: acquire at t=0, read at t=1, PAUSE [pause_from, pause_to), then write
       and release. The pause is modelled as the absence of scheduled events
       -- which is exactly what a stop-the-world pause is to the rest of the
       world. A is never told it happened.
    B: from t=b_arrives, poll for the lock every tick; on success, read one
       tick later, write the tick after that, then release.

    Both clients do the same read-modify-write: append their own letter to
    whatever the resource held when they read it.
    """
    lock, res, sch = LockService(n), Resource(fencing), Schedule()
    st = {}

    def a_acquire(t):
        st["a_token"] = lock.acquire("A", t, ttl)
        sch.say(t, f"A  acquire -> token {st['a_token']}, "
                   f"lease [{t}, {t + ttl})")

    def a_read(t):
        st["a_saw"] = res.read("A", t)
        sch.say(t, f"A  read -> {st['a_saw']!r}")

    def a_pause(t):
        sch.say(t, f"A  ---- pause begins ({pause_to - pause_from} ticks) ----")

    def a_write(t):
        sch.say(t, "A  ---- pause ends; A still believes it holds the lock ----")
        if st.get("a_token") is None:
            return
        v = st["a_saw"] + "A"
        ok = res.write("A", t, v, st["a_token"])
        sch.say(t, f"A  write {v!r} token={st['a_token']} -> "
                   f"{'ACCEPTED' if ok else 'REJECTED (stale token)'}")

    def b_poll(t):
        if st.get("b_token") is not None:
            return
        tok = lock.acquire("B", t, ttl)
        if tok is None:
            return
        st["b_token"], st["b_at"] = tok, t
        sch.say(t, f"B  acquire -> token {tok}, lease [{t}, {t + ttl})")
        sch.at(t + 1, b_read)
        sch.at(t + 2, b_write)
        sch.at(t + 3, lambda u: lock.release("B", u))

    def b_read(t):
        st["b_saw"] = res.read("B", t)
        sch.say(t, f"B  read -> {st['b_saw']!r}")

    def b_write(t):
        v = st["b_saw"] + "B"
        ok = res.write("B", t, v, st["b_token"])
        sch.say(t, f"B  write {v!r} token={st['b_token']} -> "
                   f"{'ACCEPTED' if ok else 'REJECTED (stale token)'}")

    sch.at(0, a_acquire)
    sch.at(1, a_read)
    sch.at(pause_from, a_pause)
    sch.at(pause_to, a_write)
    sch.at(pause_to + 1, lambda t: lock.release("A", t))
    for t in range(b_arrives, pause_to + 200):
        sch.at(t, b_poll)
    sch.run()
    return lock, res, st, sch.trace


def overlapping_grants(lock):
    """THE innocence assertion. Every lease every server ever issued, compared
    pairwise as half-open intervals. Empty means the lock never once let two
    clients hold it at the same instant."""
    bad = []
    for s in lock.servers:
        for i, g in enumerate(s.grants):
            for h in s.grants[i + 1:]:
                if g["start"] < h["end"] and h["start"] < g["end"]:
                    bad.append((s.id, g, h))
    return bad


def stale_overwrites(res):
    """THE corruption. An accepted write carrying a token lower than one the
    resource has already applied: an older lease's work landing on top of a
    newer lease's completed critical section."""
    hi, bad = 0, []
    for w in res.log:
        if not w["ok"]:
            continue
        if w["token"] < hi:
            bad.append(dict(t=w["t"], client=w["client"],
                            token=w["token"], applied=hi))
        hi = max(hi, w["token"])
    return bad


def unlocked_writes(lock, res):
    """Accepted writes made by a client that did not hold a quorum of valid
    leases at that instant. Related to `stale_overwrites` but not the same:
    fencing eliminates the second and not the first."""
    return [dict(t=w["t"], client=w["client"])
            for w in res.log
            if w["ok"] and not lock.held_by(w["client"], w["t"])]