#!/usr/bin/env python3
"""
DDIA Ch. 7 (Sharding): why `hash(key) % N` is a rebalancing catastrophe,
and how a fixed number of shards behind a shard->node mapping fixes it.

Two schemes assign 1,000,000 keys to nodes, then grow the cluster 10 -> 11
nodes. We measure what fraction of keys have to move (i.e. be shipped over
the network to a different node) under each.

  Scheme A -- direct: node = hash(key) % N.  Grow N: 10 -> 11.
  Scheme B -- fixed shards: shard = hash(key) % 1000, then a shard->node
              mapping. Grow by reassigning the minimum number of whole
              shards onto the new node.

Deterministic on purpose:
  * hashing is md5-based, NOT Python's built-in hash() (which is salted
    per process via PYTHONHASHSEED -- unsafe for sharding, exactly the
    trap the book warns about with Java's Object.hashCode()).
  * the one place we make a choice (which shards to move) is seeded.

Run from ddia/ch07:
    python3 code/mod_n_rebalancing.py
"""

import hashlib
import random

N_KEYS = 1_000_000
SHARDS = 1000
OLD_NODES = 10
NEW_NODES = 11
SEED = 42


def h(key):
    """Stable hash. md5 of the key as an integer.

    We deliberately do NOT use Python's built-in hash(): for str it is
    randomized per process (PYTHONHASHSEED), so `hash(key) % N` would send
    the same key to a different node on every restart. The book makes the
    same point about Java's hashCode being unsuitable for sharding.
    """
    return int(hashlib.md5(key.encode()).hexdigest(), 16)


def bar():
    print("=" * 68)


def main():
    keys = [f"key{i}" for i in range(N_KEYS)]
    hashes = [h(k) for k in keys]

    bar()
    print("DDIA Ch.7 -- rebalancing cost: hash % N   vs   fixed shards")
    print(f"  {N_KEYS:,} keys ('key0'..'key{N_KEYS-1}'), grow {OLD_NODES} -> {NEW_NODES} nodes")
    print(f"  stable md5 hashing (not Python hash()), seed={SEED}")
    bar()

    # ---------- Scheme A: node = hash(key) % N ----------
    moved_A = 0
    moved_to_old = 0  # of the moved keys, how many land on an OLD node id <= 9
    for x in hashes:
        old_node = x % OLD_NODES
        new_node = x % NEW_NODES
        if old_node != new_node:
            moved_A += 1
            if new_node <= OLD_NODES - 1:  # new_node in 0..9 -> a pre-existing node
                moved_to_old += 1
    frac_A = moved_A / N_KEYS

    print()
    print("Scheme A -- direct: node = hash(key) % N,  grow N 10 -> 11")
    print(f"    keys that move : {moved_A:>9,} of {N_KEYS:,}")
    print(f"    fraction moved : {frac_A*100:6.2f}%")
    print(f"    a key stays put ONLY when  (h % 10) == (h % 11)")
    print(f"    of the movers, {moved_to_old:,} land on an OLD node (id 0-9),")
    print(f"    not the new node -- pure churn between nodes that already existed.")

    # ---------- Scheme B: fixed number of shards ----------
    # Count how many keys live in each of the 1000 shards.
    shard_size = [0] * SHARDS
    for x in hashes:
        shard_size[x % SHARDS] += 1

    # Initial mapping for 10 nodes: shard s -> node (s % 10). Each node owns
    # exactly 100 shards.
    node_shards = {n: [] for n in range(OLD_NODES)}
    for s in range(SHARDS):
        node_shards[s % OLD_NODES].append(s)

    # Grow to 11 nodes with MINIMAL movement: target ~ 1000/11 = 90.9 shards
    # per node. 1000 = 11*90 + 10, so we hand node 10 exactly 90 shards by
    # taking 9 whole shards off each of the 10 over-full nodes (100 -> 91).
    # Nothing else moves. Which 9 is a free choice -> seed it.
    rng = random.Random(SEED)
    per_node_give_up = 9  # 10 nodes * 9 = 90 shards -> new node 10
    moved_shards = []
    for n in range(OLD_NODES):
        owned = list(node_shards[n])
        rng.shuffle(owned)
        moved_shards.extend(owned[:per_node_give_up])

    moved_B = sum(shard_size[s] for s in moved_shards)
    frac_B = moved_B / N_KEYS

    # sanity: final shard counts per node
    final_counts = {n: len(node_shards[n]) - per_node_give_up for n in range(OLD_NODES)}
    final_counts[OLD_NODES] = len(moved_shards)  # the new node

    print()
    print("Scheme B -- fixed shards: shard = hash(key) % 1000, map shards -> nodes")
    print(f"    1000 shards, 10 nodes own 100 each; grow to 11 nodes")
    print(f"    move {per_node_give_up} whole shards off each old node -> new node 10")
    print(f"    shards moved   : {len(moved_shards):>9,} of {SHARDS:,}")
    print(f"    keys moved     : {moved_B:>9,} of {N_KEYS:,}")
    print(f"    fraction moved : {frac_B*100:6.2f}%")
    print(f"    final shards/node: 10 old nodes -> 91 each, new node 10 -> {final_counts[OLD_NODES]}")

    # ---------- The magnitude ----------
    print()
    bar()
    print("The gap")
    print(f"    hash % N     moves {frac_A*100:6.2f}%  ({moved_A:,} keys)")
    print(f"    fixed shards moves {frac_B*100:6.2f}%  ({moved_B:,} keys)")
    print(f"    ratio        : {moved_A/moved_B:.1f}x more data shipped by hash % N")
    print(f"    same 1,000,000 keys, same +1 node -- {moved_A/moved_B:.1f}x the network.")
    bar()


if __name__ == "__main__":
    main()