"""Build the GPT-U-20M token shards: 2.68B tokens of the 45/35/20 mix plus 2M-token holdouts per domain.

Stages (each resumable; run through scripts/run_prepare.sh, which restarts the process between stages):
  1. holdout   data/tokens/val_<domain>.bin from holdout source files that never feed training or the tokenizer
  2. streams   each sub-stream (dclm, edu, code_py, code_other) is streamed from the Hub on its own and tokenized
               until its token budget (share x total, +0.2% slack) -> data/tokens/streams/<name>.bin
  3. assemble  documents of the four stream files are interleaved in a seeded random order -> train_XXX.bin
  4. validate  sizes, token ids, shard overlaps, holdouts; then the stream files are deleted

Why streams are tokenized one at a time: holding the four Hub streams open together needs 7-8 GB with spikes past
11 GB, and on this machine memory pressure plus disk I/O has bugchecked Windows. One stream at a time peaks at ~2.5 GB.
The interleave is still per document with proportions fixed for the whole run (a seeded permutation of all documents),
so there are no domain blocks, and because every stream holds exactly its token share the realized shares are exact.

Shard format: uint16 token stream; every document is prefixed with <|endoftext|> (id 0); no padding, nothing
truncated. Shard k holds blocks*1024+1 tokens and its last token is repeated as the first token of shard k+1, so
block j of a shard is tokens[j*1024 : j*1024+1025].
"""
import argparse
import copy
import json
import multiprocessing as mp
import os
import random
import sys
import threading
import time
from pathlib import Path

os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")  # parallelism comes from the process pool
os.environ.setdefault("ARROW_DEFAULT_MEMORY_POOL", "system")  # glibc malloc gives arrow buffers back; jemalloc hoards

import numpy as np
from tokenizers import Tokenizer

from sources import SEED, STREAM_DOMAIN, STREAM_SHARE, STREAMS, file_splits, open_stream

ROOT = Path(__file__).resolve().parents[1]
TOKENIZER_PATH = ROOT / "tokenizer" / "tokenizer.json"
SPLITS_PATH = ROOT / "data" / "file_splits.json"
VAL_DIR = ROOT / "data" / "tokens"

SEQ = 1024
STEPS, SEQS_PER_STEP = 19_840, 128
MARGIN_STEPS = 3 * 200  # room for the NaN guard's data skips (max 3 events x 200 steps) without repeating data
FULL_BLOCKS = (STEPS + MARGIN_STEPS) * SEQS_PER_STEP
STREAM_SLACK = 1.002  # each stream collects 0.2% over its share so assembly never runs short
HOLDOUT_TOKENS = 2_000_000
EVAL_WINDOWS = 200 * 8  # evaluate.py: 200 batches x 8 sequences
BATCH_DOCS = 512
BATCH_CHARS = 8_000_000  # also cap a batch by size: a few huge code files must not balloon a tokenizer worker
CKPT_EVERY = 20_000_000
READ_CHUNK = 16_000_000  # tokens per sequential read while assembling (DrvFs-friendly)
NUM_PROC = max(1, (os.cpu_count() or 3) - 2)
EOT_ID = 0
DOMAINS = ("dclm", "edu", "code")
SRC_IDX = {name: i for i, name in enumerate(STREAMS)}
RESTART_EXIT_CODE = 75  # run_prepare.sh starts a fresh process: next stage, or resume after the memory guard
MIN_AVAILABLE_GB = 3.0  # checked after every batch: checkpoint, then restart
HARD_MIN_AVAILABLE_GB = 1.5  # checked twice a second: restart at once from the last periodic checkpoint


def log(msg: str) -> None:
    print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)


def write_json(path: Path, obj: dict) -> None:
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(json.dumps(obj, indent=2))
    os.replace(tmp, path)


def load_tokenizer() -> Tokenizer:
    tok = Tokenizer.from_file(str(TOKENIZER_PATH))
    tok.encode_special_tokens = True  # a literal "<|endoftext|>" inside a document is plain text
    return tok


# ---------------------------------------------------------------- memory guard
def rss_gb(pid: int | str = "self") -> float:
    try:
        with open(f"/proc/{pid}/status") as f:
            return next(int(line.split()[1]) for line in f if line.startswith("VmRSS")) / 2**20
    except (OSError, StopIteration):  # the process is gone
        return 0.0


