"""
upload_artifacts.py  —  run THIS locally before launching the Colab notebook.

Publishes the model code + trained tokenizer + (optionally) your local
checkpoint to HuggingFace Hub so the Colab script can pull them.

  set HF_TOKEN=your_huggingface_token
  python upload_artifacts.py            # code + tokenizer
  python upload_artifacts.py --with-ckpt   # also upload latest local checkpoint
"""
import os, sys, argparse
from huggingface_hub import HfApi

ROOT = r"C:\Users\User\CalcGPU\clankerDiffusion"
CODE_REPO = "clankerDiffusion/base"
CKPT_REPO = "clankerDiffusion/checkpoints"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--with-ckpt", action="store_true")
    a = ap.parse_args()

    token = os.environ.get("HF_TOKEN")
    if not token:
        sys.exit("set HF_TOKEN first:  set HF_TOKEN=hf_xxx")
    api = HfApi(token=token)

    api.create_repo(CODE_REPO, repo_type="model", exist_ok=True)
    code_files = ["model.py", "tokenizer.py", "train.py", "prep.py", "infer.py",
                  os.path.join("data", "tokenizer.json"),
                  os.path.join("data", "tokenizer.json.meta.json")]
    for rel in code_files:
        p = os.path.join(ROOT, rel)
        if os.path.exists(p):
            api.upload_file(repo_id=CODE_REPO,
                           path_in_repo=os.path.basename(p),
                           path_or_fileobj=p)
            print("uploaded", rel)
        else:
            print("skip (missing)", rel)

    if a.with_ckpt:
        api.create_repo(CKPT_REPO, repo_type="model", exist_ok=True)
        ckpt_dir = os.path.join(ROOT, "checkpoints")
        if os.path.isdir(ckpt_dir):
            pts = sorted(f for f in os.listdir(ckpt_dir) if f.endswith(".pt"))
            if pts:
                p = os.path.join(ckpt_dir, pts[-1])
                api.upload_file(repo_id=CKPT_REPO,
                               path_in_repo=os.path.basename(p),
                               path_or_fileobj=p)
                print("uploaded checkpoint", pts[-1])
    print("DONE. Colab CODE_REPO =", CODE_REPO)


if __name__ == "__main__":
    main()