# Phase 2 builder: sources -> per-source staged shards -> interleaved mixed shards.
#
# Two stages, because one stage cannot produce a mixed corpus without re-streaming.
#
#   stage  per source: stream -> normalise -> language/quality gate -> dedup -> tokenise
#                      -> write  stage/<key>/shard-####.bin
#   merge  read all staged shards together, emit documents round-robin in proportion to each source's
#          target, and write the final mixed  mix/shard-####.bin
#
# Why the merge exists (memory/ERRORS.md E-016). The first version consumed one source to its target and
# only then started the next, and it flushed a shard when the buffer filled. Since fineweb-edu's target is
# 300M tokens = 37 shards of 8M, the first 37 shards of the mix would have been nothing but web text, then
# textbooks, then math. A shuffle inside the buffer cannot fix that -- the buffer only ever held one
# source. That is a curriculum drift across the first third of the main run, and no LR schedule repairs
# ordering like that. Mixing must happen between sources, not inside them.
#
# Shard layout is shared by both stages and carries a per-document offset table, so document boundaries
# survive to the training reader (E-016's second latent bug: the old format stored a doc *count*, from
# which boundaries cannot be recovered):
#
#   header   : <IIII  = vocab, n_docs, n_tokens, flags(0)          16 B
#   offsets  : uint32 × (n_docs + 1), cumulative, [0] = 0
#   ids      : uint16 × n_tokens
#
# doc i of a shard is ids[offsets[i] : offsets[i+1]].
#
# Constraints (docs/00-platform-notes.md): Kaggle CPU session, 4 vCPU, 30 GiB cgroup, **no swap**,
# 19.5 GB /kaggle/working. Staged (~2.6 GB) + final (~2.6 GB) + hash sets (~35 MB) fit with room. Streaming
# is ~3.3 MB/s per process and tokenising ~661k tok/s on 4 cores, so ingest is serial-per-source by design.
#
# Resumability: state is saved only *after* a file lands, and the dedup hash sets are appended in the same
# transaction -- otherwise a resume forgets what was already emitted and lets duplicates back in across
# the boundary. No credentials in this file.

import argparse
import array
import hashlib
import json
import os
import random
import struct
import sys
import time

# Sibling modules (hubsync, ounce100m_credentials) live beside this file, not in any package, and a job
# started as `python build/build_mix.py` gets only the script's own directory on sys.path -- which is
# build/, so the credentials import inside hubsync fails and the whole run refuses to be interruption-
# safe (E-021's lesson, hit again from the other direction). Add both this directory and its parent.
_HERE = os.path.dirname(os.path.abspath(__file__))
for _p in (_HERE, os.path.dirname(_HERE)):
    if _p not in sys.path:
        sys.path.insert(0, _p)

MAGIC_FLAGS = 0
TOKENS_PER_SHARD = 8_000_000    # 16 MB as uint16; ~140 shards for 1.1B tokens
MERGE_BATCH_DOCS = 32           # granularity of the round-robin; ~20k tokens, so a 262k batch spans ~13
VAL_STRIDE = 50                 # every 50th document is held out ~= 2% of tokens, sampled across
                                # sources rather than whole documents, so validation PPL is honest
VAL_TOKENS_PER_SHARD = 2_000_000
ROW_TOKEN_CAP = 16_384          # one document must never dominate a shard
SHUFFLE_SEED = 20260919         # recorded so a rebuild is reproducible
HASHES_FILE = "seen_hashes.bin"
KEYS_FILE = "seen_keys.bin"
STATE_NAME = "build_state.json"