def tree_rss_gb() -> float:
    """This process plus its tokenizer workers."""
    return rss_gb() + sum(rss_gb(child.pid) for child in mp.active_children())


def available_gb() -> float:
    with open("/proc/meminfo") as f:
        return next(int(line.split()[1]) for line in f if line.startswith("MemAvailable")) / 2**20


def start_memory_watchdog() -> None:
    """Last line of defence for spikes between batch checks (or during a resume's fast-forward, which writes
    nothing): kill the workers and exit at once; the restart redoes at most CKPT_EVERY tokens."""

    def watch() -> None:
        while True:
            available = available_gb()
            if available < HARD_MIN_AVAILABLE_GB:
                log(f"memory watchdog: WSL available {available:.2f} GB -> hard restart from the last checkpoint")
                for child in mp.active_children():
                    child.kill()
                sys.stdout.flush()
                os._exit(RESTART_EXIT_CODE)
            time.sleep(0.5)

    threading.Thread(target=watch, daemon=True, name="memory-watchdog").start()


# ---------------------------------------------------------------- tokenizer workers (spawned processes)
_tok = None


def _init_worker() -> None:
    global _tok
    _tok = load_tokenizer()


def _encode(job):
    b, texts = job
    t0 = time.perf_counter()
    encs = _tok.encode_batch_fast(texts)
    lens = np.fromiter((len(e.ids) + 1 for e in encs), dtype=np.int64, count=len(encs))
    ids = np.empty(int(lens.sum()), dtype=np.uint16)
    pos = 0
    for e, n in zip(encs, lens):
        ids[pos] = EOT_ID
        ids[pos + 1:pos + n] = e.ids
        pos += n
    return b, ids, lens, time.perf_counter() - t0


# ---------------------------------------------------------------- 1. holdout
def token_byte_lengths(tok: Tokenizer) -> np.ndarray:
    """UTF-8 bytes per token id. Each byte-level BPE character stands for exactly one byte; the special token is 0."""
    lengths = np.zeros(tok.get_vocab_size(), dtype=np.int64)
    for token, idx in tok.get_vocab().items():
        if idx != EOT_ID:
            lengths[idx] = len(token)
    return lengths


def build_holdout(splits: dict, tok: Tokenizer) -> dict:
    """2M tokens per domain from the holdout files; the code holdout holds the same 27/73 Python split as training."""
    VAL_DIR.mkdir(parents=True, exist_ok=True)
    byte_len = token_byte_lengths(tok)
    index_path = VAL_DIR / "index.json"
    val = json.loads(index_path.read_text()).get("val", {}) if index_path.exists() else {}
    need = HOLDOUT_TOKENS + 1
    for domain in DOMAINS:
        path = VAL_DIR / f"val_{domain}.bin"
        if domain in val and path.exists():
            continue
        names = [n for n in STREAMS if STREAM_DOMAIN[n] == domain]
        docs = []
        for name in names:
            budget = need * STREAM_SHARE[name] / sum(STREAM_SHARE[n] for n in names)
            got = 0
            for row in open_stream(name, splits[domain]["holdout"]):
                ids = np.array([EOT_ID] + tok.encode(row["text"]).ids, dtype=np.uint16)
                docs.append((name, ids))
                got += len(ids)
                if got >= budget:
                    break
        random.Random(SEED).shuffle(docs)
        chunks, by_stream, n = [], dict.fromkeys(names, 0), 0
        for name, ids in docs:
            take = min(len(ids), need - n)
            chunks.append(ids[:take])
            by_stream[name] += take
            n += take
            if n == need:
                break
        arr = np.concatenate(chunks)
        assert len(arr) == need, f"holdout {domain}: only {len(arr):,} tokens"
        targets = arr[1:]
        tmp = path.with_suffix(".bin.tmp")
        arr.tofile(tmp)
        os.replace(tmp, path)
        val[domain] = {
            "file": path.name, "tokens": int(len(arr)), "targets": int(len(targets)), "docs": len(chunks),
            "bytes": int(byte_len[targets].sum()),
            "eval_windows": EVAL_WINDOWS, "eval_targets": EVAL_WINDOWS * SEQ,
            "eval_bytes": int(byte_len[targets[:EVAL_WINDOWS * SEQ]].sum()),
            "tokens_by_stream": by_stream, "source_files": splits[domain]["holdout"],
        }
        log(f"holdout {domain}: {len(arr):,} tokens, {len(chunks):,} docs, {val[domain]['bytes']:,} bytes "
            f"({val[domain]['bytes'] / len(targets):.3f} bytes/token), by stream {by_stream}")
    index = json.loads(index_path.read_text()) if index_path.exists() else {}
    index["val"] = val
    write_json(index_path, index)
    return val


