"""
clankerDiffusion — push artifacts to HuggingFace (via huggingface_hub).

  python push_hf.py --what all        # code + tokenizer + rag corpus + train.bin
  python push_hf.py --what code       # .py sources + tokenizer + rag corpus
  python push_hf.py --what data       # train.bin + tokenizer + rag corpus
  python push_hf.py --what ckpt       # latest checkpoints

Repos (under your HF user):
  coderofpears/clankerDiffusion-base        code + tokenizer + rag corpus
  coderofpears/clankerDiffusion-data        train.bin  (packed training corpus)
  coderofpears/clankerDiffusion-checkpoints model checkpoints (.pt)
"""
import os
import argparse
import tempfile
import shutil
import fnmatch

HERE = os.path.dirname(os.path.abspath(__file__))
DATADIR = os.path.join(HERE, "data")
CKPTDIR = os.path.join(HERE, "checkpoints")

CODE_REPO = "coderofpears/clankerDiffusion-base"
DATA_REPO = "coderofpears/clankerDiffusion-data"
CKPT_REPO = "coderofpears/clankerDiffusion-checkpoints"

PY_SOURCES = ["model.py", "tokenizer.py", "prep.py", "train.py", "infer.py",
              "tools.py", "agent.py", "rag.py", "build_rag.py", "rag_finetune.py",
              "colab_train.py", "modal_train.py", "push_hf.py", "upload_artifacts.py"]


def _token():
    # prefer explicit env, else read from .env
    tok = os.environ.get("HF_TOKEN")
    if tok:
        return tok
    for p in (os.path.join(HERE, ".env"), os.path.join(os.path.expanduser("~"), ".env")):
        if os.path.exists(p):
            for line in open(p, encoding="utf-8"):
                line = line.strip()
                if line.startswith("HF_TOKEN"):
                    return line.split("=", 1)[1].strip().strip('"').strip("'")
    raise SystemExit("HF_TOKEN not found in env or .env")


def _upload(repo, local, patterns):
    from huggingface_hub import HfApi
    api = HfApi(token=_token())
    try:
        api.create_repo(repo_id=repo, repo_type="model", exist_ok=True)
    except Exception as e:
        print(f"[hf] create_repo note: {e}")

    stage = tempfile.mkdtemp()
    try:
        for root, _, files in os.walk(local):
            for f in files:
                rel = os.path.relpath(os.path.join(root, f), local)
                if any(fnmatch.fnmatch(rel, p) or fnmatch.fnmatch(f, p) for p in patterns):
                    dst = os.path.join(stage, rel)
                    os.makedirs(os.path.dirname(dst), exist_ok=True)
                    shutil.copy(os.path.join(root, f), dst)
        if not os.listdir(stage):
            print(f"[hf] nothing matched for {repo}")
            return
        api.upload_folder(folder_path=stage, repo_id=repo, repo_type="model")
        print(f"[hf] -> {repo}")
    finally:
        shutil.rmtree(stage, ignore_errors=True)


def push_code():
    _upload(CODE_REPO, HERE, PY_SOURCES + ["data/tokenizer.json", "data/rag_corpus.txt", "README.md"])


def push_data():
    _upload(DATA_REPO, DATADIR, ["train.bin", "tokenizer.json", "rag_corpus.txt"])


def push_ckpt():
    _upload(CKPT_REPO, CKPTDIR, ["*.pt"])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--what", default="all", choices=["all", "code", "data", "ckpt"])
    a = ap.parse_args()
    if a.what in ("all", "code"):
        push_code()
    if a.what in ("all", "data"):
        push_data()
    if a.what in ("all", "ckpt"):
        push_ckpt()


if __name__ == "__main__":
    main()