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 database_is_valid() -> bool:
    """True when world.db exists and has the core schema."""
    if not DB_PATH.exists():
        return False
    try:
        if DB_PATH.stat().st_size < 512:
            return False
        import sqlite3

        conn = sqlite3.connect(DB_PATH)
        try:
            row = conn.execute(
                "SELECT name FROM sqlite_master WHERE type='table' AND name='world_state'"
            ).fetchone()
            return row is not None
        finally:
            conn.close()
    except Exception:
        return False


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 database_is_valid():
        return False

    if DB_PATH.exists():
        DB_PATH.unlink(missing_ok=True)

    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),
        )
        downloaded_path = Path(downloaded)
        if downloaded_path.resolve() != DB_PATH.resolve():
            downloaded_path.replace(DB_PATH)
        return database_is_valid()
    except Exception as e:
        print(f"Restore failed: {e}")
        return False