src/hf_sync.py

"""
Sync a completed run's manifest to a shared Hugging Face Dataset repo.

Why this exists: HF Spaces have ephemeral local storage -- anything written
to disk during a Space session disappears on restart/sleep. Without this,
every audit run (images included) would vanish once the Space goes idle.
Instead, each run uploads its manifest -- WITH each row's image embedded
via the `datasets` library's Image feature -- as its own Parquet file into
a dataset repo. HF's dataset viewer renders those embedded images as
thumbnails directly. Multiple Parquet files in one repo are treated as one
combined dataset, so there's no read-modify-write step and no merge
conflicts if two runs finish at the same time -- same append-only
philosophy as ManifestWriter's chunking.

The local manifest.parquet (written by ManifestWriter) stays lightweight
and file-path-based -- this module builds a separate, image-embedded copy
only at upload time, so local runs/iteration aren't slowed down by it.

Auth: writing requires a write-scoped HF token. On a deployed Space this
comes from an HF_TOKEN secret set by the Space owner (never exposed to
visitors) -- a visitor's run triggers the write, but always under the
Space owner's account.
"""
import os
import uuid
from datetime import datetime, timezone


def _read_image_bytes(file_path):
    with open(file_path, "rb") as f:
        return {"bytes": f.read(), "path": None}


def sync_manifest_to_hub(df, repo_id=None, token=None, include_images=True, run_id=None,
                         on_progress=None,
                         progress_prefix="Images generated… uploading to shared dataset"):
    """
    Uploads `df` (one run's manifest) as a Parquet file in a shared
    HF dataset repo. Silently no-ops if not configured, so local/dev runs
    don't require HF credentials.

    repo_id: e.g. "niftymuch/bias-audit-manifests" (defaults to env HF_DATASET_REPO)
    token: write-scoped HF token (defaults to env HF_TOKEN)
    include_images: embed each row's image (read from its file_path) into
        the uploaded Parquet via datasets.Image, so the row is fully
        self-contained. Set False to upload metadata only.
    run_id: stable id for this run (e.g. local folder name). When set, the
        file is written to runs/{run_id}.parquet and later uploads with the
        same run_id overwrite that file (so Analyze updates Generate's save
        instead of creating a second copy). When omitted, a unique timestamp
        id is used (legacy append-only behavior).
    on_progress: optional callable(frac: float, desc: str) for UI progress
        during packing + upload (e.g. gr.Progress).
    progress_prefix: leading text for progress descriptions.

    Returns the dataset URL if the upload happened, else None.
    """
    repo_id = repo_id or os.environ.get("HF_DATASET_REPO")
    token = token or os.environ.get("HF_TOKEN")
    if not repo_id or not token:
        return None

    def _report(frac, detail=None):
        if on_progress is not None:
            desc = progress_prefix if not detail else f"{progress_prefix} ({detail})"
            on_progress(frac, desc)

    from huggingface_hub import HfApi

    _report(0.0)

    api = HfApi(token=token)
    api.create_repo(repo_id, repo_type="dataset", exist_ok=True, private=True)

    if not run_id:
        run_id = f"{datetime.now(timezone.utc).strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}"
    # Keep path-safe (local folder names are already safe; sanitize just in case)
    run_id = "".join(c if c.isalnum() or c in "-_" else "_" for c in str(run_id))

    tmp_dir = os.environ.get("TMPDIR") or os.environ.get("TEMP") or "/tmp"
    tmp_path = os.path.join(tmp_dir, f"{run_id}.parquet")

    if include_images:
        from datasets import Dataset, Image

        upload_df = df.copy()
        paths = upload_df["file_path"].tolist()
        n = len(paths)
        images = []
        # Packing images is usually the slow local part of the upload.
        for i, path in enumerate(paths):
            images.append(_read_image_bytes(path))
            if i == 0 or (i + 1) % 25 == 0 or (i + 1) == n:
                _report(
                    0.05 + 0.70 * ((i + 1) / n if n else 1),
                    f"packing {i + 1:,}/{n:,}",
                )
        upload_df["image"] = images
        _report(0.80, "writing file")
        dataset = Dataset.from_pandas(upload_df, preserve_index=False)
        dataset = dataset.cast_column("image", Image())
        dataset.to_parquet(tmp_path)
    else:
        _report(0.80, "writing file")
        df.to_parquet(tmp_path, index=False)

    try:
        _report(0.90, "sending to Hugging Face")
        api.upload_file(
            path_or_fileobj=tmp_path,
            path_in_repo=f"runs/{run_id}.parquet",
            repo_id=repo_id,
            repo_type="dataset",
        )
    finally:
        if os.path.exists(tmp_path):
            os.remove(tmp_path)

    _report(1.0, "complete")
    return f"https://huggingface.co/datasets/{repo_id}"