#!/usr/bin/env python3
"""One-shot push of the code and the newest checkpoint per run to the Hub.

hf_sync.py (the sidecar on the compute node) cannot upload: huggingface_hub is
not importable there, so it only scp's tarballs to the login node's staging
dirs. This script runs on the LOGIN node, where huggingface_hub does import, and
does the actual upload:

    code/...            every .py/.sh of the run directory, flat mirror
    runs/<tag>/steps/   the newest N checkpoint tarballs found in
                        /tmp/sudoku_hf_<tag>/

Tarball names (step_*/stage*) are unique per checkpoint, so anything already in
the repo is skipped; the rolling latest.tar is always re-sent.
"""
import argparse
import os
import sys
import time


def code_files(root):
    out = []
    for name in sorted(os.listdir(root)):
        p = os.path.join(root, name)
        if os.path.isfile(p) and name.endswith((".py", ".sh")):
            out.append((p, f"code/{name}"))
    sub = os.path.join(root, "train")
    if os.path.isdir(sub):
        for name in sorted(os.listdir(sub)):
            p = os.path.join(sub, name)
            if os.path.isfile(p) and name.endswith(".py"):
                out.append((p, f"code/train/{name}"))
    return out


def run_dirs(stage_root, prefix, only):
    out = []
    for name in sorted(os.listdir(stage_root)):
        if not name.startswith(prefix):
            continue
        tag = name[len(prefix):]
        if only and tag not in only:
            continue
        d = os.path.join(stage_root, name)
        if os.path.isdir(d):
            out.append((tag, d))
    return out


def newest_tars(d, n):
    tars = []
    for name in os.listdir(d):
        if not name.endswith(".tar"):
            continue
        p = os.path.join(d, name)
        if os.path.isfile(p):
            tars.append((os.path.getmtime(p), name, p))
    tars.sort()
    return [(name, p) for _, name, p in tars[-n:]]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", default="Avra98/Sudoku_superposition")
    ap.add_argument("--code-dir", default=os.path.dirname(
        os.path.abspath(__file__)))
    ap.add_argument("--stage-root", default="/tmp")
    ap.add_argument("--stage-prefix", default="sudoku_hf_")
    ap.add_argument("--tags", default="",
                    help="comma-separated run tags; empty = every staging dir")
    ap.add_argument("--per-run", type=int, default=1,
                    help="how many of the newest tarballs to push per run")
    ap.add_argument("--token-file", default="/scratch/users/gatmiry/.hf_token")
    ap.add_argument("--skip-code", action="store_true")
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    token = os.environ.get("HF_TOKEN", "")
    if not token and os.path.isfile(args.token_file):
        with open(args.token_file) as f:
            token = f.read().strip()
    if not token:
        print("no HF token", file=sys.stderr)
        return 1

    from huggingface_hub import HfApi
    api = HfApi(token=token)
    try:
        existing = set(api.list_repo_files(args.repo, repo_type="model"))
    except Exception as e:
        print(f"could not list repo ({e}); nothing will be skipped")
        existing = set()

    plan = []
    if not args.skip_code:
        plan += code_files(args.code_dir)
    only = {t for t in args.tags.split(",") if t}
    for tag, d in run_dirs(args.stage_root, args.stage_prefix, only):
        for name, p in newest_tars(d, args.per_run):
            sub = "stages" if name.startswith("stage") else "steps"
            dest = f"runs/{tag}/{sub}/{name}"
            if name == "latest.tar":
                dest = f"runs/{tag}/latest.tar"
            elif dest in existing:
                print(f"skip {dest} (already in repo)")
                continue
            plan.append((p, dest))

    total = sum(os.path.getsize(p) for p, _ in plan)
    print(f"{len(plan)} files, {total / 1e9:.2f} GB", flush=True)
    if args.dry_run:
        for p, dest in plan:
            print(f"  {dest}  <- {p} ({os.path.getsize(p) / 1e6:.1f} MB)")
        return 0

    done_bytes = 0
    for i, (p, dest) in enumerate(plan, 1):
        sz = os.path.getsize(p)
        t0 = time.time()
        for attempt in (1, 2, 3):
            try:
                api.upload_file(path_or_fileobj=p, path_in_repo=dest,
                                repo_id=args.repo, repo_type="model",
                                commit_message=f"push {dest}")
                break
            except Exception as e:
                print(f"  attempt {attempt} failed for {dest}: {e}", flush=True)
                if attempt == 3:
                    print(f"GIVING UP on {dest}", flush=True)
                else:
                    time.sleep(10 * attempt)
        done_bytes += sz
        print(f"[{i}/{len(plan)}] {dest} {sz / 1e6:.1f} MB in "
              f"{time.time() - t0:.1f}s ({done_bytes / 1e9:.2f} GB done)",
              flush=True)
    print("push complete", flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())