# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Find five bad rows among a million with a Merkle tree.

Two lists of 1,000,000 (key, row) pairs, five rows changed in the second.
A Merkle tree with branching factor 16 over key-sorted rows; compare from the
root down and count the hashes exchanged. Then a table hash of every row
added twice: XOR against addition modulo 2^64.
"""
import argparse, hashlib, json, random

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rows", type=int, default=1_000_000, help="Rows per side.")
p.add_argument("--changed", type=int, default=5, help="Rows changed on one side.")
p.add_argument("--fanout", type=int, default=16, help="Branching factor of the tree.")
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)
H = lambda b: hashlib.blake2b(b, digest_size=8).digest()

left = [f"{k:08d}|row {k} amount {k % 997}".encode() for k in range(args.rows)]
right = list(left)
bad = sorted(random.sample(range(args.rows), args.changed))
for k in bad: right[k] = right[k] + b" (edited)"

def tree(rows):
    levels = [[H(r) for r in rows]]
    while len(levels[-1]) > 1:
        lv = levels[-1]
        levels.append([H(b"".join(lv[i:i + args.fanout])) for i in range(0, len(lv), args.fanout)])
    return levels                                            # levels[-1] is the root

L, R = tree(left), tree(right)
exchanged, frontier = 1, [0]                                 # compare the roots first
for depth in range(len(L) - 1, 0, -1):
    nxt = []
    for i in frontier:
        if L[depth][i] == R[depth][i]: continue
        kids = range(i * args.fanout, min((i + 1) * args.fanout, len(L[depth - 1])))
        exchanged += len(kids); nxt += [k for k in kids if L[depth - 1][k] != R[depth - 1][k]]
    frontier = nxt
found = sorted(frontier)
print(f"{args.rows:,} rows, {args.changed} changed: found rows {found} after exchanging {exchanged:,} hashes "
      f"(a full diff moves {args.rows:,} rows)")

h64 = lambda r: int.from_bytes(H(r), "big")
sample = left[:1000]
xor = 0
for r in sample + sample: xor ^= h64(r)
add_twice = sum(h64(r) for r in sample + sample) % (1 << 64); add_once = sum(h64(r) for r in sample) % (1 << 64)
print(f"every row twice: XOR table hash = {xor} (same as an empty table); "
      f"addition mod 2^64 differs from the single copy: {add_twice != add_once}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"summary": [{"rows": args.rows, "changed": args.changed, "found": len(found), "exchanged": exchanged,
                   "xor_twice": xor, "sum_differs": "yes" if add_twice != add_once else "no"}]}}, f)