# Gate 2g: rebuild the filtered mix, re-audit it, verify, and PUBLISH mix-v1 with the batched uploader.
# Gate 2f rebuilt and audited the mix cleanly and then stopped at HfHubHTTPError after ~139 sequential
# commits: a Hub commit-RATE ceiling, not a network fault (E-033). publish_mix now sends the whole batch
# as ONE upload_folder commit, resumes by content sha256 (the listing reads lfs.oid, which is a dict --
# attribute access silently returned None for every shard and turned the "resume" into a 2.3 GB re-send),
# and hard-exits if any path is still absent or disagrees in sha after the commit.
import hashlib, json, os, shutil, subprocess, sys, threading, time

os.chdir("/kaggle/working")
sys.path.insert(0, "/kaggle/working")
REV = "292c9d28bb81251ce09c78a8e84bdbe69032c28e"
WANT = {
  "ounce100m_credentials.py": ("ounce100m_credentials.py",
                               "6525f62f03f2d73650a1eb4f70fcb52d1194caad4ca88b2d8bd8fd54f88339b6"),
  "hubsync.py": ("build/hubsync.py",
                 "e9696dd0c48bfcdb3bf201cc6b065d865dab988510646fbefefb49a4b35332c1"),
  "build_mix.py": ("build/build_mix.py",
                   "f0cf56256f7dffa2caac4105d59e261d775eb9c50437f142b49ce8ac18db86a1"),
  "verify_mix.py": ("build/verify_mix.py",
                    "38a6bb00a42f80cfdd85081847237edbdab22340a81d33f90e2956485ca55c01"),
  "audit_contamination.py": ("build/audit_contamination.py",
                             "58656d68977d858d763322b6d7966adc637457c4975e132df89586f3ac4d3088"),
  "publish_mix.py": ("build/publish_mix.py",
                     "2d609bb453ed73d9766d958268c858c6f5afed93d758ced8e041054a4aa6f697"),
}
BASE = "https://huggingface.co/Cion-lab/ounce100m-code/resolve/" + REV
for p, (rp, want) in sorted(WANT.items()):
    r = subprocess.run(["curl", "-sfL", BASE + "/" + rp, "-o", p], capture_output=True, text=True)
    assert r.returncode == 0, ("fetch failed", rp, (r.stderr or "")[:200])
    got = hashlib.sha256(open(p, "rb").read()).hexdigest()
    assert got == want, ("SHA MISMATCH", rp, got, want)
    print("OK", p, got[:12], flush=True)

import ounce100m_credentials as C
print("creds:", json.dumps(C.install(verify=True)), flush=True)
ROOT = "/kaggle/working/mixroot"
CACHE = "/kaggle/working/refcache"
FILT = os.path.join(ROOT, "filter")
os.makedirs(CACHE, exist_ok=True); os.makedirs(FILT, exist_ok=True)
STAGE_REPO = "Cion-lab/ounce100m-mix-stage"
MIX_REPO = "Cion-lab/ounce100m-mix-v1"
T0 = time.time()


def run(argv, label, timeout):
    """Stream a stage's output line by line: the Gate 2f version buffered everything, so a session that
    hit the notebook timeout reported nothing about the stage that was running."""
    print("=== " + label, flush=True)
    t0 = time.time()
    e = dict(os.environ); e["PYTHONPATH"] = "/kaggle/working"
    killed = []
    p = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,
                         env=e, bufsize=1)
    timer = threading.Timer(timeout, lambda: (killed.append(True), p.kill()))
    timer.daemon = True
    timer.start()
    tail = []
    try:
        for line in p.stdout:
            line = line.rstrip("\n")
            tail.append(line)
            del tail[:-80]
            print("   |", line[:280], flush=True)
    finally:
        timer.cancel()
    rc = p.wait()
    if killed:
        print("   TIMEOUT after %d s, killed" % timeout, flush=True)
        rc = -9
    print("%s_RC %s seconds %.1f elapsed %.0f" % (label, rc, time.time() - t0, time.time() - T0),
          flush=True)
    return rc, "\n".join(tail)


from huggingface_hub import HfApi, hf_hub_download, list_repo_files
have = {f for f in list_repo_files(STAGE_REPO, repo_type="dataset") if f.startswith("auditref/")}
for n in ("ref.npy", "refmask.npy", "refmeta.json", "refcounts.json", "errors.json", "srv.json"):
    if "auditref/" + n in have:
        shutil.copyfile(hf_hub_download(STAGE_REPO, "auditref/" + n, repo_type="dataset",
                                       cache_dir="/tmp/hf"), os.path.join(CACHE, n))
