#!/usr/bin/env python3
"""Stage (hardlinks), upload, and byte-verify the two Nex-N2.5-mini repos.
usage: nex_publish.py stage|upload|verify|readme <std|imat> [README path]
Upload through upload_watchdog.sh (40G cap, stall kill + retry) - see nex_cards.sh upload
Env: default HF_HOME (token); HF_XET_CACHE / TMPDIR pinned to /mnt/models by the caller."""
import hashlib, json, os, shutil, sys, urllib.request

W = "/mnt/models/nex-n2.5-mini"; N = "Nex-N2.5-mini"; U = "kingjones777"
REPOS = {"std": f"{U}/{N}-ROCmFP4-GGUF", "imat": f"{U}/{N}-ROCmFP4-imatrix-GGUF"}


def pubname(fn):
    """Published filename: the quant token must be hyphen-delimited or the Hub cannot parse the variant
    (`Q4_0_ROCMFP4_STRIX_LEAN` reads as the garbage label `Q4_0_ROCMFP`)."""
    return fn.replace("Q4_0_ROCMFP4_", "Q4_0-ROCmFP4-")


MMPROJ = {f"mmproj-{N}-BF16.gguf": (f"{W}/out", f"mmproj-{N}-BF16.gguf")}
# the stock chat template plus one line that maps enable_thinking onto reasoning_effort (served with
# --chat-template-file; see the card's Reasoning controls) - shipped in both repos
TPL = {"chat_template_enable_thinking.jinja": (f"{W}/tpl", "chat_template_enable_thinking.jinja")}
TIERS = ("Q4_0_ROCMFP4_STRIX_LEAN", "Q4_0_ROCMFP4_COHERENT", "Q4_0_ROCMFP4_FAST")
FILES = {
    "std": {**{pubname(f"{N}-{t}.gguf"): (f"{W}/out", f"{N}-{t}.gguf") for t in TIERS}, **MMPROJ, **TPL},
    "imat": {**{pubname(f"{N}-imatrix-{t}.gguf"): (f"{W}/out-imat", f"{N}-imatrix-{t}.gguf") for t in TIERS},
             f"{N}.imatrix": (f"{W}/imat", f"{N}.imatrix"), **MMPROJ, **TPL},
}
P = "recipe/pipeline/"
SCRIPTS = ("nex_download.sh", "verify_download.py", "nex_phase1.sh", "nex_phase2.sh", "diag_bf16.sh", "nex_phase2b.sh",
           "nex_harness.py", "nex_bench.py", "nex_sizing.sh",
           "nex_cachegate.py", "readback.py", "nex_aggregate.py", "nex_judge.py", "nex_render.py", "nex_publish.py",
           "nex_cards.sh", "upload_watchdog.sh", "nex_seats.sh",
           # tool-call / reasoning diagnosis (stock template 6/14) and the template fix
           "diag_tools_run.sh", "nex_tools_diag.py", "nex_reasoning_probe.py", "nex_tools_tpl.py",
           "nex_seat_default_probe.py", "run_tools_c1.sh", "run_tools_roff.sh", "nex_refresh_cards.sh", "nex_publish_finish.sh")
