init_dataset.py

"""
init_dataset.py — one-off bootstrap for the public Piclets dataset.

Creates the repo (if needed) and seeds the empty aggregate index files that the
server keeps in sync and the frontend reads directly. Safe to run once at setup.

Usage:
    HF_TOKEN=hf_xxx  python init_dataset.py            # create/seed (won't clobber)
    HF_TOKEN=hf_xxx  python init_dataset.py --verify    # print current totals
    HF_TOKEN=hf_xxx  python init_dataset.py --force      # re-seed empty indices (DESTROYS data)

The token needs WRITE access to the dataset. Set DATASET_REPO to override the
default repo id.
"""

import io
import os
import sys
import json

from huggingface_hub import HfApi, hf_hub_download, CommitOperationAdd

REPO = os.getenv("DATASET_REPO", "Fraser/Pictuary")
TOKEN = os.getenv("HF_API_KEY") or os.getenv("HF_TOKEN")

EMPTY_STATS = {"total_monsters": 0, "total_users": 0, "total_rarity_all": 0, "last_updated": None}
DATASET_README = f"""---
license: mit
tags:
  - piclets
  - game
---

# Piclets — shared monster dataset

Public database for the Piclets discovery game. Each real-world object maps to one
canonical monster, owned by its first discoverer.

Layout:
- `monsters/<key>.json` — one monster per normalized object name
- `images/<key>.webp` — the monster's art
- `users/<sub>.json` — a player's discoveries + summed rarity score
- `index/monsters.json` `index/feed.json` `index/leaderboard.json` `index/stats.json`
  — aggregate views the app reads directly

Written only by the Piclets Discovery Server. See that Space for details.
"""


def _json_add(path, obj):
    blob = json.dumps(obj, ensure_ascii=False, indent=2).encode("utf-8")
    return CommitOperationAdd(path_in_repo=path, path_or_fileobj=io.BytesIO(blob))


def verify(api):
    try:
        local = hf_hub_download(REPO, "index/stats.json", repo_type="dataset", token=TOKEN)
        with open(local, encoding="utf-8") as f:
            stats = json.load(f)
        print(f"[verify] {REPO}: {json.dumps(stats)}")
    except Exception as exc:
        print(f"[verify] could not read index/stats.json: {exc}")


def already_seeded(api) -> bool:
    try:
        hf_hub_download(REPO, "index/stats.json", repo_type="dataset", token=TOKEN)
        return True
    except Exception:
        return False


def main():
    if not TOKEN:
        sys.exit("Set HF_API_KEY (write access to the dataset) before running.")
    api = HfApi(token=TOKEN)

    if "--verify" in sys.argv:
        verify(api)
        return

    api.create_repo(REPO, repo_type="dataset", exist_ok=True, private=False)
    print(f"[init] repo ready: {REPO}")

    force = "--force" in sys.argv
    if already_seeded(api) and not force:
        print("[init] index files already exist — nothing to do. Use --force to reset (destroys data).")
        return

    api.create_commit(
        repo_id=REPO, repo_type="dataset",
        operations=[
            _json_add("index/monsters.json", []),
            _json_add("index/feed.json", []),
            _json_add("index/leaderboard.json", []),
            _json_add("index/stats.json", EMPTY_STATS),
            CommitOperationAdd(path_in_repo="README.md", path_or_fileobj=io.BytesIO(DATASET_README.encode("utf-8"))),
        ],
        commit_message="Initialize Piclets dataset",
    )
    print("[init] seeded empty indices. Done.")


if __name__ == "__main__":
    main()