download_examples.py

"""Download example media from the LatticeSemi/Demo-Examples-v1.0 dataset.

Runs at Docker build time (after ``COPY --chown=user . /app``) to populate the
demo's ``examples*/`` subfolders from the shared HF dataset.  Mirrors the
dataset layout: every file under ``<DEMO>/`` in the dataset is copied to the
same relative path locally (``<DEMO>/examples/foo.mp4`` -> ``./examples/foo.mp4``).

Auth order — same as ``install_eve.py``:

1. Docker BuildKit secret ``MODEL_ACCESS_TOKEN``
2. ``HF_TOKEN`` env / build arg
3. No token — works only if the dataset is public.

Local-dev fallback: if the download fails (offline, no token, dataset private),
but a sibling ``examples*/`` folder already contains files, succeed with a
warning so manual placement / pre-existing files keep working.  If neither
the download nor any local file is present, crash loud with the dataset URL
and the token-fix hint.
"""

import os
import shutil
import sys
from pathlib import Path

from huggingface_hub import HfApi, hf_hub_download

DEMO = "eve_gmod"
DATASET = "LatticeSemi/Demo-Examples-v1.0"
SECRET_PATH = "/run/secrets/MODEL_ACCESS_TOKEN"


def get_token() -> tuple[str | None, str]:
    """Return an HF token from the first available source, or None.

    Falls through past the secret file on PermissionError — BuildKit mounts
    the secret as root by default; the Dockerfile must pass ``uid=1000`` for
    a non-root user to read it.  Falling through lets ``HF_TOKEN`` still work
    if the mount permissions are misconfigured.
    """
    if os.path.isfile(SECRET_PATH):
        try:
            with open(SECRET_PATH) as f:
                token = f.read().strip()
                if token:
                    return token, "build secret (MODEL_ACCESS_TOKEN)"
        except PermissionError as exc:
            print(
                f"[download_examples] WARNING: cannot read {SECRET_PATH}: {exc}. "
                f"Add ``uid=1000`` to the secret mount in the Dockerfile.",
                file=sys.stderr,
            )
    token = os.environ.get("HF_TOKEN", "").strip()
    if token:
        return token, "build arg (HF_TOKEN)"
    return None, "no auth"


def has_local_examples() -> bool:
    """True iff any ``examples*`` sibling dir contains at least one file."""
    cwd = Path(".")
    if not cwd.is_dir():
        return False
    for entry in cwd.iterdir():
        if entry.is_dir() and entry.name.startswith("examples"):
            if any(p.is_file() for p in entry.rglob("*")):
                return True
    return False


def main() -> int:
    token, source = get_token()
    print(f"[download_examples] demo={DEMO} dataset={DATASET} auth={source}")

    download_error: str | None = None
    downloaded = 0
    try:
        api = HfApi()
        files = api.list_repo_files(DATASET, repo_type="dataset", token=token)
        prefix = f"{DEMO}/"
        wanted = [f for f in files if f.startswith(prefix)]
        if not wanted:
            download_error = f"dataset has no files under '{prefix}'"
        for remote in wanted:
            cached = hf_hub_download(
                repo_id=DATASET,
                repo_type="dataset",
                filename=remote,
                token=token,
            )
            local = Path(remote[len(prefix) :])
            local.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(cached, local)
            downloaded += 1
    except Exception as exc:  # noqa: BLE001 — broad catch by design at build time
        download_error = f"{type(exc).__name__}: {exc}"

    if downloaded > 0:
        print(f"[download_examples] OK: downloaded {downloaded} files")
        return 0

    if has_local_examples():
        print(
            f"[download_examples] WARNING: download failed — {download_error}.\n"
            f"  Falling back to existing local example files.",
            file=sys.stderr,
        )
        return 0

    print(
        f"FATAL: no example files for '{DEMO}'.\n"
        f"  Download from {DATASET} failed: {download_error or 'no files retrieved'}\n"
        f"  Fix one of:\n"
        f"    - HF Space: set MODEL_ACCESS_TOKEN secret with read access to the dataset.\n"
        f"    - Local dev: run `hf auth login`, or manually drop files into the demo's\n"
        f"      examples*/ subfolders from\n"
        f"      https://huggingface.co/datasets/{DATASET}/tree/main/{DEMO}",
        file=sys.stderr,
    )
    return 1


if __name__ == "__main__":
    sys.exit(main())