"""HF ckpt watcher — rate-limit 친화 버전 (HF 커밋 128/hour 제한 대응).

핵심: 커밋 수 최소화.
- 무거운 동기화(ckpt+logs+plot)는 '새 10k ckpt가 생겼을 때만' 수행.
- 샘플/ckpt는 upload_folder로 '폴더당 1커밋' (파일별 커밋 금지 — 이전 버전이 429 유발).
- ckpt: 최신 KEEP개만 유지 (로테이션). 샘플: 전부 유지.
usage: python hf_ckpt_watcher.py   (setsid nohup / harness background)
"""
import os, re, time
from huggingface_hub import HfApi

TOKEN = os.environ.get("HF_TOKEN", "<SET_HF_TOKEN_ENV>")
REPO = "a12s12/spatial-semanticist-L-migration"
BASE = "/NHNHOME/WORKSPACE/0226010398_A/sr_diffusion/clevr_sudoku/semanticist/output/tokenizer/models_l_spatial"
MODELS_DIR = os.path.join(BASE, "models")
IMAGES_DIR = os.path.join(BASE, "images")
LOGDIR = os.path.join(BASE, "logs/semanticist")
KEEP = 2
POLL_SEC = 300           # 5분 폴링 (커밋 절약)
REQUIRED = ["model.safetensors", "optimizer.bin"]

api = HfApi(token=TOKEN)


def local_complete_steps():
    out = []
    if not os.path.isdir(MODELS_DIR):
        return out
    for name in os.listdir(MODELS_DIR):
        m = re.fullmatch(r"step(\d+)", name)
        if m:
            d = os.path.join(MODELS_DIR, name)
            if all(os.path.exists(os.path.join(d, r)) for r in REQUIRED):
                out.append(int(m.group(1)))
    return sorted(out)


def hf_steps():
    steps = set()
    for f in api.list_repo_files(REPO):
        m = re.match(r"step(\d+)/", f)
        if m:
            steps.add(int(m.group(1)))
    return sorted(steps)


def sync_samples():
    """IMAGES_DIR 전체를 samples_all_steps/ 로 upload_folder (변경분만 1커밋)."""
    if os.path.isdir(IMAGES_DIR):
        api.upload_folder(folder_path=IMAGES_DIR, path_in_repo="samples_all_steps",
                          repo_id=REPO, repo_type="model",
                          allow_patterns=["*.jpg"],
                          commit_message="sync recon samples")


def sync_logs():
    """loss 그래프 갱신 + tensorboard event (새 ckpt milestone에서만 호출)."""
    try:
        import make_loss_plot
        make_loss_plot.main()
        png = os.path.join(BASE, "loss_curves.png")
        if os.path.exists(png):
            api.upload_file(path_or_fileobj=png, path_in_repo="loss_curves.png",
                            repo_id=REPO, repo_type="model",
                            commit_message="update loss curves")
    except Exception as e:
        print(f"[watcher] plot skip: {e}", flush=True)
    if os.path.isdir(LOGDIR):
        api.upload_folder(folder_path=LOGDIR, path_in_repo="logs",
                          repo_id=REPO, repo_type="model",
                          allow_patterns=["events*"],
                          commit_message="update tb logs")


def sync_ckpts():
    """최신 KEEP개 로컬 ckpt만 HF에 유지. 반환: 최신 step (없으면 None)."""
    local = local_complete_steps()
    if not local:
        return None
    on_hf = set(hf_steps())
    target = set(local[-KEEP:])
    for s in sorted(target - on_hf):
        print(f"[watcher] uploading step{s} ...", flush=True)
        api.upload_folder(folder_path=os.path.join(MODELS_DIR, f"step{s}"),
                          path_in_repo=f"step{s}", repo_id=REPO, repo_type="model",
                          commit_message=f"ckpt step{s}")
        print(f"[watcher] step{s} uploaded", flush=True)
    pruned = sorted(on_hf - target)
    for old in pruned:
        print(f"[watcher] pruning HF step{old}", flush=True)
        api.delete_folder(path_in_repo=f"step{old}", repo_id=REPO, repo_type="model",
                          commit_message=f"prune step{old} (keep newest {KEEP})")
    # ★ private repo 저장한도 대응: 로테이션으로 지운 ckpt는 git history에 LFS blob이
    #   남아 용량을 잡아먹는다. prune이 있었으면 history를 squash해서 실제로 회수.
    if pruned:
        try:
            api.super_squash_history(repo_id=REPO, repo_type="model")
            print(f"[watcher] squashed history (freed rotated ckpt LFS)", flush=True)
        except Exception as e:
            print(f"[watcher] squash skip: {str(e)[:120]}", flush=True)
    return local[-1]


def main():
    print(f"[watcher] start (rate-safe). keep={KEEP} poll={POLL_SEC}s repo={REPO}", flush=True)
    last_milestone = None
    while True:
        try:
            local = local_complete_steps()
            newest = local[-1] if local else None
            # 새 10k ckpt가 생겼을 때만 무거운 전체 동기화 (커밋 절약)
            if newest is not None and newest != last_milestone:
                print(f"[watcher] new milestone step{newest} -> full sync", flush=True)
                new_top = sync_ckpts()
                sync_samples()
                sync_logs()
                last_milestone = new_top
                print(f"[watcher] synced. HF ckpt newest {KEEP} up to step{new_top}", flush=True)
            else:
                # milestone 사이엔 샘플만 가볍게 (upload_folder=변경분 1커밋)
                sync_samples()
        except Exception as e:
            print(f"[watcher] error (다음 폴링에 재시도): {str(e)[:200]}", flush=True)
        time.sleep(POLL_SEC)


if __name__ == "__main__":
    main()