#!/usr/bin/env python3
"""
DDIA Ch. 10 -- The CAP theorem, made countable.

Two regions, R1 and R2, each a {client, replica}, joined by one inter-region
link. The link goes DOWN (a partition -- an event imposed on the system, not a
setting anyone chose). While it is down, the SAME fixed workload is replayed
against two deployments of the exact same two regions:

  (a) CP  -- single-leader. The leader lives in R1. R2 is a follower.
  (b) AP  -- multi-leader. Both regions accept writes locally.

During the partition:
  * CP keeps every replica consistent, but R2 (the non-leader side) cannot serve
    writes or linearizable reads -- it must REFUSE them. We count the refusals.
  * AP serves every request at both regions -- zero refusals -- but the two
    regions' states DIVERGE. On heal we count the keys that conflict and need
    reconciliation.

Nothing here is fabricated: run it and the two costs are printed as measured.
Deterministic: fixed workload, RNG seeded so the interleaving reproduces exactly.
"""

import random

# ---------------------------------------------------------------------------
# The two regions and their replicas
# ---------------------------------------------------------------------------

class Replica:
    """One region's copy of the keyspace: key -> (value, version)."""
    def __init__(self, keys):
        # both regions start replicated and in agreement, at version 0
        self.store = {k: ("v0", 0) for k in keys}

    def local_write(self, key, value):
        _, ver = self.store[key]
        self.store[key] = (value, ver + 1)

    def local_read(self, key):
        return self.store[key]


# ---------------------------------------------------------------------------
# The workload -- one fixed sequence, replayed identically against CP and AP
# ---------------------------------------------------------------------------
# Each request is (region, kind, key, value). kind is one of:
#   "write"             -- mutate a key
#   "read_linearizable" -- must return the latest value (needs the leader)
#   "read_stale"        -- may return a possibly-stale local value

KEYS = ["x", "y", "z", "w", "a", "b", "c"]

# R1 (the leader side) writes three keys during the partition.
R1_REQUESTS = [
    ("R1", "write", "x", "r1"),
    ("R1", "write", "y", "r1"),
    ("R1", "write", "z", "r1"),
]

def build_r2_requests():
    """R2's 20 requests during the partition, order fixed by a seeded shuffle."""
    reqs = []
    # 7 writes -- three (x, y, z) collide with R1's writes; four (w, a, b, c) don't
    for k in ["x", "y", "z", "w", "a", "b", "c"]:
        reqs.append(("R2", "write", k, "r2"))
    # 5 linearizable reads
    for k in ["x", "y", "a", "w", "z"]:
        reqs.append(("R2", "read_linearizable", k, None))
    # 8 stale-tolerant reads
    for k in ["x", "y", "z", "w", "a", "b", "c", "x"]:
        reqs.append(("R2", "read_stale", k, None))
    assert len(reqs) == 20
    rng = random.Random(1010)          # seeded -> deterministic interleaving
    rng.shuffle(reqs)
    return reqs

WORKLOAD = R1_REQUESTS + build_r2_requests()


# ---------------------------------------------------------------------------
# Deployment (a): CP -- single leader in R1
# ---------------------------------------------------------------------------

def run_cp():
    """
    Leader in R1. Link is DOWN, so R2 (follower) cannot reach the leader.
    R2 must refuse writes and linearizable reads; it may still serve stale reads
    from its local replica. R1 serves everything (leader is local).
    Consistency is preserved: only the leader takes writes, so nothing diverges.
    """
    r1 = Replica(KEYS)
    r2 = Replica(KEYS)  # follower, frozen at v0 for the partition's duration

    refused = 0
    served = 0
    r2_refused = []
    r2_stale_example = None

    for region, kind, key, value in WORKLOAD:
        if region == "R1":
            # leader side, reachable locally
            if kind == "write":
                r1.local_write(key, value)
            served += 1
            continue

        # region == "R2": follower, cut off from the leader
        if kind in ("write", "read_linearizable"):
            refused += 1
            r2_refused.append((kind, key))
        else:  # read_stale -- served from the local (possibly stale) replica
            served += 1
            if r2_stale_example is None:
                r2_stale_example = (key, r2.local_read(key))

    return {
        "refused": refused,
        "served": served,
        "refused_ops": r2_refused,
        "stale_example": r2_stale_example,
        "diverged_keys": [],           # only the leader wrote -> no divergence
    }


# ---------------------------------------------------------------------------
# Deployment (b): AP -- multi-leader, both regions accept writes locally
# ---------------------------------------------------------------------------

