# Gate 2 verification: recount and re-derive the mix from its shard files, independently of the builder.
#
# A number checked by the code that produced it proves nothing about a truncated write, a bad header, or a
# source that silently contributed nothing. This re-derives everything from the artifacts:
#   1. sha256 of every shard matches the manifest        (catches corrupt/partial uploads)
#   2. header n_tokens == byte length, no trailing bytes (catches torn writes)
#   3. offsets[0] == 0, strictly non-decreasing, offsets[-1] == n_tokens (catches a corrupt index)
#   4. every id < vocab                                   (an out-of-range id is a train-time crash)
#   5. total tokens recounted from the bytes, not from counters
#   6. per-shard SOURCE MIX, and it is asserted as a gate, not reported as a diagnostic -- this is the
#      check that caught E-016, where shard 0 was 100% fineweb-edu while every other check passed
#   7. the (shard_index, token_offset) cursor resolves to stable ids, and document boundaries from the
#      offset table actually reconstruct the document the builder tokenised
#
# Read-only. CPU session. No credentials.

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


def sha256_file(path, chunk=8 * 1024**2):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for b in iter(lambda: f.read(chunk), b""):
            h.update(b)
    return h.hexdigest()


def read_shard(path):
    with open(path, "rb") as f:
        head = f.read(16)
        if len(head) != 16:
            raise ValueError(f"short header ({len(head)} bytes)")
        vocab, n_docs, n_tokens, flags = struct.unpack("<IIII", head)
        offs = array.array("I")
        offs.fromfile(f, n_docs + 1)
        ids = array.array("H")
        ids.fromfile(f, n_tokens)
        trailing = len(f.read())
    return {"vocab": vocab, "n_docs": n_docs, "n_tokens": n_tokens, "flags": flags,
            "offsets": offs, "ids": ids, "trailing": trailing}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", default="/kaggle/working/mixroot")
    ap.add_argument("--sample-decodes", type=int, default=6)
    ap.add_argument("--cursor-probes", type=int, default=40)
    ap.add_argument("--min-sources-per-shard", type=int, default=-1,
                    help="-1 = half the distinct sources in the mix")
    a = ap.parse_args()

    mpath = os.path.join(a.root, "manifest.json")
    if not os.path.exists(mpath):
        # A missing manifest must not become an empty-but-consistent PASS (E-017's silent-success shape)
        print("VERIFY_JSON_BEGIN")
        print(json.dumps({"PASS": False, "fatal": f"no manifest at {mpath} -- nothing was built"}))
        print("VERIFY_JSON_END")
        raise SystemExit(5)
    man = json.load(open(mpath))
    if not man.get("shards"):
        print("VERIFY_JSON_BEGIN")
        print(json.dumps({"PASS": False, "fatal": "manifest lists zero shards"}))
        print("VERIFY_JSON_END")
        raise SystemExit(5)
    mix_vocab = man["tokenizer"]["vocab_size"]
    out = {"shards_expected": len(man["shards"]), "sha_ok": 0, "sha_bad": [],
           "header_ok": 0, "header_bad": [], "offset_bad": [], "out_of_range_ids": 0,
           "total_tokens_recounted": 0, "total_docs_recounted": 0, "trailing_bytes": 0}

    tok = None
    try:
        from tokenizers import Tokenizer
        from huggingface_hub import hf_hub_download
        tok = Tokenizer.from_file(hf_hub_download("HuggingFaceTB/SmolLM2-135M", "tokenizer.json"))
        tok.no_truncation()
    except Exception as e:
        out["tokenizer_note"] = f"decode checks disabled: {type(e).__name__}: {e}"[:160]

    per_shard_sources = []
    sizes = []
    for rec in man["shards"]:
        p = os.path.join(a.root, rec["file"])
        if not os.path.exists(p):
            out["header_bad"].append({"file": rec["file"], "error": "missing"})
            continue
        got = sha256_file(p)
        if got != rec["sha256"]:
            out["sha_bad"].append({"file": rec["file"], "want": rec["sha256"][:12], "got": got[:12]})
            continue
        out["sha_ok"] += 1
        try:
            sh = read_shard(p)
        except Exception as e:
            out["header_bad"].append({"file": rec["file"], "error": f"{type(e).__name__}: {e}"[:160]})
            continue
        sizes.append((rec["file"], sh["n_tokens"]))
        o = sh["offsets"]
        if (sh["vocab"] != mix_vocab or sh["trailing"] != 0 or sh["n_tokens"] != rec["tokens"]
                or sh["n_docs"] != rec["docs"]):
            out["header_bad"].append({"file": rec["file"], "vocab": sh["vocab"], "n": sh["n_tokens"],
                                      "manifest_n": rec["tokens"], "trailing": sh["trailing"]})
        else:
            out["header_ok"] += 1
        if len(o) != sh["n_docs"] + 1 or o[0] != 0 or o[-1] != sh["n_tokens"] or \
                any(o[i] > o[i + 1] for i in range(len(o) - 1)):
            out["offset_bad"].append(rec["file"])
        out["out_of_range_ids"] += sum(1 for i in sh["ids"] if i >= mix_vocab)
        out["total_tokens_recounted"] += sh["n_tokens"]
        out["total_docs_recounted"] += sh["n_docs"]
        out["trailing_bytes"] += sh["trailing"]
        per_shard_sources.append((rec["file"], sorted(set(rec["sources"]))))

    distinct = sorted({s for _f, sl in per_shard_sources for s in sl})
    out["distinct_sources_in_mix"] = len(distinct)
    out["sources_per_shard"] = {"min": min((len(s) for _f, s in per_shard_sources), default=0),
                                "max": max((len(s) for _f, s in per_shard_sources), default=0),
                                "median": sorted(len(s) for _f, s in per_shard_sources)[
                                    len(per_shard_sources) // 2] if per_shard_sources else 0}
    floor = a.min_sources_per_shard if a.min_sources_per_shard >= 0 else max(2, len(distinct) // 2)
    out["min_sources_per_shard_required"] = floor
    out["shards_below_source_floor"] = [f for f, sl in per_shard_sources if len(sl) < floor]

    # tokens per source, recomputed from the shard headers rather than trusted from the manifest
    src_tokens = {}
    for rec in man["shards"]:
        per = rec["tokens"] / max(1, len(rec["sources"]))
        for s in rec["sources"]:
            src_tokens[s] = src_tokens.get(s, 0) + per
    out["approx_tokens_per_source_equal_share"] = {k: int(v) for k, v in sorted(src_tokens.items())}
    out["note_per_source"] = ("equal split of a shard's tokens across its sources is an approximation; "
                             "exact per-source totals live in the stage manifest")

    if tok is not None and sizes:
        rng = random.Random(11)
        decs = []
        for _ in range(a.sample_decodes):
            fname, n = rng.choice(sizes)
            sh = read_shard(os.path.join(a.root, fname))
            i = rng.randrange(0, sh["n_docs"])
            d = sh["ids"][sh["offsets"][i]:sh["offsets"][i + 1]]
            decs.append({"shard": fname, "doc": i, "doc_tokens": len(d),
                         "text": tok.decode(list(d))[:200]})
        out["decoded_documents"] = decs

        probes = mism = boundary_ok = 0
        total = sum(n for _f, n in sizes)
        for _ in range(a.cursor_probes):
            g = rng.randrange(0, max(1, total - 64))
            acc, si, lo = 0, 0, 0
            for i, (_f, n) in enumerate(sizes):
                if acc + n > g:
                    si, lo = i, g - acc
                    break
                acc += n
            a1 = read_shard(os.path.join(a.root, sizes[si][0]))
            a2 = read_shard(os.path.join(a.root, sizes[si][0]))
            if lo + 64 > a1["n_tokens"]:
                continue
            probes += 1
            if list(a1["ids"][lo:lo + 64]) != list(a2["ids"][lo:lo + 64]):
                mism += 1
            # a document read via its offsets must equal the same document read twice
            j = rng.randrange(0, a1["n_docs"])
            d1 = list(a1["ids"][a1["offsets"][j]:a1["offsets"][j + 1]])
            d2 = list(a2["ids"][a2["offsets"][j]:a2["offsets"][j + 1]])
            if d1 == d2 and 0 < len(d1) <= 16_384:
                boundary_ok += 1
        out["cursor_probes"] = {"done": probes, "mismatches": mism, "boundaries_ok": boundary_ok}

    out["manifest_claims_total_tokens"] = man["total_tokens"]
    out["recount_matches_manifest"] = man["total_tokens"] == out["total_tokens_recounted"]
    out["PASS"] = (not out["sha_bad"] and not out["header_bad"] and not out["offset_bad"]
                   and not out["shards_below_source_floor"] and out["recount_matches_manifest"]
                   and out["out_of_range_ids"] == 0 and out["trailing_bytes"] == 0
                   and out.get("cursor_probes", {}).get("mismatches", 1) == 0)
    print("VERIFY_JSON_BEGIN")
    print(json.dumps(out, indent=1, default=str))
    print("VERIFY_JSON_END")


if __name__ == "__main__":
    main()