# ---------------------------------------------------------------- 2. one sub-stream -> token file
def tokenize_stream(name: str, files: list[str], budget: int, max_rss_gb: float, stream_dir: Path) -> dict:
    """Stream one sub-stream from the Hub and append its EOT-prefixed documents to <stream_dir>/<name>.bin until the
    token budget is reached (whole documents only). Checkpoints every CKPT_EVERY tokens."""
    stream_dir.mkdir(parents=True, exist_ok=True)
    out_path, meta_path = stream_dir / f"{name}.bin", stream_dir / f"{name}.json"
    tmp_path, ckpt_path = stream_dir / f"{name}.bin.tmp", stream_dir / f"{name}.state.json"
    state = json.loads(ckpt_path.read_text()) if ckpt_path.exists() else None
    ds = open_stream(name, files)
    if state:
        ds.load_state_dict(state["stream"])
    tokens, docs = (state["tokens"], state["docs"]) if state else (0, 0)
    f = open(tmp_path, "r+b" if state and tmp_path.exists() else "wb")
    f.truncate(tokens * 2)
    f.seek(tokens * 2)
    log(f"{name}: budget {budget / 1e6:,.1f}M tokens" + (f", resuming at {tokens / 1e6:,.1f}M" if state else ""))

    stop = threading.Event()
    sem = threading.Semaphore(2 * NUM_PROC)  # keeps the pool from pulling the whole stream into memory
    stream_states = {}

    def jobs():
        texts, n_chars, b = [], 0, 0
        for row in ds:
            if not row["text"]:
                continue
            texts.append(row["text"])
            n_chars += len(row["text"])
            if len(texts) == BATCH_DOCS or n_chars >= BATCH_CHARS:
                sem.acquire()
                if stop.is_set():
                    return
                stream_states[b] = copy.deepcopy(ds.state_dict())
                yield b, texts
                texts, n_chars, b = [], 0, b + 1

    def checkpoint(stream_state: dict) -> None:
        f.flush()
        os.fsync(f.fileno())
        write_json(ckpt_path, {"stream": stream_state, "tokens": tokens, "docs": docs})

    start_tokens, enc_secs, t0 = tokens, 0.0, time.time()
    next_ckpt, next_log = tokens + CKPT_EVERY, tokens + 100_000_000
    done = False
    with mp.get_context("spawn").Pool(NUM_PROC, initializer=_init_worker) as pool:
        for b, ids, lens, secs in pool.imap(_encode, jobs()):
            sem.release()
            stream_state = stream_states.pop(b)
            if tokens + ids.size >= budget:  # stop after the document that crosses the budget
                cut = int(np.searchsorted(np.cumsum(lens), budget - tokens)) + 1
                ids, lens = ids[:int(lens[:cut].sum())], lens[:cut]
                done = True
            f.write(ids.tobytes())
            tokens += int(ids.size)
            docs += len(lens)
            enc_secs += secs
            if done:
                break
            if tokens >= next_ckpt:
                checkpoint(stream_state)
                next_ckpt += CKPT_EVERY
            if tree_rss_gb() > max_rss_gb or available_gb() < MIN_AVAILABLE_GB:
                checkpoint(stream_state)
                log(f"memory guard: prep rss {tree_rss_gb():.1f} GB, WSL available {available_gb():.1f} GB at "
                    f"{tokens / 1e6:,.0f}M {name} tokens -> checkpoint written, restarting")
                stop.set()
                for _ in range(4 * NUM_PROC):
                    sem.release()
                pool.terminate()
                sys.stdout.flush()
                os._exit(RESTART_EXIT_CODE)  # streaming threads can stall a normal interpreter exit
            if tokens >= next_log:
                rate = (tokens - start_tokens) / (time.time() - t0)
                log(f"{name}: {tokens / 1e6:,.0f}M / {budget / 1e6:,.0f}M tokens | {rate:,.0f} tok/s end-to-end | "
                    f"tokenizer capacity {(tokens - start_tokens) / enc_secs * NUM_PROC:,.0f} tok/s | "
                    f"ETA {(budget - tokens) / rate / 60:.0f} min | prep rss {tree_rss_gb():.1f} GB")
                next_log += 100_000_000
        stop.set()
        for _ in range(4 * NUM_PROC):
            sem.release()
        pool.terminate()
    if not done:
        raise RuntimeError(f"{name}: source files ran out at {tokens:,} tokens (budget {budget:,})")

    f.flush()
    os.fsync(f.fileno())
    f.close()
    os.replace(tmp_path, out_path)
    secs = time.time() - t0
    meta = {"name": name, "tokens": tokens, "docs": docs, "budget": budget, "seconds_last_run": round(secs),
            "tok_s_end_to_end": round((tokens - start_tokens) / secs),
            "tok_s_tokenizer_capacity": round((tokens - start_tokens) / enc_secs * NUM_PROC), "num_proc": NUM_PROC}
    write_json(meta_path, meta)
    ckpt_path.unlink(missing_ok=True)
    log(f"{name}: done, {tokens:,} tokens, {docs:,} docs, {meta['tok_s_end_to_end']:,} tok/s end-to-end, "
        f"tokenizer capacity {meta['tok_s_tokenizer_capacity']:,} tok/s")
    return meta