want_masks = sorted(f for f in have if f.startswith("auditref/filter/") and f.endswith(".u8"))
for f in want_masks + (["auditref/filter/filter.json"] if "auditref/filter/filter.json" in have else []):
    shutil.copyfile(hf_hub_download(STAGE_REPO, f, repo_type="dataset", cache_dir="/tmp/hf"),
                    os.path.join(FILT, f.split("/")[-1]))
nm = len([x for x in os.listdir(FILT) if x.endswith(".u8")])
print("masks restored:", nm, "refcache files:", len(os.listdir(CACHE)), flush=True)
if nm < 15 or not os.path.exists(os.path.join(FILT, "filter.json")):
    print("VERDICT GATE2G_STOP incomplete mask archive", flush=True)
    raise SystemExit(3)

rcm, _ = run([sys.executable, "build_mix.py", "--root", ROOT, "--merge-only", "--hub-repo",
              STAGE_REPO, "--exclude-dir", FILT, "--rebuild-merged"], "MERGE_FILTERED", 3600)
man = json.load(open(os.path.join(ROOT, "manifest.json")))
print("MANIFEST", json.dumps({k: man.get(k) for k in
      ("total_tokens", "val_tokens", "n_shards", "distinct_sources_per_shard_min",
       "contamination_filter")}, default=str)[:900], flush=True)
# Gate 2e measured this exact figure after dropping 5,878 contaminated documents. The merge is a
# deterministic round-robin over the same staged shards with the same masks, so equal here means the
# rebuild is byte-for-byte the mix that was audited; different means something in the inputs moved.
print("REMERGE_MATCHES_2E", man.get("total_tokens") == 1109714831, man.get("total_tokens"), flush=True)
if (man.get("n_shards") or 0) < 100 or (man.get("total_tokens") or 0) < 1_000_000_000:
    print("VERDICT GATE2G_STOP rebuild did not produce a real mix", flush=True)
    raise SystemExit(4)

rcb, _ = run([sys.executable, "audit_contamination.py", "--root", ROOT, "--ref-cache", CACHE,
              "--out", os.path.join(ROOT, "audit.json")], "AUDIT", 5400)
aud = json.load(open(os.path.join(ROOT, "audit.json")))
print("AUDIT_SUMMARY", json.dumps({k: aud.get(k) for k in
      ("overlap_total", "mix_documents_audited", "mix_documents_with_any_hit",
       "mix_documents_majority_hit", "val_documents_audited", "val_hits", "tasks_covered",
       "reference_loader_problems", "mix_bytes_sha256", "MEASURED", "AUDIT_PASSED")},
      default=str)[:900], flush=True)
rcv, _ = run([sys.executable, "verify_mix.py", "--root", ROOT], "VERIFY", 1800)
if aud.get("AUDIT_PASSED") is not True or rcv != 0:
    print("VERDICT GATE2G_STOP audit or verify failed -- not publishing", flush=True)
    raise SystemExit(5)

rcp, _ = run([sys.executable, "publish_mix.py", "--root", ROOT, "--repo", MIX_REPO], "PUBLISH", 5400)
print("PUBLISH_RC", rcp, flush=True)
# Read the published bytes back over the public resolve endpoint with no token in the process
# environment at all: that is what "the Hub serves this mix to anyone" means, and it is the only check
# that catches a commit that reported success but left a path absent.
ANON = """
import hashlib, json, os, urllib.request
man = json.load(open("@ROOT@/manifest.json"))
def read(path, want_sha):
    d = urllib.request.urlopen("@MIX@/resolve/main/" + path, timeout=600).read()
    ok = hashlib.sha256(d).hexdigest() == want_sha
    print("ANON_READBACK", path, len(d), "sha_match", ok, flush=True)
    assert ok, path
idxs = [0, 7, len(man["shards"]) // 2, len(man["shards"]) - 1]
for i in idxs:
    read(man["shards"][i]["file"], man["shards"][i]["sha256"])
for v in (man.get("val_shards") or [])[:2]:
    read(v["file"], v["sha256"])
import huggingface_hub as hh
print("REMOTE_FILE_COUNT", len(hh.list_repo_files("@REPO@", repo_type="dataset")), flush=True)
""".replace("@ROOT@", ROOT).replace("@REPO@", MIX_REPO) \
   .replace("@MIX@", "https://huggingface.co/datasets/" + MIX_REPO)
open("anonread.py", "w").write(ANON)
rcz, _ = run([sys.executable, "anonread.py"], "ANON_VERIFY", 1800)
print("VERDICT GATE2G_DONE publish_rc", rcp, "anon_rc", rcz, "tokens", man.get("total_tokens"),
      "overlap", aud.get("overlap_total"), "seconds", round(time.time() - T0, 1), flush=True)