persistence/backup.py

"""HF Dataset backup and restore for world state."""

from __future__ import annotations

import os
from pathlib import Path

from world.database import DB_PATH

HF_TOKEN = os.environ.get("HF_TOKEN", "")
HF_DATASET_REPO = os.environ.get("HF_DATASET_REPO", "")


def backup_database() -> bool:
    if os.environ.get("AG_SKIP_BACKUP") == "1":
        return False

    if not HF_TOKEN or not HF_DATASET_REPO:
        return False

    if not DB_PATH.exists():
        return False

    try:
        from huggingface_hub import HfApi
        api = HfApi(token=HF_TOKEN)

        try:
            api.create_repo(HF_DATASET_REPO, repo_type="dataset", private=True, exist_ok=True)
        except Exception:
            pass

        api.upload_file(
            path_or_fileobj=str(DB_PATH),
            path_in_repo="world.db",
            repo_id=HF_DATASET_REPO,
            repo_type="dataset",
            commit_message="World state backup",
        )
        return True
    except Exception as e:
        print(f"Backup failed: {e}")
        return False


def restore_database() -> bool:
    if not HF_TOKEN or not HF_DATASET_REPO:
        return False

    if DB_PATH.exists():
        return False

    try:
        from huggingface_hub import hf_hub_download

        DB_PATH.parent.mkdir(parents=True, exist_ok=True)
        downloaded = hf_hub_download(
            repo_id=HF_DATASET_REPO,
            filename="world.db",
            repo_type="dataset",
            token=HF_TOKEN,
            local_dir=str(DB_PATH.parent),
        )
        Path(downloaded).rename(DB_PATH)
        return True
    except Exception as e:
        print(f"Restore failed: {e}")
        return False