# ---------------------------------------------------------------- 3. assemble shards
class ShardWriter:
    def __init__(self, out_dir: Path, total_blocks: int, shard_tokens: int):
        self.out_dir = out_dir
        self.total_blocks = total_blocks
        self.shard_blocks = shard_tokens // SEQ
        self.shards = []
        self.blocks_done = 0
        self.cur = 0
        self.cur_counts = np.zeros(len(STREAMS), dtype=np.int64)
        self._open()

    def _target(self) -> tuple[int, int]:
        blocks = min(self.shard_blocks, self.total_blocks - self.blocks_done)
        return blocks, blocks * SEQ + 1

    def _open(self) -> None:
        self.tmp_path = self.out_dir / f"train_{len(self.shards):03d}.bin.tmp"
        self.f = open(self.tmp_path, "wb")

    def add(self, ids: np.ndarray, src_per_token: np.ndarray) -> bool:
        """Append tokens; returns True once all blocks are written."""
        pos = 0
        while pos < len(ids):
            blocks, target = self._target()
            take = min(target - self.cur, len(ids) - pos)
            self.f.write(ids[pos:pos + take].tobytes())
            self.cur_counts += np.bincount(src_per_token[pos:pos + take], minlength=len(STREAMS))
            self.cur += take
            pos += take
            if self.cur == target:
                self._close_shard(blocks, int(ids[pos - 1]))
                if self.blocks_done == self.total_blocks:
                    return True
        return False

    def _close_shard(self, blocks: int, last_token: int) -> None:
        self.f.flush()
        os.fsync(self.f.fileno())
        self.f.close()
        name = self.tmp_path.name.removesuffix(".tmp")
        os.replace(self.tmp_path, self.out_dir / name)
        self.shards.append({"file": name, "tokens": self.cur, "blocks": blocks,
                            "tokens_by_stream": dict(zip(STREAMS, self.cur_counts.tolist()))})
        self.blocks_done += blocks
        log(f"shard {name}: {blocks:,} blocks, {self.cur:,} tokens")
        if self.blocks_done < self.total_blocks:
            self._open()
            self.f.write(np.array([last_token], dtype=np.uint16).tobytes())  # overlap token, counted in prior shard
            self.cur = 1
            self.cur_counts = np.zeros(len(STREAMS), dtype=np.int64)


def doc_lengths(path: Path, n_tokens: int) -> np.ndarray:
    """Every document starts with <|endoftext|> and id 0 occurs nowhere else, so its positions delimit documents."""
    starts = []
    with open(path, "rb") as f:
        for off in range(0, n_tokens, READ_CHUNK):
            chunk = np.frombuffer(f.read(min(READ_CHUNK, n_tokens - off) * 2), dtype=np.uint16)
            starts.append(np.flatnonzero(chunk == EOT_ID) + off)
    starts = np.concatenate(starts)
    assert starts[0] == 0, f"{path.name} does not start with <|endoftext|>"
    return np.diff(np.append(starts, n_tokens))


