src/hf_storage.py

"""
Free persistence for HuggingFace Spaces via a PRIVATE Dataset repo (Option B).

HF Spaces have an ephemeral filesystem — anything written at runtime (run
history JSON + generated resumes) is wiped on every restart / rebuild / sleep.
To keep that data at zero cost we mirror the local `data/output/` tree into a
PRIVATE Dataset repo on the user's own HF account and sync it back down on
startup.

Activation:
    Set an HF *write* token as a Space secret named HF_TOKEN
    (or HUGGINGFACEHUB_API_TOKEN). Optionally set HF_RUNS_DATASET to choose the
    repo id; otherwise it defaults to "<your-username>/jaa-run-history".

Everything here is BEST-EFFORT: when no token is present (local dev or a
token-less Space) every function silently no-ops, and any network/API failure
is logged and swallowed so it can never break the pipeline.

Dataset layout mirrors the local tree under DATA_ROOT:
    run_history/run_*.json
    resumes/<date>/<file>.docx|.pdf
"""

from __future__ import annotations

import os
import logging
from pathlib import Path
from typing import Iterable, Optional

log = logging.getLogger("hf_storage")

DATA_ROOT = Path("data/output")          # local root mirrored to the dataset
_TOKEN_ENV = ("HF_TOKEN", "HUGGINGFACEHUB_API_TOKEN", "HF_API_TOKEN")
_repo_cache: Optional[str] = None


def _token() -> str:
    for k in _TOKEN_ENV:
        v = os.getenv(k)
        if v:
            return v.strip()
    return ""


def is_enabled() -> bool:
    """True only when an HF token is available (otherwise everything no-ops)."""
    return bool(_token())


def _repo_id(api=None) -> str:
    """Resolve the dataset repo id (cached for the process lifetime)."""
    global _repo_cache
    if _repo_cache is not None:
        return _repo_cache
    rid = os.getenv("HF_RUNS_DATASET", "").strip()
    if not rid:
        try:
            from huggingface_hub import HfApi
            api = api or HfApi()
            who = api.whoami(token=_token())
            name = who.get("name") or who.get("fullname") or ""
            rid = f"{name}/jaa-run-history" if name else ""
        except Exception as e:
            log.warning(f"hf_storage whoami failed: {e}")
            rid = ""
    _repo_cache = rid
    return rid


def _ensure_repo(api) -> str:
    rid = _repo_id(api)
    if not rid:
        return ""
    try:
        from huggingface_hub import create_repo
        create_repo(rid, repo_type="dataset", private=True,
                    exist_ok=True, token=_token())
    except Exception as e:
        log.warning(f"hf_storage create_repo failed: {e}")
    return rid


def sync_down(dest: Path = DATA_ROOT) -> bool:
    """Download the dataset snapshot into local data/output (merge). Best-effort."""
    if not is_enabled():
        return False
    try:
        from huggingface_hub import snapshot_download, HfApi
        api = HfApi()
        rid = _repo_id(api)
        if not rid:
            return False
        dest.mkdir(parents=True, exist_ok=True)
        snapshot_download(
            repo_id=rid, repo_type="dataset",
            local_dir=str(dest), token=_token(),
            allow_patterns=["run_history/*", "resumes/**", "reports/*"],
        )
        log.info(f"hf_storage synced run data down from {rid}")
        return True
    except Exception as e:
        # RepositoryNotFound on first ever run is expected → just skip.
        log.warning(f"hf_storage sync_down skipped: {e}")
        return False


def push_run(run_json_path: str, resume_paths: Iterable[str] = (),
             root: Path = DATA_ROOT) -> bool:
    """Upload one run's history JSON + its resume files in a single commit.

    Paths are stored relative to `root` so they restore to the same local
    location on sync_down. Best-effort: returns False (and logs) on any failure.
    """
    if not is_enabled():
        return False
    try:
        from huggingface_hub import HfApi, CommitOperationAdd
        api = HfApi()
        rid = _ensure_repo(api)
        if not rid:
            return False

        ops = []
        seen = set()
        for f in [run_json_path, *(resume_paths or [])]:
            if not f:
                continue
            p = Path(f)
            if not p.exists() or str(p) in seen:
                continue
            seen.add(str(p))
            try:
                rel = p.resolve().relative_to(root.resolve()).as_posix()
            except Exception:
                rel = f"resumes/{p.name}"
            ops.append(CommitOperationAdd(path_in_repo=rel, path_or_fileobj=str(p)))

        if not ops:
            return False
        api.create_commit(
            repo_id=rid, repo_type="dataset", operations=ops,
            commit_message=f"run {Path(run_json_path).stem} ({len(ops)} files)",
            token=_token(),
        )
        log.info(f"hf_storage pushed {len(ops)} file(s) to {rid}")
        return True
    except Exception as e:
        log.warning(f"hf_storage push_run skipped: {e}")
        return False