# /// script
# requires-python = ">=3.10"
# dependencies = ["zstandard"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Does a trained dictionary rescue small messages?

Small JSON messages compressed one at a time with Zstd, without and with a
dictionary trained on 10,000 samples; then one message decompressed with the
wrong dictionary, and with none.
"""
import argparse, json, random
import zstandard as zstd

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--messages", type=int, default=20_000, help="Messages (half train the dictionary).")
p.add_argument("--dict-kb", type=int, default=16, help="Dictionary size, KB.")
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)
def msg(i):
    return json.dumps({"parcel_id": f"P{random.randint(0, 10**7)}", "scanner_id": f"S{random.randint(1, 900)}",
                       "status": random.choice(["SCANNED", "LOADED", "DELIVERED", "DELAYED"]),
                       "depot": random.choice(["RTM", "LYN", "HAM", "WAW"]), "weight_grams": random.randint(100, 30000)}).encode()
samples = [msg(i) for i in range(args.messages)]
half = args.messages // 2
d = zstd.train_dictionary(args.dict_kb * 1024, samples[:half])
plain, withd = zstd.ZstdCompressor(level=3), zstd.ZstdCompressor(level=3, dict_data=d)
test = samples[half:]; raw = sum(len(m) for m in test)
r_plain = raw / sum(len(plain.compress(m)) for m in test); r_dict = raw / sum(len(withd.compress(m)) for m in test)
print(f"no dictionary : {r_plain:.2f}x   (average message {raw / len(test):.0f} bytes)")
print(f"dictionary    : {r_dict:.2f}x")
blob = withd.compress(test[0])
other = zstd.train_dictionary(args.dict_kb * 1024, [msg(i) for i in range(half)])
outcomes = {}
for label, dec in (("the right dictionary", zstd.ZstdDecompressor(dict_data=d)), ("a different dictionary", zstd.ZstdDecompressor(dict_data=other)),
                   ("no dictionary", zstd.ZstdDecompressor())):
    try: outcomes[label] = "ok" if dec.decompress(blob) == test[0] else "wrong bytes"
    except zstd.ZstdError as e: outcomes[label] = f"error: {str(e)[:60]}"
    print(f"decompress with {label}: {outcomes[label]}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"summary": [{"plain": r_plain, "dict": r_dict, "avg_bytes": raw / len(test)}],
                   "decode": [{"decoder": k, "outcome": v} for k, v in outcomes.items()]}}, f)