#!/usr/bin/env python3
"""Push Ingot (SN97) build-pod artifacts to HF repo eyad-silx/progress.

Modes: --small (docs/logs/scripts), --big (GGUFs + QAT states), default = both.
Only uploads new/changed files (size+mtime_ns fingerprint). Files rejected by the
free-plan private-storage cap are retried on the overflow repo (silx-ai/Progress,
org quota) and, if that also fails, recorded as pending with sha256 in the
MANIFEST (bytes stay on the pod + local rescue cache).
State: /root/.hf_push_state.json. Log: /root/hf_push.log.
"""
import argparse, fcntl, glob, hashlib, json, os, sys
from datetime import datetime, timezone
from pathlib import Path
from huggingface_hub import HfApi

REPO = "eyad-silx/progress"
# Free-plan private-storage caps (~20 GiB user eyad-silx, org silx-ai also full).
# Files that hit the cap overflow to this private org repo (token has repo.write):
OVERFLOW_REPO = "silx-ai/Progress"
OVERFLOW_PREFIX = "ingot-backup/"

TOKEN = Path("/root/.hf_token").read_text().strip()
STATE_P = Path("/root/.hf_push_state.json")
LOG_P = Path("/root/hf_push.log")

README = """# Ingot (SN97) challenger build — progress proof

Automated backup of build/train/eval artifacts from the H100 pod (RunPod, 1x H100 NVL).
Pushed by `hf_push.py` — checks hourly and uploads any new/changed artifact.

- `king-ix/` — incumbent king GGUF (Ingot-IX-27B-Q2) + provenance files (baseline)
- `chal-identity/`, `chal-v1/`, `chal-n2/`, `chal-norms/` — challenger GGUF iterations
  (scale-refit / QAT variants; same 27B Q2 architecture as the king)
- `states/` — QAT state checkpoints (`t0..tN.safetensors`); the quota-blocked ones
  (t2–t9, t15) are stored as `states/missing_t2-t9_t15.tar.zst` (restore: `tar --zstd -xf`)
- `chal-n2/chal-norms.gguf.zst` — zstd-19 archive of the chal-n2 GGUF, uploaded compressed
  to fit the free-plan private-storage cap (raw sha256 in MANIFEST; restore: `zstd -d`)
- `artifacts/` — `HANDOFF.md`, build/train/eval scripts, run logs, sample draws,
  measurement results (`measure*.jsonl`, `pick*.jsonl`, `drawn.jsonl`)
- `MANIFEST.md` — size + sha256 + upload time of every artifact (regenerated each push),
  including a "Pending" section for files blocked by the free-plan private-storage cap
  (their bytes are preserved on the build pod and in a local rescue cache)
- `PROGRESS.md` — push history

The GGUFs are ~6.8 GB each (PQ2_0 quantized 27B linear-attention models).
Note: chal-identity's GGUF is byte-identical to king-ix's (sha256 in MANIFEST) — the
identity refit test round-trips the king exactly.
"""

def now():
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")

def log(evt, **kw):
    line = json.dumps({"t": now(), "event": evt, **kw})
    with open(LOG_P, "a") as f:
        f.write(line + "\n")
    print(line, flush=True)

def sha256(p):
    h = hashlib.sha256()
    with open(p, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 24), b""):
            h.update(chunk)
    return h.hexdigest()

def load_state():
    return json.loads(STATE_P.read_text()) if STATE_P.exists() else {}

def save_state(s):
    STATE_P.write_text(json.dumps(s, indent=1, sort_keys=True))

def small_tracks():
    t = [("/root/HANDOFF.md", "artifacts/HANDOFF.md")]
    for pat, dest in [("/root/*.py", "artifacts/scripts/"),
                      ("/root/*.sh", "artifacts/scripts/"),
                      ("/root/*.jsonl", "artifacts/"),
                      ("/root/*.json", "artifacts/"),
                      ("/root/*.log", "artifacts/logs/")]:
        for p in sorted(glob.glob(pat)):
            name = os.path.basename(p)
            if name in ("hf_push.log", "hf_push_nohup.log"):
                continue
            t.append((p, dest + name))
    for extra in ("README.md", "LICENSE", "NOTICE.txt"):
        p = f"/root/models/king-ix/{extra}"
        if os.path.exists(p):
            t.append((p, f"king-ix/{extra}"))
    return t

