# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Retention window against storage (a model of a table format).

A model of an open table format's file bookkeeping, not Iceberg or Delta: a
table of 100 data files; every simulated day 10% of them are rewritten (new
files written, old ones dropped from the current snapshot). Expiry then
vacuum keep every file referenced by a snapshot inside the retention window.
Storage after 10 days for windows of 1, 3 and 7 days, against
size × (1 + f·R). Then a query that started on an old snapshot reads after a
vacuum with a shorter window.
"""
import argparse, json, random

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--files", type=int, default=100, help="Data files in the table.")
p.add_argument("--file-mb", type=float, default=128, help="MB per file.")
p.add_argument("--f", type=float, default=0.1, help="Share of the table rewritten per day.")
p.add_argument("--days", type=int, default=10, help="Simulated days.")
p.add_argument("--windows", default="1,3,7", help="Retention windows (days).")
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()

def simulate(R):
    rng = random.Random(args.seed); nxt = args.files
    current = list(range(args.files)); snapshots = {0: set(current)}; stored = set(current)
    for day in range(1, args.days + 1):
        out = rng.sample(range(len(current)), int(args.f * args.files))
        for i in out: current[i] = nxt; stored.add(nxt); nxt += 1          # rewrite: new files in, old ones out of the snapshot
        snapshots[day] = set(current)
        snapshots = {d: s for d, s in snapshots.items() if d >= day - R}     # expire snapshots older than the window
        stored = set().union(*snapshots.values())                          # vacuum: delete files no kept snapshot references
    return len(stored) * args.file_mb, snapshots
base = args.files * args.file_mb; rows = []
for R in map(int, args.windows.split(",")):
    mb, _ = simulate(R)
    rows.append({"window_days": R, "storage_gb": mb / 1024, "formula_gb": base * (1 + args.f * R) / 1024, "table_gb": base / 1024})
_, snaps = simulate(1)
query_snapshot = args.days - 3                                              # a long query started on the snapshot of day 7
needed = None
rng = random.Random(args.seed); current = list(range(args.files)); hist = {0: set(current)}; nxt = args.files
for day in range(1, args.days + 1):
    for i in rng.sample(range(len(current)), int(args.f * args.files)): current[i] = nxt; nxt += 1
    hist[day] = set(current)
needed = hist[query_snapshot]; kept = set().union(*snaps.values())
missing = len(needed - kept)
summary = {"query_snapshot_day": query_snapshot, "files_needed": len(needed), "files_missing": missing,
           "outcome": f"file not found ({missing} of {len(needed)} files deleted)" if missing else "ok"}
for r in rows: print(f"retention {r['window_days']} day(s): storage {r['storage_gb']:.2f} GB (formula {r['formula_gb']:.2f} GB, table {r['table_gb']:.2f} GB)")
print(f"query on the day-{query_snapshot} snapshot after a 1-day vacuum: {summary['outcome']}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"summary": summary, "tables": {"windows": rows}}, f)