# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Escrow removes coordination.

Ten nodes sell a global budget of 1,000 discount codes under skewed demand.
Version 1 asks one coordinator for every sale; version 2 gives each node its
own share and moves half of the richest node's remainder when a node runs out.
"""
import argparse, json, random

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--nodes", type=int, default=10, help="Selling nodes.")
p.add_argument("--codes", type=int, default=1000, help="Global budget of codes.")
p.add_argument("--demand", type=int, default=1500, help="Purchase attempts (more than the budget, so it runs out).")
p.add_argument("--skew", type=float, default=1.2, help="Zipf exponent of demand across nodes.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
random.seed(args.seed)
weights = [1 / (i + 1) ** args.skew for i in range(args.nodes)]
attempts = random.choices(range(args.nodes), weights, k=args.demand)

# Version 1: a coordinator holds the count; every sale is a request and a reply
left, msgs1, sold1 = args.codes, 0, 0
for _ in attempts:
    msgs1 += 2
    if left > 0: left -= 1; sold1 += 1

# Version 2: escrow — each node holds a share, and only a node that runs out sends messages
share = [args.codes // args.nodes] * args.nodes
msgs2, sold2, transfers = 0, 0, 0
for n in attempts:
    if share[n] == 0:
        donor = max(range(args.nodes), key=lambda i: share[i])
        msgs2 += 2                                        # ask the richest node, it replies with half
        if share[donor] > 0:
            give = (share[donor] + 1) // 2
            share[donor] -= give; share[n] += give; transfers += 1
    if share[n] > 0: share[n] -= 1; sold2 += 1

print(f"{args.demand:,} purchase attempts, skewed across {args.nodes} nodes, {args.codes:,} codes")
print(f"coordinator: {msgs1:,} messages, {sold1:,} codes sold")
print(f"escrow:      {msgs2:,} messages ({transfers} transfers), {sold2:,} codes sold (never more than {args.codes:,})")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"versions": [
            {"version": "one coordinator", "messages": msgs1, "sold": sold1},
            {"version": "escrow", "messages": msgs2, "sold": sold2}],
            "summary": [{"msgs1": msgs1, "msgs2": msgs2, "transfers": transfers, "sold2": sold2, "codes": args.codes}]}}, f)