# Corrected config and column names, from the Phase 2 inventory (docs/02-mix-plan.md §9). The names in
# §2 came from dataset cards, and five of them did not exist.
SOURCES = [
    dict(repo="HuggingFaceFW/fineweb-edu", config="sample-10BT", col="text", tokens=340_000_000,
         min_score=4, why="quality-scored general English"),
    dict(repo="HuggingFaceFW/finepdfs-edu", config="eng_Latn", col="text", tokens=180_000_000,
         en_filter=True, why="long-form educational PDFs; needs own en filter"),
    dict(repo="HuggingFaceTB/cosmopedia", config="stanford", col="text", tokens=60_000_000,
         why="textbook expository"),
    dict(repo="HuggingFaceTB/cosmopedia", config="openstax", col="text", tokens=30_000_000,
         why="textbook science"),
    dict(repo="HuggingFaceTB/cosmopedia", config="khanacademy", col="text", tokens=10_000_000,
         why="worked steps"),
    dict(repo="HuggingFaceTB/cosmopedia", config="auto_math_text", col="text", tokens=30_000_000,
         why="synthetic math exposition"),
    dict(repo="HuggingFaceTB/cosmopedia", config="wikihow", col="text", tokens=30_000_000,
         why="procedural / physical-world steps"),
    dict(repo="HuggingFaceFW/finewiki", config="en", col="rewritten", tokens=80_000_000,
         why="rewritten encyclopedic prose"),
    dict(repo="omarkamali/wikipedia-monthly", config="20260101.en", col="text", tokens=60_000_000,
         en_filter=True, why="current encyclopedic breadth"),
    dict(repo="HuggingFaceTB/finemath", config="finemath-4plus", col="text", tokens=160_000_000,
         why="math; upstream-decontaminated against GSM8K/MMLU/ARC"),
    dict(repo="open-web-math/open-web-math", config="default", col="text", tokens=70_000_000,
         en_filter=True, why="forum/webbook math prose"),
    dict(repo="HuggingFaceFW/finephrase", config="tutorial", col="text", tokens=60_000_000,
         en_filter=True, why="stepwise tutorial register"),
    dict(repo="HuggingFaceFW/finephrase", config="faq", col="text", tokens=30_000_000,
         en_filter=True, why="FAQ register"),
    dict(repo="HuggingFaceFW/finephrase", config="table", col="text", tokens=30_000_000,
         en_filter=True, why="tabular->prose"),
    dict(repo="SimpleStories/SimpleStories", config="default", col="story", tokens=40_000_000,
         why="long-range simple narrative coherence"),
    dict(repo="common-pile/arxiv_abstracts", config="default", col="text", tokens=20_000_000,
         en_filter=True, why="CC0 scientific abstracts"),
    # No code source yet: stack-v3-train has only a `default` config and its rows are whole repos with a
    # nested files struct, so it needs its own probe (docs/02-mix-plan.md §9.3).
    # common-pile/libretexts dropped: `text` holds URLs, 93% of sampled rows failed the en test.
]

STOP = (" the ", " and ", " of ", " to ", " that ", " is ", " in ", " for ", " with ", " as ")
KEY_COLS = ("url", "id", "file_path", "doc_id", "generation_id", "repo_name")


def src_key(src):
    return f"{src['repo'].split('/')[1]}__{src['config']}"


def normalize(s):
    # Whitespace-collapse only: aggressive folding collides genuinely different documents (tables, code),
    # and a false merge is worse than a missed near-duplicate.
    return " ".join(s.split())


def looks_english(s):
    low = s[:600].lower()
    return sum(w in low for w in STOP) >= 3


def h64(s):
    """64-bit digest of a string. digest()[:8] is 8 BYTES; hexdigest()[:16] is 16 *characters*
    and broke the first version (E-015)."""
    return struct.unpack("<Q", hashlib.sha256(s.encode("utf-8")).digest()[:8])[0]


def key_of(row, cols):
    for c in KEY_COLS:
        if c in cols:
            v = row.get(c)
            if isinstance(v, str) and v:
                return v
    return None


# --------------------------------------------------------------------- shard I/O
def write_shard(path, vocab, docs):
    """docs: iterable of array('H')/list of token ids. Returns (n_docs, n_tokens, sha256, bytes)."""
    offs = array.array("I", [0])
    ids = array.array("H")
    for d in docs:
        ids.extend(d)
        offs.append(len(ids))
    tmp = path + ".tmp"
    with open(tmp, "wb") as f:
        f.write(struct.pack("<IIII", vocab, len(offs) - 1, len(ids), MAGIC_FLAGS))
        offs.tofile(f)
        ids.tofile(f)
    os.replace(tmp, path)
    digest = hashlib.sha256(open(path, "rb").read()).hexdigest()
    return len(offs) - 1, len(ids), digest, os.path.getsize(path)


def read_shard_header(path):
    with open(path, "rb") as f:
        vocab, n_docs, n_tokens, flags = struct.unpack("<IIII", f.read(16))
    return vocab, n_docs, n_tokens, flags


def iter_shard_docs(path):
    """Yield documents as array('H'). Offsets are read in full (4 B/doc) which is fine at 8M tok/shard."""
    with open(path, "rb") as f:
        vocab, n_docs, n_tokens, _flags = struct.unpack("<IIII", f.read(16))
        offs = array.array("I")
        offs.fromfile(f, n_docs + 1)
        ids = array.array("H")
        ids.fromfile(f, n_tokens)
    for i in range(n_docs):
        yield ids[offs[i]:offs[i + 1]]


