#!/usr/bin/env python3
"""
Assemble the Hugging Face bucket and publish it.

    python code/upload_bucket.py --repo adeerkhan/msq-plurality-gap [--dry-run]

Stages, in order:
  1. dataset/   the 42 items (jsonl + csv + card)
  2. code/      the full harness
  3. results/   every generation and every metric
  4. figures/   the figures used in the README
  5. model/     the evaluated model's weights (openbmb/MiniCPM5-2B, Apache-2.0)
  6. README.md / FINDINGS.md / LICENSE / requirements.txt
"""
from __future__ import annotations

import argparse
import os
import pathlib
import shutil
import sys

ROOT = pathlib.Path(__file__).resolve().parent.parent
STAGE = ROOT / ".bucket"

MODEL_REPO = "openbmb/MiniCPM5-2B"
MODEL_FILES = ["config.json", "generation_config.json", "model.safetensors",
                "model.safetensors.index.json", "tokenizer.json",
                "tokenizer_config.json", "special_tokens_map.json",
                "chat_template.jinja"]

README_TEXT = """# MiniCPM5-2B (evaluated checkpoint)

Unmodified weights of [`openbmb/MiniCPM5-2B`](https://huggingface.co/openbmb/MiniCPM5-2B),
mirrored here so the bucket is self-contained for re-running the evaluation.

* Architecture: `LlamaForCausalLM`, 42 layers, GQA 16 Q / 2 KV heads
* Parameters: 2,516,756,480 (1,981,982,720 non-embedding)
* Context length: 131,072
* Licence: Apache-2.0
* SHA of the upstream revision this was taken from: see `UPSTREAM_REVISION.txt`

The evaluation in this bucket was run against these weights in bfloat16 (and,
for the cross-model comparison, in NF4 4-bit for every model so that the
protocol is identical across models). The weights here are the unmodified
upstream bf16 checkpoint.
"""


def stage_model() -> None:
    from huggingface_hub import snapshot_download
    print("resolving model snapshot ...")
    snap = pathlib.Path(snapshot_download(MODEL_REPO,
                                          allow_patterns=MODEL_FILES + ["*.md"]))
    dst = STAGE / "model"
    dst.mkdir(parents=True, exist_ok=True)
    for f in sorted(snap.iterdir()):
        if f.name.endswith(".md"):
            continue
        target = dst / f.name
        if target.exists():
            continue
        shutil.copyfile(f.resolve(), target)
        print(f"  + model/{f.name} ({target.stat().st_size/1e6:.0f} MB)")
    # revision provenance
    from huggingface_hub import HfApi
    info = HfApi().model_info(MODEL_REPO)
    (dst / "UPSTREAM_REVISION.txt").write_text(
        f"upstream repo : https://huggingface.co/{MODEL_REPO}\n"
        f"revision sha  : {info.sha}\n"
        f"last modified : {info.last_modified}\n"
        f"licence       : {info.card_data.get('license') if info.card_data else 'apache-2.0'}\n")
    (dst / "README.md").write_text(README_TEXT)


EXCLUDE_CODE = {"stop_all.sh", "stop_fetches.sh", "stop_runs.sh", "fetch_one.sh",
                "__pycache__"}


def stage_tree() -> None:
    STAGE.mkdir(parents=True, exist_ok=True)
    for name in ("dataset", "code", "results", "figures"):
        src = ROOT / name
        if not src.exists():
            continue
        dst = STAGE / name
        if dst.exists():
            shutil.rmtree(dst)
        shutil.copytree(src, dst, ignore=shutil.ignore_patterns("__pycache__", "*.pyc"))
        if name == "code":
            for f in EXCLUDE_CODE:
                (dst / f).unlink(missing_ok=True)
        print(f"  + {name}/")
    for f in ("README.md", "FINDINGS.md", "LICENSE", "requirements.txt"):
        if (ROOT / f).exists():
            shutil.copyfile(ROOT / f, STAGE / f)
            print(f"  + {f}")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", default="adeerkhan/msq-plurality-gap")
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--skip-model", action="store_true")
    ap.add_argument("--no-stage", action="store_true",
                    help="reuse the already-staged tree (for retries)")
    args = ap.parse_args()

    if STAGE.exists() and args.no_stage:
        print(f"reusing staged tree at {STAGE}")
    else:
        if STAGE.exists():
            shutil.rmtree(STAGE)
        print("staging tree ...")
        stage_tree()
        if not args.skip_model:
            stage_model()

    total = sum(p.stat().st_size for p in STAGE.rglob("*") if p.is_file())
    print(f"staged {total/1e9:.2f} GB at {STAGE}")

    if args.dry_run:
        for p in sorted(STAGE.rglob("*")):
            if p.is_file():
                print(f"  {p.relative_to(STAGE)}  {p.stat().st_size/1e6:.1f} MB")
        return

    from huggingface_hub import HfApi, upload_folder
    api = HfApi()
    api.create_repo(repo_id=args.repo, repo_type="model", exist_ok=True,
                    private=False)
    print(f"uploading to https://huggingface.co/{args.repo}")
    upload_folder(repo_id=args.repo, repo_type="model",
                  folder_path=str(STAGE), commit_message="MSQ: the unstated premise")
    print("done")


if __name__ == "__main__":
    main()