class SequentialReader:
    def __init__(self, path: Path):
        self.f = open(path, "rb")
        self.buf = np.empty(0, dtype=np.uint16)
        self.pos = 0

    def take(self, n: int) -> np.ndarray:
        while len(self.buf) - self.pos < n:
            more = np.frombuffer(self.f.read(READ_CHUNK * 2), dtype=np.uint16)
            if more.size == 0:
                raise EOFError(self.f.name)
            self.buf = np.concatenate([self.buf[self.pos:], more])
            self.pos = 0
        out = self.buf[self.pos:self.pos + n]
        self.pos += n
        return out


def assemble(metas: dict, out_dir: Path, total_blocks: int, shard_tokens: int) -> dict:
    """Interleave the documents of the four stream files in a seeded random order (per-document sampling at
    proportions fixed for the whole run) and pack them into shards. Fast and local, so it always starts over."""
    for p in out_dir.glob("train_*.bin*"):
        p.unlink()
    stream_dir = out_dir / "streams"
    lengths = {}
    for name in STREAMS:
        lengths[name] = doc_lengths(stream_dir / f"{name}.bin", metas[name]["tokens"])
        assert len(lengths[name]) == metas[name]["docs"], f"{name}: document count mismatch"
    order = np.concatenate([np.full(len(lengths[n]), SRC_IDX[n], dtype=np.uint8) for n in STREAMS])
    np.random.default_rng(SEED).shuffle(order)
    readers = {n: SequentialReader(stream_dir / f"{n}.bin") for n in STREAMS}
    used = dict.fromkeys(STREAMS, 0)
    writer = ShardWriter(out_dir, total_blocks, shard_tokens)
    t0, done = time.time(), False
    chunks, srcs, doc_lens, pending = [], [], [], 0
    for s in order.tolist():
        name = STREAMS[s]
        n = int(lengths[name][used[name]])
        used[name] += 1
        chunks.append(readers[name].take(n))
        srcs.append(s)
        doc_lens.append(n)
        pending += n
        if pending >= 4_000_000:
            done = writer.add(np.concatenate(chunks), np.repeat(np.array(srcs, dtype=np.uint8), doc_lens))
            chunks, srcs, doc_lens, pending = [], [], [], 0
            if done:
                break
    if not done and chunks:
        done = writer.add(np.concatenate(chunks), np.repeat(np.array(srcs, dtype=np.uint8), doc_lens))
    if not done:
        raise RuntimeError("stream files are too short for the requested number of blocks")
    log(f"assembled {writer.blocks_done:,} blocks in {time.time() - t0:.0f}s; documents used {used}")

    by_stream = {n: sum(s["tokens_by_stream"][n] for s in writer.shards) for n in STREAMS}
    total = sum(by_stream.values())
    by_domain = {d: sum(v for n, v in by_stream.items() if STREAM_DOMAIN[n] == d) for d in DOMAINS}
    index_path = out_dir / "index.json"
    index = json.loads(index_path.read_text()) if index_path.exists() else {}
    for stale in ("probabilities", "token_caps", "skipped_docs_after_cap"):
        index.pop(stale, None)
    index.update({
        "tokenizer": str(TOKENIZER_PATH.relative_to(ROOT)), "dtype": "uint16", "seq_len": SEQ, "block": SEQ + 1,
        "stride": SEQ, "eot_id": EOT_ID, "seed": SEED,
        "mix": {"method": "each sub-stream tokenized to its token budget, documents interleaved in a seeded random "
                          "order (per-document sampling at fixed proportions)",
                "target_share_by_stream": STREAM_SHARE, "streams": metas, "documents_used": used},
        "blocks": writer.blocks_done, "train_steps": STEPS, "seqs_per_step": SEQS_PER_STEP,
        "margin_steps": MARGIN_STEPS if total_blocks == FULL_BLOCKS else 0,
        "shards": writer.shards,
        "total": {
            "tokens": total, "tokens_by_stream": by_stream, "tokens_by_domain": by_domain,
            "share_by_stream": {n: by_stream[n] / total for n in STREAMS},
            "share_by_domain": {d: by_domain[d] / total for d in DOMAINS},
            "python_share_of_code": by_stream["code_py"] / by_domain["code"],
        },
        "throughput": {
            "tok_s_end_to_end_by_stream": {n: metas[n]["tok_s_end_to_end"] for n in STREAMS},
            "tok_s_tokenizer_capacity_by_stream": {n: metas[n]["tok_s_tokenizer_capacity"] for n in STREAMS},
            "num_proc": NUM_PROC,
        },
    })
    write_json(index_path, index)
    return index


