"""Stage and upload the campaign, then verify it reads back before anything is destroyed.

/workspace is not a volume on this instance, so the container filesystem is the
only copy until this runs. This project has already lost two batches exactly that
way -- the olmix-policy/ft-as20k runs behind policy-mixture.md, and 256 raw
validation dumps from ranking-transfer.md -- so the verify step is not optional
and the instance is not destroyed on an unverified upload.

Uploads every measurement and provenance record plus the final checkpoint of each
run. Intermediate depth checkpoints (14 per run, ~100 GB) are deliberately left
behind: their probe scores are archived, and they are reproducible from the
recipe, the pinned mixture vector and the seed.
"""
import argparse
import hashlib
import json
import shutil
from pathlib import Path

from huggingface_hub import HfApi

STAGE = Path("/workspace/upload")
RUNS = Path("/workspace/runs")
KEEP = ("run.json", "config.yaml", "status.json", "events.jsonl")


def sha(p: Path, n=1 << 20) -> str:
    h = hashlib.sha256()
    with open(p, "rb") as f:
        while chunk := f.read(n):
            h.update(chunk)
    return h.hexdigest()


def stage() -> dict:
    if STAGE.exists():
        shutil.rmtree(STAGE)
    manifest = {}
    for campaign in ("mae-64", "mae-seed", "mae-lr", "mae-validate"):
        root = RUNS / campaign
        if not root.exists():
            continue
        for run in sorted(root.iterdir()):
            if not run.is_dir():
                continue
            dest = STAGE / campaign / run.name
            dest.mkdir(parents=True, exist_ok=True)
            for name in KEEP:
                if (run / name).exists():
                    shutil.copy2(run / name, dest / name)
            exports = sorted((run / "exports").glob("step_*"))
            for exp in exports:
                # every readout at every depth, but weights only at the last
                for j in list(exp.glob("probe.json")) + list(exp.glob("lora*.json")):
                    d = dest / "exports" / exp.name
                    d.mkdir(parents=True, exist_ok=True)
                    shutil.copy2(j, d / j.name)
            if exports:
                d = dest / "exports" / exports[-1].name
                d.mkdir(parents=True, exist_ok=True)
                for name in ("model.safetensors", "config.yaml"):
                    if (exports[-1] / name).exists():
                        shutil.copy2(exports[-1] / name, d / name)
    # analysis outputs, campaign specs, and the code that produced them
    for src, sub in ((Path("/workspace/analysis"), "analysis"),
                     (Path("/workspace/scripts"), "scripts"),
                     (Path("/workspace/configs"), "configs")):
        for f in sorted(src.rglob("*")):
            if f.is_file() and f.suffix in (".json", ".py", ".yaml", ".csv"):
                d = STAGE / sub / f.relative_to(src)
                d.parent.mkdir(parents=True, exist_ok=True)
                shutil.copy2(f, d)
    shutil.copy2("/workspace/code/eat-map-regmix/eatmap/mae.py", STAGE / "scripts" / "mae.py")
    for f in sorted(STAGE.rglob("*")):
        if f.is_file():
            manifest[str(f.relative_to(STAGE))] = {"size": f.stat().st_size, "sha256": sha(f)}
    (STAGE / "MANIFEST.json").write_text(json.dumps(manifest, indent=1, sort_keys=True))
    return manifest


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", default="quinnlue/mae-cross-objective")
    ap.add_argument("--private", action="store_true")
    ap.add_argument("--verify-only", action="store_true")
    args = ap.parse_args()

    api = HfApi()
    if not args.verify_only:
        m = stage()
        total = sum(v["size"] for v in m.values())
        print(f"staged {len(m)} files, {total/2**30:.2f} GiB -> {args.repo}")
        api.create_repo(args.repo, repo_type="model", private=args.private, exist_ok=True)
        api.upload_folder(folder_path=str(STAGE), repo_id=args.repo, repo_type="model",
                          commit_message="MAE cross-objective mixture-ranking campaign")
        print("upload complete")

    # ---- verify: every staged file present on the Hub at the right size, and a
    # random sample byte-identical after a round trip.
    manifest = json.loads((STAGE / "MANIFEST.json").read_text())
    remote = {f.rfilename: f for f in api.repo_info(
        args.repo, repo_type="model", files_metadata=True).siblings}
    missing = [k for k in manifest if k not in remote]
    wrong = [k for k, v in manifest.items()
             if k in remote and remote[k].size not in (None, v["size"])]
    print(f"\nverify: {len(manifest)} staged | {len(remote)} on hub | "
          f"missing {len(missing)} | size-mismatch {len(wrong)}")
    if missing:
        print("  MISSING:", missing[:10])
    if wrong:
        print("  MISMATCH:", wrong[:10])

    import random
    from huggingface_hub import hf_hub_download
    random.seed(0)
    sample = random.sample([k for k in manifest if k in remote],
                           min(8, len(manifest)))
    bad = []
    for k in sample:
        p = hf_hub_download(args.repo, k, repo_type="model",
                            cache_dir="/workspace/.verify")
        if sha(Path(p)) != manifest[k]["sha256"]:
            bad.append(k)
    print(f"round-trip sha256 on {len(sample)} sampled files: "
          f"{'ALL MATCH' if not bad else f'MISMATCH {bad}'}")
    ok = not missing and not wrong and not bad
    print(f"\nARCHIVE {'VERIFIED' if ok else 'NOT VERIFIED -- DO NOT DESTROY'}")
    return 0 if ok else 1


if __name__ == "__main__":
    raise SystemExit(main())