def run_ap():
    """
    Both regions are leaders. Every request is served locally -> zero refusals.
    But each side writes without the other, so their replicas diverge. On heal
    we compare them: a key written on BOTH sides with different values conflicts.
    """
    r1 = Replica(KEYS)
    r2 = Replica(KEYS)
    written = {"R1": {}, "R2": {}}     # key -> value each side actually wrote

    refused = 0
    served = 0

    for region, kind, key, value in WORKLOAD:
        replica = r1 if region == "R1" else r2
        if kind == "write":
            replica.local_write(key, value)
            written[region][key] = value
        else:
            replica.local_read(key)    # always answerable from the local copy
        served += 1

    # Heal: reconcile. A key written on both sides with differing values conflicts.
    both = sorted(set(written["R1"]) & set(written["R2"]))
    diverged = [k for k in both if written["R1"][k] != written["R2"][k]]

    return {
        "refused": refused,
        "served": served,
        "diverged_keys": diverged,
        "r1_values": {k: r1.local_read(k) for k in diverged},
        "r2_values": {k: r2.local_read(k) for k in diverged},
    }


# ---------------------------------------------------------------------------
# Transcript
# ---------------------------------------------------------------------------

def main():
    total = len(WORKLOAD)
    r2_total = sum(1 for r in WORKLOAD if r[0] == "R2")

    print("=" * 70)
    print("CAP, made countable -- two regions, one link, the SAME workload twice")
    print("=" * 70)
    print(f"R1 = {{client, replica}}   R2 = {{client, replica}}   link: DOWN (partition)")
    print(f"workload: {total} requests during the partition "
          f"({len(R1_REQUESTS)} at R1, {r2_total} at R2)")
    print()
    print("The partition is NOT a setting we chose -- the link is cut. Both")
    print("deployments below face the identical outage; they differ only in what")
    print("they give up while it lasts.")
    print()

    # ---- CP -----------------------------------------------------------------
    cp = run_cp()
    print("-" * 70)
    print("(a) CP  --  single-leader, leader in R1  =>  choose CONSISTENCY")
    print("-" * 70)
    print("R2 is a follower cut off from the leader. It CANNOT serve writes or")
    print("linearizable reads, so it refuses them; it still serves stale reads.")
    print()
    print(f"   R2 requests refused (writes + linearizable reads): "
          f"{cp['refused']} of {r2_total}")
    print(f"   R2 requests served (stale-tolerant reads only):    "
          f"{cp['served'] - len(R1_REQUESTS)} of {r2_total}")
    ex_key, (ex_val, ex_ver) = cp["stale_example"]
    print(f"   example stale read at R2:  read('{ex_key}') -> "
          f"('{ex_val}', v{ex_ver})   (leader has moved on; R2 is behind)")
    print(f"   keys diverged (need reconciliation on heal):       "
          f"{len(cp['diverged_keys'])}")
    print()
    print(f"   COST:  {cp['refused']} refused requests. "
          f"availability sacrificed on the minority side; consistency intact.")
    print()

    # ---- AP -----------------------------------------------------------------
    ap = run_ap()
    print("-" * 70)
    print("(b) AP  --  multi-leader, both regions accept writes  =>  choose AVAILABILITY")
    print("-" * 70)
    print("Every request is served locally at whichever region issued it.")
    print("Nothing is refused -- but the two replicas drift apart.")
    print()
    print(f"   requests refused:                                  "
          f"{ap['refused']} of {total}")
    print(f"   requests served:                                   "
          f"{ap['served']} of {total}")
    print(f"   keys diverged (need reconciliation on heal):       "
          f"{len(ap['diverged_keys'])}  -> {ap['diverged_keys']}")
    for k in ap["diverged_keys"]:
        v1, ver1 = ap["r1_values"][k]
        v2, ver2 = ap["r2_values"][k]
        print(f"       key '{k}':  R1 has ('{v1}', v{ver1})   R2 has ('{v2}', v{ver2})"
              f"   <- conflict")
    print()
    print(f"   COST:  {len(ap['diverged_keys'])} conflicting keys to reconcile. "
          f"consistency sacrificed; availability intact.")
    print()

    # ---- Takeaway -----------------------------------------------------------
    print("=" * 70)
    print("TAKEAWAY")
    print("=" * 70)
    print(f"   CP:  refused = {cp['refused']:>2}   diverged keys = {len(cp['diverged_keys'])}")
    print(f"   AP:  refused = {ap['refused']:>2}   diverged keys = {len(ap['diverged_keys'])}")
    print()
    print("   The partition happened to BOTH deployments -- it was never optional.")
    print("   You did not 'pick 2 of 3'. You chose, WHILE partitioned, between")
    print("   refusing service (CP) and diverging (AP), and each cost is countable.")


if __name__ == "__main__":
    main()