LOGS = (
    "D2_verify_download.log", "phase1.log", "phase2.log", "Q_sizes.log", "N6_bench.log", "N7_sizing.log",
    "b_n-vision-q106-faon.log", "b_n-vision-q106-faoff.log",
    "C1_convert.log", "C2_mmproj.log", "C_readback.log",
    "Q1_q106.log", "Q1_q102.log", "Q1_q103.log", "Q_readback.log",
    "N1c_ppl_bf16_cpu.log", "N2c_imatrix_cpu.log",
    # the stopped first attempt on ROCm0 and the backend diagnosis that followed (the Known issues evidence)
    "N1_ppl_bf16.log", "diag_ppl_q106_rocm_c4.log", "diag_bf16.log", "diag_bf16_rocm_faoff.log", "diag_bf16_vk_faon.log",
    "diag_q106_vk_faon.log", "diag_bf16_cpu.log", "diag_bf16_purecpu_c1.log",
    "N3_q106i.log", "N3_q102i.log", "N3_q103i.log", "N3_readback.log",
    *(f"{p}_kld_{t}.log" for p in ("N4", "N4v") for t in ("q106", "q102", "q103", "q106i", "q102i", "q103i")),
    "N5_kld_q106_repeat.log", "N5v_kld_q106_repeat.log", "b_n-c3-q106.log",
    "b_n-tools-q106.log", "diag_tools.log", "diag_tools_server.log",
    "probe_reasoning.log", *(f"probe_reasoning_{c}.log" for c in
                             ("default", "fmt-deepseek", "srv-kwargs-high", "reasoning-on", "tpl-enable-thinking")),
    "N6t_tools_tpl.log", "b_n-tools-q106-tpl.log", "b_n-tools-q106-tpl-probe.log",
    "N6t_tools_tpl_medium.log", "b_n-tools-q106-tpl-medium.log", "b_n-tools-q106-tpl-medium-probe.log",
    "probe_seat_default.log", "probe_seat_default_C1.log", "probe_seat_default_C2.log",
    "N6t_tools_c1.log", "b_n-tools-q106-c1.log", "b_n-tools-q106-c1-probe.log", "b_n-vision-q106-c1-faon.log",
    "N6t_tools_roff.log", "b_n-tools-q106-roff.log", "b_n-tools-q106-roff-r2.log", "b_n-tools-q106-roff-r3.log",
    "b_n-tools-q106-roff-probe.log", "b_n-vision-q106-roff-faon.log",
    "N8_unice.log", "N8a_seats.log", "N8b_seats.log", "N8c_seats.log", "N8d_seats.log",
)
RAW = ("nex_repeat.jsonl", "nex_reference.jsonl", "nex_bench.jsonl", "nex_sizing.jsonl", "hub_check.json",
       "nex_tools_diag.json", "nex_reasoning_probe.json", "nex_seat_default_probe.json", "nex_template_shim.json",
       "nex_template_shim_medium.json", "nex_template_shim_c1.json", "nex_template_shim_roff.json",
       "nex_seats.jsonl", "nex_seats_plan.json")
TEMPLATES = ("chat_template_enable_thinking_medium.jinja", "chat_template_enable_thinking_v2.jinja")  # tested, not used
RECIPE = {
    "recipe/results_summary.json": "results/summary.json",
    **{P + x: x for x in SCRIPTS},
    **{f"recipe/raw/{x}": f"results/{x}" for x in RAW},
    **{f"recipe/logs/{x}": f"logs/{x}" for x in LOGS},
    **{f"recipe/templates/{x}": f"tpl/{x}" for x in TEMPLATES},
    "recipe/templates/chat_template_stock.jinja": "hf/chat_template.jinja",
}


def separation_errors(kind, names):
    """The imatrix build is its OWN repo (King): no imatrix weights in the standard repo, no standard weights in the
    imatrix repo. The vision projector is the only model file both carry."""
    errs = []
    for fn in names:
        if not fn.endswith(".gguf") or fn.startswith("mmproj-"):
            continue
        is_imat = "-imatrix-" in fn
        if kind == "std" and is_imat:
            errs.append(f"imatrix model file in the standard repo: {fn}")
        if kind == "imat" and not is_imat:
            errs.append(f"standard model file in the imatrix repo: {fn}")
    if kind == "imat" and f"{N}.imatrix" not in names:
        errs.append("imatrix repo is missing the .imatrix file")
    return errs


assert REPOS["std"] != REPOS["imat"]
assert set(FILES["std"]) & set(FILES["imat"]) == set(MMPROJ) | set(TPL), "only the projector + template may be shared"
for _k in FILES:
    assert not separation_errors(_k, list(FILES[_k])), separation_errors(_k, list(FILES[_k]))


def git_blob_sha1(p):
    """What the Hub reports as `oid` for a file stored in plain git (no LFS/Xet pointer)."""
    h = hashlib.sha1(b"blob %d\0" % os.path.getsize(p))
    with open(p, "rb") as fh:
        for b in iter(lambda: fh.read(64 << 20), b""):
            h.update(b)
    return h.hexdigest()