# ---------------------------------------------------------------- 4. validation
def validate(out_dir: Path) -> None:
    index = json.loads((out_dir / "index.json").read_text())
    tok = Tokenizer.from_file(str(TOKENIZER_PATH))
    vocab = tok.get_vocab_size()
    rng = np.random.default_rng(0)
    prev_last = None
    for s in index["shards"]:
        arr = np.fromfile(out_dir / s["file"], dtype=np.uint16)  # sequential read; mmap over DrvFs is slow
        assert len(arr) == s["tokens"] == s["blocks"] * SEQ + 1, f"{s['file']}: size mismatch"
        assert int(arr.max()) < vocab, f"{s['file']}: token id out of range"
        if prev_last is not None:
            assert int(arr[0]) == prev_last, f"{s['file']}: overlap token mismatch"
        prev_last = int(arr[-1])
        eot = int((arr == EOT_ID).sum())
        j = int(rng.integers(s["blocks"]))
        text = tok.decode(arr[j * SEQ:j * SEQ + SEQ + 1].tolist(), skip_special_tokens=False)
        log(f"{s['file']} ok: {s['blocks']:,} blocks, {eot:,} docs (mean {len(arr) / max(eot, 1):,.0f} tokens/doc) | "
            f"block {j}: {text[:160]!r}")
        del arr
    for d, v in index.get("val", {}).items():
        arr = np.fromfile(VAL_DIR / v["file"], dtype=np.uint16)
        assert len(arr) == v["tokens"] and int(arr.max()) < vocab
        log(f"val_{d} ok: {len(arr):,} tokens, {v['bytes']:,} bytes, by stream {v['tokens_by_stream']}")
    t = index["total"]
    log(f"blocks {index['blocks']:,} (training needs {STEPS * SEQS_PER_STEP:,}) | shares by domain: "
        + ", ".join(f"{d} {100 * x:.2f}%" for d, x in t["share_by_domain"].items())
        + f" | python in code {100 * t['python_share_of_code']:.1f}%")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", type=Path, default=VAL_DIR)
    ap.add_argument("--blocks", type=int, default=FULL_BLOCKS)
    ap.add_argument("--shard-tokens", type=int, default=500_000_000)
    ap.add_argument("--validate", action="store_true", help="only validate an existing output directory")
    ap.add_argument("--max-rss-gb", type=float, default=8.0,
                    help="memory guard: checkpoint and restart when this process plus its workers exceed it")
    ap.add_argument("--keep-streams", action="store_true", help="keep data/tokens/streams after validation")
    args = ap.parse_args()
    out_dir = args.out if args.out.is_absolute() else ROOT / args.out

    if args.validate:
        validate(out_dir)
        return

    out_dir.mkdir(parents=True, exist_ok=True)
    start_memory_watchdog()
    tok = load_tokenizer()
    splits = file_splits(SPLITS_PATH)
    build_holdout(splits, tok)

    target_tokens = args.blocks * SEQ + 1
    stream_dir = out_dir / "streams"
    metas = {}
    for name in STREAMS:
        meta_path = stream_dir / f"{name}.json"
        if meta_path.exists():
            metas[name] = json.loads(meta_path.read_text())
            continue
        budget = int(STREAM_SHARE[name] * target_tokens * STREAM_SLACK) + 1
        tokenize_stream(name, splits[STREAM_DOMAIN[name]]["train"], budget, args.max_rss_gb, stream_dir)
        sys.stdout.flush()
        os._exit(RESTART_EXIT_CODE)  # next stream in a fresh process, with a fresh heap

    index = assemble(metas, out_dir, args.blocks, args.shard_tokens)
    log(json.dumps({k: index[k] for k in ("total", "throughput")}, indent=2))
    validate(out_dir)
    if not args.keep_streams:
        for p in stream_dir.glob("*"):
            p.unlink()
        stream_dir.rmdir()
        log("stream files deleted")
    sys.stdout.flush()
    os._exit(0)  # streaming threads can stall a normal interpreter exit


if __name__ == "__main__":
    main()