tipibot/scripts/reset_pb_collections.py

"""Destructively recreate economy PocketBase collections for dev + economy profiles.

Usage:
    python scripts/reset_pb_collections.py --confirm

This will DELETE and recreate the collections configured by:
- PB_ECONOMY_COLLECTION_DEV
- PB_ECONOMY_COLLECTION_ECONOMY
"""

from __future__ import annotations

import argparse
import asyncio
import sys
from pathlib import Path

import aiohttp
from dotenv import load_dotenv

sys.path.insert(0, str(Path(__file__).parent.parent))
load_dotenv()

import config  # noqa: E402

PB_URL = config.PB_URL
PB_ADMIN_EMAIL = config.PB_ADMIN_EMAIL
PB_ADMIN_PASSWORD = config.PB_ADMIN_PASSWORD


def _text_field(name: str, required: bool = False) -> dict:
    return {
        "name": name,
        "type": "text",
        "required": required,
        "options": {"min": None, "max": None, "pattern": ""},
    }


def _number_field(name: str) -> dict:
    return {
        "name": name,
        "type": "number",
        "required": False,
        "options": {"min": None, "max": None, "noDecimal": False},
    }


def _bool_field(name: str) -> dict:
    return {"name": name, "type": "bool", "required": False}


def _json_field(name: str) -> dict:
    return {"name": name, "type": "json", "required": False}


def _collection_payload(name: str) -> dict:
    fields = [
        _text_field("user_id", required=True),
        _number_field("balance"),
        _number_field("exp"),
        _number_field("daily_streak"),
        _text_field("last_daily"),
        _text_field("last_work"),
        _text_field("last_beg"),
        _text_field("last_crime"),
        _text_field("last_rob"),
        _text_field("last_heist"),
        _text_field("last_streak_date"),
        _text_field("jailed_until"),
        _text_field("last_fish"),
        _json_field("items"),
        _json_field("item_uses"),
        _json_field("reminders"),
        _json_field("prestige_upgrades"),
        _json_field("fish_book"),
        _json_field("fish_inventory"),
        _bool_field("eco_banned"),
        _bool_field("jailbreak_used"),
        _number_field("heist_global_cd_until"),
        _number_field("peak_balance"),
        _number_field("lifetime_earned"),
        _number_field("lifetime_lost"),
        _number_field("work_count"),
        _number_field("beg_count"),
        _number_field("total_wagered"),
        _number_field("biggest_win"),
        _number_field("biggest_loss"),
        _number_field("slots_jackpots"),
        _number_field("crimes_attempted"),
        _number_field("crimes_succeeded"),
        _number_field("times_jailed"),
        _number_field("total_bail_paid"),
        _number_field("heists_joined"),
        _number_field("heists_won"),
        _number_field("total_given"),
        _number_field("total_received"),
        _number_field("best_daily_streak"),
        _number_field("prestige_level"),
        _number_field("prestige_points"),
        _number_field("season_total_exp"),
        _number_field("total_fish_caught"),
    ]

    return {
        "name": name,
        "type": "base",
        "fields": fields,
        "listRule": None,
        "viewRule": None,
        "createRule": None,
        "updateRule": None,
        "deleteRule": None,
    }


async def _auth_token(session: aiohttp.ClientSession) -> str:
    async with session.post(
        f"{PB_URL}/api/collections/_superusers/auth-with-password",
        json={"identity": PB_ADMIN_EMAIL, "password": PB_ADMIN_PASSWORD},
    ) as resp:
        if resp.status != 200:
            raise RuntimeError(f"Auth failed ({resp.status}): {await resp.text()}")
        return (await resp.json())["token"]


async def _delete_if_exists(session: aiohttp.ClientSession, headers: dict[str, str], name: str) -> None:
    async with session.get(f"{PB_URL}/api/collections/{name}", headers=headers) as resp:
        if resp.status == 404:
            print(f"[SKIP] {name} does not exist")
            return
        if resp.status != 200:
            raise RuntimeError(f"Could not fetch {name} ({resp.status}): {await resp.text()}")

    async with session.delete(f"{PB_URL}/api/collections/{name}", headers=headers) as resp:
        if resp.status not in (200, 204):
            raise RuntimeError(f"Delete failed for {name} ({resp.status}): {await resp.text()}")
    print(f"[DELETE] {name}")


async def _create_collection(session: aiohttp.ClientSession, headers: dict[str, str], name: str) -> None:
    payload = _collection_payload(name)
    async with session.post(f"{PB_URL}/api/collections", json=payload, headers=headers) as resp:
        if resp.status not in (200, 201):
            raise RuntimeError(f"Create failed for {name} ({resp.status}): {await resp.text()}")
    print(f"[CREATE] {name}")


async def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--confirm", action="store_true", help="Required flag to run destructive reset")
    args = parser.parse_args()

    if not args.confirm:
        raise SystemExit("Refusing to run without --confirm (this operation deletes collections).")

    targets = []
    for name in [config.PB_ECONOMY_COLLECTION_DEV, config.PB_ECONOMY_COLLECTION_ECONOMY]:
        if name and name not in targets:
            targets.append(name)

    if not targets:
        raise SystemExit("No target collections configured.")

    timeout = aiohttp.ClientTimeout(total=20)
    async with aiohttp.ClientSession(timeout=timeout) as session:
        token = await _auth_token(session)
        headers = {"Authorization": token}

        for name in targets:
            await _delete_if_exists(session, headers, name)
            await _create_collection(session, headers, name)

    print("\nDone. Collections recreated:")
    for name in targets:
        print(f" - {name}")


if __name__ == "__main__":
    asyncio.run(main())