def sha(p):
    h = hashlib.sha256()
    with open(p, "rb") as fh:
        for b in iter(lambda: fh.read(64 << 20), b""):
            h.update(b)
    return h.hexdigest()


kind = sys.argv[2]; stage = f"{W}/hf-upload/{kind}"; repo = REPOS[kind]
if sys.argv[1] == "stage":
    missing = [f"{d}/{src}" for fn, (d, src) in FILES[kind].items() if not os.path.exists(f"{d}/{src}")] + \
              [src for src in RECIPE.values() if not os.path.exists(f"{W}/{src}")]
    if missing:
        raise SystemExit(f"STAGE ABORT - missing: {missing}")
    shutil.rmtree(stage, ignore_errors=True)
    for fn, (d, src) in FILES[kind].items():
        os.makedirs(stage, exist_ok=True); os.link(f"{d}/{src}", f"{stage}/{fn}")
    for dst, src in RECIPE.items():
        os.makedirs(os.path.dirname(f"{stage}/{dst}"), exist_ok=True); shutil.copy2(f"{W}/{src}", f"{stage}/{dst}")
    sums = {fn: sha(f"{stage}/{fn}") for fn in sorted(FILES[kind])}
    with open(f"{stage}/SHA256SUMS", "w") as fh:
        for fn, h in sums.items():
            fh.write(f"{h}  {fn}\n")
    sizes = {x: os.path.getsize(f"{stage}/{x}") for x in sums}
    json.dump(sizes, open(f"{stage}/../staged_sizes_{kind}.json", "w"))
    print(json.dumps({"stage": stage, "files": len(sums), "bytes": sum(sizes.values()), "recipe_files": len(RECIPE)}))
elif sys.argv[1] == "upload":
    from huggingface_hub import HfApi
    api = HfApi()
    print("whoami:", api.whoami()["name"], flush=True)
    api.create_repo(repo, repo_type="model", private=False, exist_ok=True)
    api.upload_large_folder(repo_id=repo, repo_type="model", folder_path=stage,
                            ignore_patterns=["README.md", ".cache/**"], num_workers=2)
    print("UPLOAD_DONE", repo, flush=True)
elif sys.argv[1] == "verify":
    tree = json.load(urllib.request.urlopen(f"https://huggingface.co/api/models/{repo}/tree/main?recursive=true", timeout=60))
    remote = {t["path"]: t for t in tree if t.get("type") == "file"}
    local = dict(l.split()[::-1] for l in open(f"{stage}/SHA256SUMS").read().splitlines())
    bad = []
    for fn, h in local.items():
        r = remote.get(fn)
        if not r:
            bad.append((fn, "MISSING remote")); continue
        if r["size"] != os.path.getsize(f"{stage}/{fn}"):
            bad.append((fn, f"size {r['size']}")); continue
        lfs = r.get("lfs")
        if lfs:
            if lfs.get("oid") != h:
                bad.append((fn, f"sha {lfs.get('oid')} != {h[:12]}"))
        elif r.get("oid") != git_blob_sha1(f"{stage}/{fn}"):
            bad.append((fn, f"git oid {r.get('oid')} does not match the staged bytes"))
    for fn in list(RECIPE) + ["SHA256SUMS"]:
        if fn not in remote:
            bad.append((fn, "MISSING remote"))
    stale = sorted(x for x in remote if x not in local and x not in RECIPE and x not in ("SHA256SUMS", "README.md", ".gitattributes"))
    bad += [(x, "SEPARATION") for x in separation_errors(kind, list(remote))]
    print(json.dumps({"repo": repo, "checked": len(local), "bad": bad, "unexpected_remote_files": stale,
                      "result": "PASS" if not bad and not stale else "FAIL"}))
    sys.exit(0 if not bad and not stale else 1)
elif sys.argv[1] == "readme":
    from huggingface_hub import HfApi
    HfApi().upload_file(path_or_fileobj=sys.argv[3], path_in_repo="README.md", repo_id=repo, repo_type="model",
                        commit_message="Model card: measured results, reproduction, known issues")
    print("README_DONE", repo)
else:
    raise SystemExit(f"usage: {sys.argv[0]} stage|upload|verify|readme <std|imat> [README path]")