def big_tracks():
    # priority order: newest challenger first (pod lifetime is uncertain)
    t = [("/root/models/chal-identity/Ingot-IX-27B-Q2.gguf", "chal-identity/Ingot-IX-27B-Q2.gguf"),
         ("/root/models/chal-v1/chal-v1.gguf", "chal-v1/chal-v1.gguf"),
         ("/root/models/chal-n2/chal-norms.gguf", "chal-n2/chal-norms.gguf"),
         ("/root/models/chal-norms/chal-norms.gguf", "chal-norms/chal-norms.gguf"),
         ("/root/models/king-ix/Ingot-IX-27B-Q2.gguf", "king-ix/Ingot-IX-27B-Q2.gguf")]
    # NOTE: zstd-19 archives (chal-n2/chal-norms.gguf.zst, states_missing_t2-t9_t15.tar.zst)
    # exist on the pod but are NOT tracked: the quota check rejects any new LFS blob right now,
    # even 348 MB (tested 04:37 UTC Oct 26). Track them again once a cap is lifted — or just
    # let the raw files above push, which the cap will then allow too.
    for p in sorted(glob.glob("/root/states/*.safetensors")):
        t.append((p, "states/" + os.path.basename(p)))
    return t

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--small", action="store_true")
    ap.add_argument("--big", action="store_true")
    args = ap.parse_args()
    if args.small or args.big:
        tracks = (small_tracks() if args.small else []) + (big_tracks() if args.big else [])
    else:
        tracks = small_tracks() + big_tracks()

    lockf = open("/root/.hf_push.lock", "w")
    try:
        fcntl.flock(lockf, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except BlockingIOError:
        log("skipped", reason="another push run in progress")
        sys.exit(0)

    state = load_state()
    pending = state.get("_pending", {})
    api = HfApi(token=TOKEN)
    results = []
    for local, repo in tracks:
        if not os.path.exists(local):
            continue
        st = os.stat(local)
        fp = f"{st.st_size}:{st.st_mtime_ns}"
        prev = state.get(repo)
        if prev and prev.get("fp") == fp:
            results.append((repo, "unchanged"))
            continue
        d = None
        try:
            d = sha256(local)
            dest = REPO
            try:
                api.upload_file(path_or_fileobj=local, path_in_repo=repo, repo_id=REPO,
                                repo_type="model", commit_message=f"push {repo}")
            except Exception as e1:
                if "storage limit" not in str(e1).lower():
                    raise
                api.upload_file(path_or_fileobj=local, path_in_repo=OVERFLOW_PREFIX + repo,
                                repo_id=OVERFLOW_REPO, repo_type="model",
                                commit_message=f"overflow push {repo} (main repo private storage full)")
                dest = OVERFLOW_REPO
                log("overflowed", local=local, repo=repo, to=OVERFLOW_REPO, sha256=d, size=st.st_size)
            state[repo] = {"fp": fp, "sha256": d, "size": st.st_size, "uploaded_at": now(), "dest": dest}
            pending.pop(repo, None)
            state["_pending"] = pending
            save_state(state)
            if dest == REPO:
                log("uploaded", local=local, repo=repo, sha256=d, size=st.st_size)
            results.append((repo, "uploaded"))
        except Exception as e:
            entry = {"fp": fp, "size": st.st_size, "blocked_at": now()}
            if d:
                entry["sha256"] = d
            pending[repo] = entry
            state["_pending"] = pending
            save_state(state)
            log("error", local=local, repo=repo, err=repr(e)[:400])
            results.append((repo, "FAILED"))

    # regenerate docs (manifest / progress / readme) from state
    try:
        manifest = ["# Manifest — eyad-silx/progress (+ overflow repo silx-ai/Progress)", "",
                    "Files with `stored in` = silx-ai/Progress live at `ingot-backup/<repo path>` in that repo.",
                    "Free-plan private-storage caps (~20 GiB user eyad-silx; org silx-ai also at cap) block",
                    "the rest — see the Pending section; those bytes live on the build pod and in a local",
                    "rescue cache, and push automatically once a cap is lifted (plan upgrade / space freed).", "",
                    "| repo path | size (bytes) | sha256 | uploaded (UTC) | stored in |",
                    "|---|---|---|---|---|"]
        for repo in sorted(k for k in state if not k.startswith("_")):
            v = state[repo]
            manifest.append(f"| `{repo}` | {v['size']:,} | `{v['sha256']}` | {v['uploaded_at']} | {v.get('dest', REPO)} |")
        if pending:
            manifest += ["", "## Pending — blocked by private-storage quota", "",
                         "| repo path | size (bytes) | sha256 | blocked (UTC) |",
                         "|---|---|---|---|"]
            for repo in sorted(pending):
                v = pending[repo]
                manifest.append(f"| `{repo}` | {v['size']:,} | `{v.get('sha256', '(hash pending)')}` | {v['blocked_at']} |")
        Path("/tmp/MANIFEST.md").write_text("\n".join(manifest) + "\n")
        api.upload_file(path_or_fileobj="/tmp/MANIFEST.md", path_in_repo="MANIFEST.md",
                        repo_id=REPO, repo_type="model", commit_message="update MANIFEST.md")

        hist = [json.loads(l) for l in LOG_P.read_text().splitlines()[-80:] if l.strip()]
        prog = ["# Progress — push history", "", f"Last run: {now()}", ""]
        prog += [f"- `{h['t']}` **{h['event']}** {h.get('repo', h.get('local', ''))}" for h in hist]
        Path("/tmp/PROGRESS.md").write_text("\n".join(prog) + "\n")
        api.upload_file(path_or_fileobj="/tmp/PROGRESS.md", path_in_repo="PROGRESS.md",
                        repo_id=REPO, repo_type="model", commit_message="update PROGRESS.md")

        Path("/tmp/README_REPO.md").write_text(README)
        api.upload_file(path_or_fileobj="/tmp/README_REPO.md", path_in_repo="README.md",
                        repo_id=REPO, repo_type="model", commit_message="README")
    except Exception as e:
        log("error", step="docs", err=repr(e)[:500])

    n_up = sum(1 for r in results if r[1] == "uploaded")
    n_un = sum(1 for r in results if r[1] == "unchanged")
    n_fail = sum(1 for r in results if r[1] == "FAILED")
    log("run_done", uploaded=n_up, unchanged=n_un, failed=n_fail,
        pending=len(pending))
    sys.exit(1 if n_fail else 0)

if __name__ == "__main__":
    main()