# --------------------------------------------------------------------- state
class State:
    def __init__(self, root):
        self.root = root
        self.stage_dir = os.path.join(root, "stage")
        self.mix_dir = os.path.join(root, "mix")
        os.makedirs(self.stage_dir, exist_ok=True)
        os.makedirs(self.mix_dir, exist_ok=True)
        self.path = os.path.join(root, STATE_NAME)
        self.hashes_path = os.path.join(root, HASHES_FILE)
        self.keys_path = os.path.join(root, KEYS_FILE)
        self.d = {"staged": {}, "merged_shards": [], "version": 3}
        self.seen_exact = set()
        self.seen_key = set()
        if os.path.exists(self.path):
            self.d = json.load(open(self.path))
            self.seen_exact = self._load_u64(self.hashes_path)
            self.seen_key = self._load_u64(self.keys_path)
            print(f"resuming: {len(self.d['staged'])} sources staged, "
                  f"{len(self.d['merged_shards'])} mixed shards, "
                  f"{len(self.seen_exact):,} hashes", flush=True)

    @staticmethod
    def _load_u64(p):
        if not os.path.exists(p):
            return set()
        a = array.array("Q")
        with open(p, "rb") as f:
            a.fromfile(f, os.path.getsize(p) // 8)
        return set(a)

    def commit(self, new_hashes, new_keys):
        """Append the dedup sets FIRST, then the state file: if the job dies between the two, the worst
        case is a hash set slightly ahead of the state (a document re-emitted is a duplicate, which the
        next run's dedup catches; a hash set behind the state is a permanent silent duplicate)."""
        if new_hashes:
            with open(self.hashes_path, "ab") as f:
                array.array("Q", new_hashes).tofile(f)
        if new_keys:
            with open(self.keys_path, "ab") as f:
                array.array("Q", new_keys).tofile(f)
        tmp = self.path + ".tmp"
        with open(tmp, "w") as f:
            json.dump(self.d, f, indent=1)
        os.replace(tmp, self.path)

    def staged_tokens(self, key):
        return sum(s["tokens"] for s in self.d["staged"].get(key, {}).get("shards", []))

    def staged_tokens_total(self):
        return sum(v.get("tokens", 0) for v in self.d["staged"].values())

    def merged_tokens(self):
        return sum(s["tokens"] for s in self.d["merged_shards"])


# --------------------------------------------------------------------- stage
def stage(args, st, tokenizer):
    remote = HS.remote_stage_files(args.hub_repo) if (HS and args.hub_repo
                                    and not args.merge_only) else set()
    if args.hub_repo:
        print(f"hub: {len(remote)} staged files already on {args.hub_repo}", flush=True)
    for src in SOURCES:
        key = src_key(src)
        rec = st.d["staged"].get(key, {})
        if rec.get("complete"):
            print(f"stage: skip {key} (complete)", flush=True)
            continue
        if HS and args.hub_repo:
            # A source whose shards AND record are already on the Hub is done; pulling the record is
            # cheaper than re-staging and keeps the merge's proportional weights honest.
            have = [p for p in remote if p.startswith(f"stage/{key}/") and p.endswith(".bin")]
            if have and f"stage/{key}/record.json" in remote:
                try:
                    HS.download_missing(st.root, [key], args.hub_repo, remote=remote)
                    rec_json = json.load(open(os.path.join(st.stage_dir, key, "record.json")))
                    st.d["staged"][key] = {"complete": True, "shards": rec_json["shards"],
                                           "tokens": rec_json["tokens"],
                                           "docs_seen": rec_json["docs_seen"],
                                           "dropped": rec_json["dropped"],
                                           "target": rec_json["tokens"], "restored_from_hub": True}
                    st.commit([], [])
                    print(f"stage: restored {key} from Hub ({rec_json['tokens']:,} tok), skipping",
                          flush=True)
                    continue
                except Exception as e:
                    print(f"!! restore failed for {key}: {type(e).__name__}: {str(e)[:140]}", flush=True)
        want = min(src["tokens"], args.cap_tokens_per_source or src["tokens"])
        if args.budget_tokens and st.staged_tokens_total() >= args.budget_tokens:
            print("stage: overall budget reached", flush=True)
            break
        import datasets

        try:
            ds = datasets.load_dataset(src["repo"], src["config"], split="train", streaming=True)
        except Exception as e:
            print(f"!! stage: {key} could not load: {type(e).__name__}: {str(e)[:200]}", flush=True)
            rec["error"] = f"{type(e).__name__}: {str(e)[:200]}"
            st.d["staged"][key] = rec
            st.commit([], [])
            continue

        got = docs_seen = 0
        drops = {"short_or_low_score": 0, "non_english": 0, "exact_dup": 0, "key_dup": 0,
                 "too_few_tokens": 0, "no_column": 0}
        cols = None
        buf, new_h, new_k = [], [], []
        sd = os.path.join(st.stage_dir, key)
        os.makedirs(sd, exist_ok=True)
        shards = rec.get("shards", [])
        t0 = time.monotonic()
        for row in ds:
            if got >= want:
                break
            if args.sec_per_source and time.monotonic() - t0 > args.sec_per_source:
                break
            if cols is None:
                cols = set(row.keys())
                if src["col"] not in cols:
                    print(f"!! stage: {key} has no column {src['col']!r}; seen {sorted(cols)[:10]}",
                          flush=True)
                    drops["no_column"] = -1
                    break
            raw = row.get(src["col"])
            if not isinstance(raw, str):
                continue
            docs_seen += 1
            s = normalize(raw)
            if len(s) < args.min_chars:
                drops["short_or_low_score"] += 1
                continue
            if src.get("min_score") is not None:
                sc = row.get("int_score", row.get("score"))
                try:
                    if sc is not None and float(sc) < src["min_score"]:
                        drops["short_or_low_score"] += 1
                        continue
                except (TypeError, ValueError):
                    pass
            if src.get("en_filter") and not looks_english(s):
                drops["non_english"] += 1
                continue
            h = h64(s)
            if h in st.seen_exact:
                drops["exact_dup"] += 1
                continue
            u = key_of(row, cols)
            ku = h64(u) if u else None
            if ku is not None and ku in st.seen_key:
                drops["key_dup"] += 1
                continue
            ids = tokenizer.encode(s[:args.max_chars], add_special_tokens=False).ids[:ROW_TOKEN_CAP]
            if len(ids) < 16:
                drops["too_few_tokens"] += 1
                continue
            st.seen_exact.add(h)
            new_h.append(h)
            if ku is not None:
                st.seen_key.add(ku)
                new_k.append(ku)
            buf.append(array.array("H", ids))
            got += len(ids)
            if sum(len(b) for b in buf) >= TOKENS_PER_SHARD:
                idx = len(shards)
                p = os.path.join(sd, f"shard-{idx:04d}.bin")
                nd, nt, dg, bs = write_shard(p, tokenizer.get_vocab_size(), buf)
                shards.append({"file": os.path.relpath(p, st.root), "docs": nd, "tokens": nt,
                               "sha256": dg, "bytes": bs})
                st.seen_exact.update(new_h)
                st.d["staged"][key] = {"complete": False, "shards": shards, "tokens": got,
                                       "docs_seen": docs_seen, "dropped": drops}
                st.commit(new_h, new_k)
                new_h, new_k = [], []
                buf = []
                print(f"  staged {key} shard {idx:04d}: {nt:,} tok [{time.monotonic() - t0:.0f}s]",
                      flush=True)

        if buf:
            idx = len(shards)
            p = os.path.join(sd, f"shard-{idx:04d}.bin")
            nd, nt, dg, bs = write_shard(p, tokenizer.get_vocab_size(), buf)
            shards.append({"file": os.path.relpath(p, st.root), "docs": nd, "tokens": nt,
                           "sha256": dg, "bytes": bs})
            print(f"  staged {key} shard {idx:04d}: {nt:,} tok (tail)", flush=True)

        complete = got >= want
        st.d["staged"][key] = {"complete": bool(complete), "shards": shards, "tokens": got,
                               "docs_seen": docs_seen, "dropped": dict(drops),
                               "target": want, "seconds": round(time.monotonic() - t0, 1)}
        st.commit(new_h, new_k)
        # /kaggle/working is destroyed when the session ends, so a source that took 40 minutes to stage
        # must not exist only there. Published per source, because that is the unit a resume can skip.
        if complete and HS and args.hub_repo:
            with open(os.path.join(st.stage_dir, key, "record.json"), "w") as f:
                json.dump({"key": key, "src": {k: v for k, v in src.items()},
                           "shards": shards, "tokens": got, "docs_seen": docs_seen,
                           "dropped": dict(drops)}, f, indent=1)
            try:
                n = HS.upload_source(st.root, key, args.hub_repo)
                print(f"  published {key} -> {args.hub_repo}/stage/{key} ({n} files)", flush=True)
            except Exception as e:
                print(f"!! publish FAILED for {key}: {type(e).__name__}: {str(e)[:160]} -- "
                      f"this source will have to be re-staged if the session dies", flush=True)
        el = time.monotonic() - t0
        print(f"stage: {key} {got:,} tok of {want:,} from {docs_seen:,} docs; "
              f"dup={drops['exact_dup']} keydup={drops['key_dup']} lang={drops['non_english']} "
              f"short={drops['short_or_low_score']} [{got / max(1e-9, el):,.0f} tok/s]", flush=True)


# --------------------------------------------------------------------- merge
def merge(args, st, tokenizer):
    """Interleave staged documents proportionally to each source's staged token total.

    --exclude-dir applies per-source uint8 masks written by audit_contamination --write-filter: a 1 byte
    means "this document's 13-grams appear in a benchmark train/validation/dev split, do not emit it"
    (§3.3). Masks are indexed by document ordinal within the source, which is stable because both the
    audit and the merge replay the source's own record.json shard order.
    """
    vocab = tokenizer.get_vocab_size()
    masks, dropped, filter_digest = {}, {}, None
    if args.exclude_dir:
        fj = os.path.join(args.exclude_dir, "filter.json")
        if not os.path.exists(fj):
            print(f"FATAL: --exclude-dir {args.exclude_dir} contains no filter.json", flush=True)
            sys.exit(5)
        filter_digest = hashlib.sha256(open(fj, "rb").read()).hexdigest()
        fj_obj = json.load(open(fj))
        if not fj_obj.get("scanned_sources") and not fj_obj.get("sources"):
            print("FATAL: filter.json names no sources at all -- nothing was scanned, so merging with it "
                  "would publish an unfiltered mix under a filter claim", flush=True)
            sys.exit(5)
        if not fj_obj.get("scanned_sources"):
            # The completeness field post-dates some archived masks. It is a proxy; the real invariant is
            # "a mask exists for every source this merge will read", which is checked per-source below.
            print(f"WARNING: filter.json predates the scanned_sources field; trusting its "
                  f"{len(fj_obj.get('sources', {}))} per-source entries because every merged source must "
                  "still carry a mask or the merge refuses", flush=True)
        dropped = fj_obj.get("sources", {})
        for key in sorted(os.listdir(args.exclude_dir)):
            if key.endswith(".u8"):
                with open(os.path.join(args.exclude_dir, key), "rb") as fh:
                    masks[key[:-3]] = fh.read()
        print(f"exclude: {len(masks)} mask(s), "
              f"{sum(len(m) for m in masks.values()):,} docs covered, "
              f"{sum(int(v.get('dropped_docs', 0)) for v in dropped.values()):,} flagged", flush=True)
    streams = []
    for src in SOURCES:
        key = src_key(src)
        rec = st.d["staged"].get(key)
        if not rec or not rec.get("shards"):
            continue
        staged_docs = sum(int(s.get("docs", 0)) for s in rec["shards"])
        m = masks.get(key)
        if m is not None and len(m) != staged_docs:
            # A short mask would silently keep exactly the documents the audit did not reach.
            print(f"FATAL: mask for {key} covers {len(m)} docs but {staged_docs} are staged", flush=True)
            sys.exit(5)
        weight = rec.get("tokens") or sum(s["tokens"] for s in rec["shards"])
        weight = weight - int((dropped.get(key) or {}).get("dropped_tokens", 0))
        if weight <= 0:
            print(f"  {key}: fully excluded, dropped from the merge", flush=True)
            continue
        iters = [iter_shard_docs(os.path.join(st.root, s["file"])) for s in rec["shards"]]

        def gen(its=iters, m=m):
            if m is None:
                for it in its:
                    for d in it:
                        yield d
                return
            i = 0
            for it in its:
                for d in it:
                    if not m[i]:
                        yield d
                    i += 1

        streams.append({"key": key, "weight": float(weight), "credit": 0.0, "it": gen(iters),
                        "pending": None, "emitted_tokens": 0, "exhausted": False})
    if not streams:
        # The cold-resume rehearsal exited 0 here while producing an empty mix, and verify_mix then
        # reported recount_matches_manifest: true because 0 == 0. An empty mix is the worst possible
        # silent success in this pipeline, so it is a failure and not a return.
        print("FATAL merge: nothing staged. Refusing to emit an empty mix.", flush=True)
        sys.exit(4)
    if args.exclude_dir:
        # A missing mask must never mean "this source is clean" (review finding): it means the audit
        # never scanned it, and the default behaviour would have been to merge every contaminated doc.
        unscanned = [s["key"] for s in streams if s["key"] not in masks]
        if unscanned:
            print(f"FATAL merge: --exclude-dir has no mask for {unscanned}. Refusing to merge sources "
                  "the contamination audit did not read.", flush=True)
            sys.exit(5)

    # A re-merge is only correct if it starts from nothing. The loop below stops when merged_tokens()
    # reaches the input total, so on a mixroot that already holds 141 shards it would emit ZERO new
    # shards, exit 0, and let the manifest claim a filter that never happened (found in review while
    # Gate 2c was running against exactly this path). Resume-after-a-kill stays legal -- it is the same
    # inputs -- so the test is on a digest of the merge's inputs, not on whether shards exist.
    digest = hashlib.sha256(json.dumps({
        "streams": [[s["key"], round(s["weight"]) ] for s in streams], "seed": SHUFFLE_SEED,
        "tokens_per_shard": TOKENS_PER_SHARD, "val_stride": VAL_STRIDE,
        "filter": filter_digest}, sort_keys=True).encode()).hexdigest()
    prior = st.d.get("merge_inputs")
    if prior is not None and prior != digest and st.d["merged_shards"]:
        if not (args.rebuild_merged or args.exclude_dir):
            print("FATAL merge: the mix root already holds shards merged from DIFFERENT inputs "
                  f"(recorded {prior[:12]}, now {digest[:12]}). Continuing would be a silent no-op. "
                  "Pass --rebuild-merged to discard them deliberately.", flush=True)
            sys.exit(6)
        n_old = len(st.d["merged_shards"]) + len(st.d.get("val_shards", []))
        for grp in ("merged_shards", "val_shards"):
            for rec in st.d.get(grp, []):
                p = os.path.join(st.root, rec["file"])
                if os.path.exists(p):
                    os.remove(p)
        st.d["merged_shards"] = []
        st.d["val_shards"] = []
        print(f"merge: inputs changed ({prior[:12]} -> {digest[:12]}); discarded {n_old} old shard(s) "
              "and rebuilding from the staged sources", flush=True)
    st.d["merge_inputs"] = digest
    total_w = sum(s["weight"] for s in streams)
    print(f"merge: {len(streams)} sources, "
          f"{sum(s['weight'] for s in streams) / 1e6:,.0f}M staged tokens", flush=True)

    def next_doc(s):
        if s["pending"] is None and not s["exhausted"]:
            try:
                s["pending"] = next(s["it"])
            except StopIteration:
                s["exhausted"] = True
        return s["pending"]

    cap = args.merge_budget_tokens or int(total_w)
    buf, buf_src, buf_tokens = [], [], 0
    val_buf, val_tokens, docs_seq = [], 0, 0
    out_dir = os.path.join(st.root, "mix")
    os.makedirs(out_dir, exist_ok=True)
    t0 = time.monotonic()

    def flush_val():
        nonlocal val_buf, val_tokens
        if not val_buf:
            return
        idx = len(st.d.get("val_shards", []))
        p = os.path.join(st.root, "val", f"shard-{idx:04d}.bin")
        os.makedirs(os.path.dirname(p), exist_ok=True)
        tmp = p + ".tmp"
        offs = array.array("I", [0])
        ids = array.array("H")
        for d in val_buf:
            ids.extend(d)
            offs.append(len(ids))
        with open(tmp, "wb") as f:
            f.write(struct.pack("<IIII", vocab, len(val_buf), len(ids), MAGIC_FLAGS))
            offs.tofile(f)
            ids.tofile(f)
        os.replace(tmp, p)
        st.d.setdefault("val_shards", []).append({
            "file": os.path.relpath(p, st.root), "docs": len(val_buf), "tokens": len(ids),
            "sha256": hashlib.sha256(open(p, "rb").read()).hexdigest(), "bytes": os.path.getsize(p)})
        st.commit([], [])
        print(f"  val shard {idx:04d}: {len(ids):,} tok / {len(val_buf):,} docs", flush=True)
        val_buf, val_tokens = [], 0

    def flush():
        nonlocal buf, buf_src, buf_tokens
        if not buf:
            return
        random.Random(SHUFFLE_SEED + len(st.d["merged_shards"])).shuffle(buf)
        docs = [d for d, _ in buf]
        p = os.path.join(out_dir, f"shard-{len(st.d['merged_shards']):04d}.bin")
        nd, nt, dg, bs = write_shard(p, vocab, docs)
        st.d["merged_shards"].append({"file": os.path.relpath(p, st.root), "docs": nd, "tokens": nt,
                                      "sha256": dg, "bytes": bs,
                                      "sources": sorted(set(s for _, s in buf))})
        st.commit([], [])
        print(f"  mix shard {len(st.d['merged_shards']) - 1:04d}: {nt:,} tok / {nd:,} docs "
              f"from {len(st.d['merged_shards'][-1]['sources'])} sources "
              f"[{time.monotonic() - t0:.0f}s]", flush=True)
        buf, buf_src, buf_tokens = [], [], 0

    while st.merged_tokens() + buf_tokens < cap:
        live = [s for s in streams if next_doc(s) is not None]
        if not live:
            break
        # Largest-deficit scheduling: whoever is furthest below its proportional share of what has been
        # emitted goes next. O(n_sources) per batch of MERGE_BATCH_DOCS, not per document.
        emitted = sum(s["emitted_tokens"] for s in streams) or 1
        pick = max(live, key=lambda s: s["weight"] / total_w - s["emitted_tokens"] / emitted)
        took = 0
        while took < MERGE_BATCH_DOCS:
            d = next_doc(pick)
            if d is None:
                break
            docs_seq += 1
            if docs_seq % VAL_STRIDE == 0:
                # held out of training entirely: routed to the validation split instead of the buffer
                val_buf.append(d)
                val_tokens += len(d)
                if val_tokens >= VAL_TOKENS_PER_SHARD:
                    flush_val()
                pick["emitted_tokens"] += len(d)
                pick["pending"] = None
                took += 1
                continue
            buf.append((d, pick["key"]))
            buf_tokens += len(d)
            pick["emitted_tokens"] += len(d)
            pick["pending"] = None
            took += 1
            if buf_tokens >= TOKENS_PER_SHARD:
                flush()
                break
        if all(s["pending"] is None and s["exhausted"] for s in streams):
            break
    flush()
    flush_val()
    print(f"merge: {len(st.d['val_shards']) if 'val_shards' in st.d else 0} val shards, "
          f"{sum(s['tokens'] for s in st.d.get('val_shards', [])):,} held-out tokens", flush=True)
    print(f"merge: {len(st.d['merged_shards'])} shards, "
          f"{sum(s['tokens'] for s in st.d['merged_shards']):,} tokens", flush=True)


# --------------------------------------------------------------------- entry
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", default="/kaggle/working/mixroot")
    ap.add_argument("--budget-tokens", type=int, default=1_310_000_000, help="stage cap")
    ap.add_argument("--merge-budget-tokens", type=int, default=0, help="0 = take everything staged")
    ap.add_argument("--cap-tokens-per-source", type=int, default=0, help="smoke test knob")
    ap.add_argument("--min-chars", type=int, default=300)
    ap.add_argument("--max-chars", type=int, default=100_000)
    ap.add_argument("--sec-per-source", type=int, default=0)
    ap.add_argument("--only", default="", help="comma list of config names to stage")
    ap.add_argument("--hub-repo", default="Cion-lab/ounce100m-mix-stage",
                    help="staged sources are published here so a session kill cannot lose them; "
                         "empty string disables sync")
    ap.add_argument("--stage-only", action="store_true")
    ap.add_argument("--merge-only", action="store_true")
    ap.add_argument("--exclude-dir", default="",
                    help="directory of per-source <key>.u8 contamination masks + filter.json written by "
                         "audit_contamination --write-filter; documents flagged there are not merged (§3.3)")
    ap.add_argument("--rebuild-merged", action="store_true",
                    help="discard shards already recorded in build_state.json before merging. Implied by "
                         "--exclude-dir, because a filter changes every shard's contents")
    args = ap.parse_args()

    if args.only:
        keep = set(args.only.split(","))
        SOURCES[:] = [s for s in SOURCES if s["config"] in keep or src_key(s) in keep]
        print(f"--only: {len(SOURCES)} source(s)", flush=True)

    from tokenizers import Tokenizer
    from huggingface_hub import hf_hub_download
    tok_path = hf_hub_download("HuggingFaceTB/SmolLM2-135M", "tokenizer.json")
    tk = Tokenizer.from_file(tok_path)
    tk.no_truncation()
    assert tk.get_vocab_size() <= 65536, "uint16 shard layout breaks above 65,536"
    print(f"tokenizer vocab={tk.get_vocab_size()} "
          f"sha256={hashlib.sha256(open(tok_path, 'rb').read()).hexdigest()[:12]}", flush=True)

    global HS
    HS = None
    # Sync is best-effort to SET UP but not to USE: if the Hub cannot be reached or written, turning it
    # off silently would reintroduce the exact failure this exists to prevent (a killed session losing
    # hours of streaming). So a configured repo that cannot be initialised is fatal unless --no-hub.
    if args.hub_repo:
        try:
            import hubsync as HS_mod
            HS = HS_mod
            HS.ensure_repo(args.hub_repo)
        except Exception as e:
            print(f"FATAL hub sync could not initialise: {type(e).__name__}: {str(e)[:200]}",
                  flush=True)
            print("Refusing to run a multi-hour build that is not interruption-safe. Resolve the Hub "
                  "write path, or pass --hub-repo '' only if you accept losing all staged work when the "
                  "session ends.", flush=True)
            sys.exit(3)

    st = State(args.root)
    if args.hub_repo and not st.d["staged"]:
        # The state file lives on /kaggle/working too, so in a fresh session it is empty even when the
        # Hub holds a complete stage. Rebuild the stage bookkeeping from the Hub's record.json files --
        # without this, --merge-only on a new instance would find nothing and quietly emit an empty mix.
        try:
            remote = HS.remote_stage_files(args.hub_repo)
            keys = sorted({p.split("/")[1] for p in remote if p.startswith("stage/") and "/" in p})
            pulled = HS.download_missing(args.root, keys, args.hub_repo, remote=remote)
            for k in keys:
                rp = os.path.join(args.root, "stage", k, "record.json")
                if not os.path.exists(rp):
                    continue
                j = json.load(open(rp))
                st.d["staged"][k] = {"complete": True, "shards": j["shards"], "tokens": j["tokens"],
                                     "docs_seen": j["docs_seen"], "dropped": j["dropped"],
                                     "target": j["tokens"], "restored_from_hub": True}
            st.commit([], [])
            print(f"hub: restored {len(st.d['staged'])} staged source(s) "
                  f"({st.staged_tokens_total():,} tokens)", flush=True)
        except Exception as e:
            print(f"FATAL hub restore failed: {type(e).__name__}: {str(e)[:200]}", flush=True)
            sys.exit(3)

    t0 = time.monotonic()
    if not args.merge_only:
        stage(args, st, tk)
    if not args.stage_only:
        merge(args, st, tk)

    manifest = {
        "tokenizer": {"id": "HuggingFaceTB/SmolLM2-135M:tokenizer.json",
                      "vocab_size": tk.get_vocab_size(),
                      "sha256": hashlib.sha256(open(tok_path, "rb").read()).hexdigest()},
        "layout": {"dtype": "uint16", "header": "<IIII (vocab, n_docs, n_tokens, flags)",
                   "offsets": "uint32 × (n_docs+1) cumulative",
                   "doc_slice": "ids[offsets[i]:offsets[i+1]]",
                   "tokens_per_shard_target": TOKENS_PER_SHARD, "row_token_cap": ROW_TOKEN_CAP,
                   "merge_batch_docs": MERGE_BATCH_DOCS, "shuffle_seed_base": SHUFFLE_SEED,
                   "resume_cursor": "(shard_index, token_offset) into the concatenated id stream"},
        "stage": {k: {kk: vv for kk, vv in v.items() if kk != "shards"}
                  for k, v in st.d["staged"].items()},
        "staged_tokens": st.staged_tokens_total(),
        "shards": st.d["merged_shards"],
        "n_shards": len(st.d["merged_shards"]),
        "val_shards": st.d.get("val_shards", []),
        "total_tokens": sum(s["tokens"] for s in st.d["merged_shards"]),
        "val_tokens": sum(s["tokens"] for s in st.d.get("val_shards", [])),
        "val_holdout": {"stride": VAL_STRIDE, "note": "every 50th merged document, sampled across "
                         "sources rather than whole documents, excluded from training shards"},
        "distinct_sources_per_shard_min": min((len(s["sources"]) for s in st.d["merged_shards"]),
                                              default=0),
        "seconds": round(time.monotonic() - t0, 1),
    }
    if args.exclude_dir:
        # Provenance: a filtered mix is only reproducible if the exclusion set travels with it, so the
        # counts go in the manifest and the mask bytes go in mixroot/filter/ (publish_mix ships them).
        try:
            fj = json.load(open(os.path.join(args.exclude_dir, "filter.json")))
            manifest["contamination_filter"] = {
                "criterion": fj.get("criterion"), "k": fj.get("k"),
                "docs_dropped": sum(int(v.get("dropped_docs", 0)) for v in fj["sources"].values()),
                "tokens_dropped": sum(int(v.get("dropped_tokens", 0) or 0)
                                      for v in fj["sources"].values()),
                "per_source": {k2: {"docs": v2.get("dropped_docs"),
                                    "tokens": v2.get("dropped_tokens")}
                               for k2, v2 in fj["sources"].items()}}
        except Exception as e:
            print(f"FATAL: --exclude-dir set but filter.json unreadable: {e}", flush=True)
            sys.exit(5)
    with open(os.path.join(args.root, "manifest.json"), "w") as f:
        json.dump(manifest, f, indent=1)
    # Nothing may be printed between the JSON and its END marker: the caller parses this span with
    # json.loads, and a stray status line turned a passing rehearsal red (E-019).
    print("BUILD_JSON_BEGIN")
    print(json.dumps({k: manifest[k] for k in
                      ("staged_tokens", "total_tokens", "val_tokens", "n_shards",
                       "distinct_sources_per_shard_min", "seconds")}, indent=1))
    print("BUILD_JSON_END")


if __name__ == "__main__":
    main()