#!/usr/bin/env python3
"""Auto-sync the newest GPU training checkpoint to the Hugging Face repo (add-only, never deletes).

Every --interval seconds (default 3 h):
  1. pick the active run dir: the subdir of checkpoints/ (excluding init*, smoke*, quickpod*) whose newest *.pt
     is the most recent (so a new OUT dir, e.g. a bilingual resume, is picked up automatically);
  2. take its newest checkpoint (highest step), skip it if it is still being written
     (size/mtime must be stable for --stable-secs and the zip central directory must be readable);
  3. if checkpoints/gpu_<run_dir>/<file> is not on the Hub yet, upload it plus the run's progress.tsv in one commit
     whose message carries the step and the matching val_ce / val_field_alone.
Network errors are retried with backoff; a failed cycle never stops the loop.
Auth: huggingface_hub reads ~/.cache/huggingface/token (mode 600); the token is never printed or logged.
Run:  nohup nice -n 19 ionice -c3 ~/hfsync-venv/bin/python -u scripts/hf_ckpt_sync.py --loop >> logs/hf_sync.log 2>&1 < /dev/null &
      (or: bash scripts/hf_ckpt_sync.sh start|status|stop|once)
"""
import argparse, csv, datetime, fcntl, glob, json, os, re, sys, time, zipfile

os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
os.environ.setdefault("OMP_NUM_THREADS", "1")
ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
EXCLUDE = re.compile(r"^(init|smoke|quickpod)", re.I)
STEP = re.compile(r"atom_native_step_(\d+)_")


def log(msg):
    print(f"{datetime.datetime.now(datetime.timezone.utc).isoformat(timespec='seconds')} {msg}", flush=True)


def step_of(path):
    m = STEP.search(os.path.basename(path))
    return int(m.group(1)) if m else -1


def active_run_dir(ckroot):
    best = None
    for d in sorted(glob.glob(os.path.join(ckroot, "*/"))):
        name = os.path.basename(d.rstrip("/"))
        if EXCLUDE.match(name):
            continue
        pts = glob.glob(os.path.join(d, "*.pt"))
        if not pts:
            continue
        newest = max(os.path.getmtime(p) for p in pts)
        if best is None or newest > best[0]:
            best = (newest, d.rstrip("/"))
    return best[1] if best else None


def newest_ckpt(run_dir):
    pts = glob.glob(os.path.join(run_dir, "*.pt"))
    return max(pts, key=lambda p: (step_of(p), os.path.getmtime(p))) if pts else None


def is_complete(path, stable_secs):
    try:
        s1 = os.stat(path)
        time.sleep(stable_secs)
        s2 = os.stat(path)
        if (s1.st_size, s1.st_mtime) != (s2.st_size, s2.st_mtime) or s2.st_size == 0:
            return False
        with zipfile.ZipFile(path) as z:  # torch .pt = zip; readable central directory => fully written
            z.namelist()
        return True
    except Exception as e:  # noqa: BLE001
        log(f"not ready: {os.path.basename(path)} ({type(e).__name__})")
        return False


def val_for_step(tsv, step):
    """(val_ce, val_field_alone, tick) of the row for `step` (else the last row)."""
    try:
        with open(tsv, newline="", encoding="utf-8") as f:
            rows = list(csv.DictReader(f, delimiter="\t"))
    except OSError:
        return None, None, None
    if not rows:
        return None, None, None
    row = next((r for r in rows if r.get("tick") == str(step)), rows[-1])
    return row.get("val_ce"), row.get("val_field_alone"), row.get("tick")


def retry(fn, what, tries=6, base=30):
    for i in range(tries):
        try:
            return fn()
        except Exception as e:  # noqa: BLE001
            msg = str(e).splitlines()[0][:300] if str(e) else ""
            log(f"{what} failed (attempt {i + 1}/{tries}): {type(e).__name__}: {msg}")
            if i == tries - 1:
                raise
            time.sleep(min(base * 2 ** i, 1800))


def cycle(args, api):
    ckroot = os.path.join(ROOT, "checkpoints")
    run_dir = active_run_dir(ckroot)
    if not run_dir:
        log("no run dir with checkpoints yet")
        return
    ck = newest_ckpt(run_dir)
    run = os.path.basename(run_dir)
    dest_dir = f"checkpoints/gpu_{run}"
    dest = f"{dest_dir}/{os.path.basename(ck)}"
    remote = set(retry(lambda: api.list_repo_files(args.repo), "list_repo_files"))
    if dest in remote:
        log(f"up to date: {dest} already on the Hub")
        return
    if not is_complete(ck, args.stable_secs):
        log(f"skip this cycle: {os.path.basename(ck)} still being written")
        return
    from huggingface_hub import CommitOperationAdd
    step = step_of(ck)
    tsv = os.path.join(run_dir, "progress.tsv")
    vce, vfa, row_tick = val_for_step(tsv, step)
    ops = [CommitOperationAdd(path_in_repo=dest, path_or_fileobj=ck)]
    if os.path.exists(tsv):
        ops.append(CommitOperationAdd(path_in_repo=f"{dest_dir}/progress.tsv", path_or_fileobj=tsv))
    vtxt = f"val_ce {vce} / field-alone {vfa} (progress row tick {row_tick})" if vce else "val_ce n/a"
    msg = f"GPU auto-sync: {run} step {step}, {vtxt}"
    size = os.path.getsize(ck)
    t0 = time.time()
    info = retry(lambda: api.create_commit(args.repo, operations=ops, commit_message=msg), "create_commit")
    remote = set(retry(lambda: api.list_repo_files(args.repo), "list_repo_files (verify)"))
    ok = dest in remote
    log(f"{'UPLOADED' if ok else 'COMMIT MADE BUT FILE MISSING ON HUB (check .gitignore)'}: {dest} "
        f"({size} B, {time.time() - t0:.0f}s) step={step} {vtxt} commit={getattr(info, 'oid', '')}")
    with open(os.path.join(ROOT, "logs", "hf_sync_state.json"), "w") as f:
        json.dump({"last": dest, "step": step, "val_ce": vce, "val_field_alone": vfa, "ok": ok,
                   "time": datetime.datetime.now(datetime.timezone.utc).isoformat(timespec="seconds")}, f)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", default="thefinalboss/atom-ai")
    ap.add_argument("--interval", type=int, default=3 * 3600)
    ap.add_argument("--stable-secs", type=int, default=30)
    ap.add_argument("--loop", action="store_true", help="run forever (default: one cycle)")
    args = ap.parse_args()
    try:
        os.nice(19)
    except OSError:
        pass
    os.makedirs(os.path.join(ROOT, "logs"), exist_ok=True)
    lock = open(os.path.join(ROOT, "logs", "hf_sync.lock"), "w")
    try:
        fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except OSError:
        log("another hf_ckpt_sync holds logs/hf_sync.lock; exiting")
        sys.exit(1)
    if args.loop:
        with open(os.path.join(ROOT, "logs", "hf_sync.pid"), "w") as f:
            f.write(str(os.getpid()))
    from huggingface_hub import HfApi
    api = HfApi()
    who = retry(lambda: api.whoami(), "whoami")
    log(f"start pid={os.getpid()} repo={args.repo} user={who.get('name')} interval={args.interval}s loop={args.loop}")
    while True:
        try:
            cycle(args, api)
        except Exception as e:  # noqa: BLE001
            log(f"cycle error: {type(e).__name__}: {str(e).splitlines()[0][:300] if str(e) else ''}")
        if not args.loop:
            break
        time.sleep(args.interval)


if __name__ == "__main__":
    main()