"""Full-project backup to Hugging Face (continuity snapshot).

Pushes the entire workspace (code, data, notes, pipeline, checkpoints) to
FerrellSyntheticIntelligence/fsi-anomaly so work can continue on another
machine. Skips .venv and python caches. Resume-safe: a local manifest
(logs/hf_backup_manifest.json) records uploaded files by sha256, and files
already present on the Hub are skipped, so re-running after an interruption
continues where it stopped. Progress is visible per commit.

Usage:
  HF_TOKEN=hf_xxx .venv/bin/python hf_backup.py            # everything
  HF_TOKEN=hf_xxx .venv/bin/python hf_backup.py --stage ckpt   # checkpoints only
"""

import argparse
import hashlib
import json
import os
import sys
from pathlib import Path

from huggingface_hub import CommitOperationAdd, HfApi

HERE = Path(__file__).resolve().parent
EXCLUDED_DIRS = {".venv", "__pycache__", ".pytest_cache"}
EXCLUDED_SUFFIXES = {".pyc"}
MANIFEST = HERE / "logs" / "hf_backup_manifest.json"
BATCH_FILES = 100
BATCH_BYTES = 800_000_000  # ~800MB per commit (safer on tablet network)


def sha256(path: Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def iter_files(stage: str):
    for p in sorted(HERE.rglob("*")):
        if not p.is_file():
            continue
        rel = p.relative_to(HERE).as_posix()
        parts = rel.split("/")
        if any(part in EXCLUDED_DIRS for part in parts):
            continue
        if p.suffix in EXCLUDED_SUFFIXES:
            continue
        if rel == "logs/hf_backup_manifest.json":
            continue
        if stage == "small" and parts[0] == "ckpt":
            continue
        if stage == "ckpt" and parts[0] != "ckpt":
            continue
        yield rel, p


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", default="FerrellSyntheticIntelligence/fsi-anomaly")
    ap.add_argument("--stage", choices=["all", "small", "ckpt"], default="all")
    args = ap.parse_args()

    token = os.environ.get("HF_TOKEN")
    if not token:
        sys.exit("HF_TOKEN env var required")
    api = HfApi(token=token)

    try:
        api.repo_info(args.repo, repo_type="model")
        print(f"repo exists: {args.repo}", flush=True)
    except Exception:
        api.create_repo(args.repo, private=True, repo_type="model")
        print(f"created repo: {args.repo} (private)", flush=True)

    manifest = {}
    if MANIFEST.exists():
        try:
            manifest = json.loads(MANIFEST.read_text())
        except json.JSONDecodeError:
            manifest = {}

    remote = set(api.list_repo_files(args.repo, repo_type="model"))
    print(f"remote files already present: {len(remote)}", flush=True)

    ops = []
    batch_bytes = 0
    n_uploaded = 0
    n_skipped = 0

    def flush(reason):
        nonlocal ops, batch_bytes, n_uploaded
        if not ops:
            return
        api.create_commit(
            repo_id=args.repo,
            operations=ops,
            commit_message=f"backup {args.stage}: {len(ops)} files ({reason})",
            repo_type="model",
        )
        for op in ops:
            manifest[op.path_in_repo] = sha256(Path(op.path_or_fileobj))
        MANIFEST.write_text(json.dumps(manifest, indent=0))
        n_uploaded += len(ops)
        print(f"committed {len(ops)} files -> {n_uploaded} total ({reason})", flush=True)
        ops = []
        batch_bytes = 0

    for rel, p in iter_files(args.stage):
        if rel in remote or manifest.get(rel) == sha256(p):
            n_skipped += 1
            continue
        ops.append(CommitOperationAdd(path_in_repo=rel, path_or_fileobj=str(p)))
        batch_bytes += p.stat().st_size
        if len(ops) >= BATCH_FILES or batch_bytes >= BATCH_BYTES:
            flush("batch")

    flush("final")
    print(f"DONE stage={args.stage}: uploaded={n_uploaded} skipped={n_skipped}", flush=True)
    print(f"repo: https://huggingface.co/{args.repo}", flush=True)


if __name__ == "__main__":
    main()