Files
tipibot/core/economy/store.py
2026-07-27 00:49:34 +03:00

385 lines
14 KiB
Python

"""Shared foundation: user records, locks, time, cooldowns, txn log."""
from __future__ import annotations
import asyncio
import functools
import logging
from datetime import datetime, timedelta, timezone
from typing import TypedDict
import aiohttp
from .. import pb_client
from ..pb_client import DatabaseError
from ..emoji import EMOJI as E
def _clock() -> datetime:
"""Actual time source - a seam so tests can freeze time everywhere at once."""
return datetime.now(tz=timezone.utc)
def _now() -> datetime:
return _clock()
_txn_log = logging.getLogger("tipiCOIN.txn")
def _txn(event: str, **fields) -> None:
"""Log a single economy transaction to the transactions logger."""
body = " ".join(f"{k}={v}" for k, v in fields.items())
_txn_log.info("%-16s %s", event, body)
# Per-profile emoji values live in core/emoji.py; add new IDs there.
COIN = E["TipiCOIN"]
PP_EMOJI = E["TipiFIRE"]
# ---------------------------------------------------------------------------
# Prestige shop catalogue
# ---------------------------------------------------------------------------
class PrestigeItem(TypedDict):
emoji: str
max_level: int
pp_cost: int
effect: float
PRESTIGE_SHOP: dict[str, PrestigeItem] = {
"coin_mult": {
"emoji": E["TipiCOIN"],
"max_level": 5,
"pp_cost": 5,
"effect": 0.08,
},
"exp_mult": {
"emoji": "",
"max_level": 5,
"pp_cost": 5,
"effect": 0.08,
},
"daily_plus": {
"emoji": "📅",
"max_level": 3,
"pp_cost": 7,
"effect": 0.20,
},
"work_plus": {
"emoji": "💼",
"max_level": 3,
"pp_cost": 7,
"effect": 0.20,
},
}
# ---------------------------------------------------------------------------
# Cooldowns
# ---------------------------------------------------------------------------
COOLDOWNS: dict[str, timedelta] = {
"daily": timedelta(hours=20),
"work": timedelta(hours=1),
"beg": timedelta(minutes=5),
"crime": timedelta(hours=2),
"rob": timedelta(hours=2),
"fish": timedelta(minutes=2),
}
JAIL_DURATION = timedelta(minutes=30)
HEIST_JAIL = timedelta(hours=1, minutes=30)
# ---------------------------------------------------------------------------
# User schema
# ---------------------------------------------------------------------------
class UserData(TypedDict, total=False):
balance: int
exp: int # lifetime EXP (resets each season)
last_daily: str | None
last_work: str | None
last_beg: str | None
last_crime: str | None
last_rob: str | None
last_heist: str | None
daily_streak: int
last_streak_date: str | None # ISO date "YYYY-MM-DD"
items: list[str]
item_uses: dict # {item_id: remaining_uses} for consumables
jailed_until: str | None # ISO datetime or None
jailbreak_used: bool
reminders: list[str] # command names user wants DM reminders for
eco_banned: bool # if True, user cannot use any economy commands
# Lifetime statistics
peak_balance: int
lifetime_earned: int
lifetime_lost: int
work_count: int
beg_count: int
total_wagered: int
biggest_win: int
biggest_loss: int
slots_jackpots: int
crimes_attempted: int
crimes_succeeded: int
times_jailed: int
total_bail_paid: int
heists_joined: int
heists_won: int
total_given: int
total_received: int
best_daily_streak: int
heist_global_cd_until: float
# Prestige system
prestige_level: int
prestige_points: int
season_total_exp: int # cumulative EXP this season (survives prestige resets)
prestige_upgrades: dict # {upgrade_id: level}
# Fishing system
last_fish: str | None
fish_book: dict # {fish_id: times_caught}
total_fish_caught: int
fish_inventory: list # [{fish_id, weight, value}] - survives prestige
# Quest system
quest_daily: dict # {"date": "YYYY-MM-DD", "quests": {qid: {snap, claimed}}}
quest_weekly: dict # {"week": "YYYY-Www", "quests": {qid: {snap, claimed}}}
def _default_user() -> UserData:
return {
"balance": 0,
"exp": 0,
"last_daily": None,
"last_work": None,
"last_beg": None,
"last_crime": None,
"last_rob": None,
"last_heist": None,
"daily_streak": 0,
"last_streak_date": None,
"items": [],
"item_uses": {},
"jailed_until": None,
"jailbreak_used": False,
"reminders": ["daily", "work", "beg", "crime", "rob"],
"eco_banned": False,
# ── Lifetime stats ──────────────────────────────────────────────────
"peak_balance": 0,
"lifetime_earned": 0,
"lifetime_lost": 0,
"work_count": 0,
"beg_count": 0,
"total_wagered": 0,
"biggest_win": 0,
"biggest_loss": 0,
"slots_jackpots": 0,
"crimes_attempted": 0,
"crimes_succeeded": 0,
"times_jailed": 0,
"total_bail_paid": 0,
"heists_joined": 0,
"heists_won": 0,
"total_given": 0,
"total_received": 0,
"best_daily_streak": 0,
"heist_global_cd_until": 0.0,
# ── Prestige ─────────────────────────────────────────────────────────
"prestige_level": 0,
"prestige_points": 0,
"season_total_exp": 0,
"prestige_upgrades": {},
# ── Fishing ──────────────────────────────────────────────────────────
"last_fish": None,
"fish_book": {},
"total_fish_caught": 0,
"fish_inventory": [],
# ── Quests ───────────────────────────────────────────────────────────
"quest_daily": {},
"quest_weekly": {},
}
# ---------------------------------------------------------------------------
# Persistence (PocketBase backend)
# ---------------------------------------------------------------------------
_log = logging.getLogger("tipiCOIN.economy")
# ---------------------------------------------------------------------------
# Per-user write locks
# ---------------------------------------------------------------------------
# Every mutation is a read-modify-write cycle (get_user → mutate → _commit);
# without serialization, two concurrent commands for the same user overwrite
# each other's commit. Locking rules that keep this deadlock-free:
# - a decorated function must never call another decorated function
# - house balance changes go through _credit_house (an atomic PocketBase
# increment, no lock), so they are safe while holding user locks
_user_locks: dict[int, asyncio.Lock] = {}
def _user_lock(user_id: int) -> asyncio.Lock:
lock = _user_locks.get(user_id)
if lock is None:
lock = _user_locks[user_id] = asyncio.Lock()
return lock
def _locked_by(*arg_positions: int):
"""Serialize the decorated function per user id found at the given
positional-argument indices. Multiple ids are acquired in sorted order so
two-user functions (do_give, do_rob) cannot deadlock each other."""
def decorator(fn):
@functools.wraps(fn)
async def wrapper(*args, **kwargs):
locks = [_user_lock(uid) for uid in sorted({args[pos] for pos in arg_positions})]
for lock in locks:
await lock.acquire()
try:
return await fn(*args, **kwargs)
finally:
for lock in reversed(locks):
lock.release()
return wrapper
return decorator
# ---------------------------------------------------------------------------
# Public helpers
# ---------------------------------------------------------------------------
async def missing_schema_fields() -> list[str]:
"""Compare the live PocketBase collection schema against every field the
bot persists. PocketBase silently drops writes to undeclared fields, so
any name returned here means broken features without error messages."""
live = await pb_client.get_collection_fields()
expected = set(_default_user()) | {"user_id"}
return sorted(expected - live)
async def get_all_users_raw() -> dict[str, "UserData"]:
"""Return a snapshot of all user records."""
records = await pb_client.list_all_records()
result: dict[str, UserData] = {}
for record in records:
uid = record.get("user_id", "")
if not uid:
continue
user = _default_user()
for key in list(user.keys()):
if key in record:
user[key] = record[key] # type: ignore[literal-required]
user["_pb_id"] = record["id"] # type: ignore[typeddict-unknown-key]
result[uid] = user
return result
def _parse_dt(s: str | None) -> datetime | None:
if not s:
return None
dt = datetime.fromisoformat(s)
# Ensure timezone-aware
return dt if dt.tzinfo else dt.replace(tzinfo=timezone.utc)
def _cooldown_remaining(
user: UserData, action: str, override_cd: timedelta | None = None
) -> timedelta | None:
"""Return remaining cooldown, or None if the action is ready."""
last = _parse_dt(user.get(f"last_{action}"))
if last is None:
return None
cd = override_cd if override_cd is not None else COOLDOWNS[action]
remaining = cd - (_now() - last)
return remaining if remaining.total_seconds() > 0 else None
def _is_jailed(user: UserData) -> timedelta | None:
"""Return remaining jail time, or None if free."""
until = _parse_dt(user.get("jailed_until"))
if until is None:
return None
remaining = until - _now()
return remaining if remaining.total_seconds() > 0 else None
def jailed_remaining(user: UserData) -> timedelta | None:
"""Public wrapper - return remaining jail time, or None if free."""
return _is_jailed(user)
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def format_td(td: timedelta) -> str:
"""Human-readable timedelta: '1t 23m' / '45m 12s' / '8s'."""
total = int(td.total_seconds())
h, rem = divmod(total, 3600)
m, s = divmod(rem, 60)
if h:
return f"{h}t {m}m"
if m:
return f"{m}m {s}s"
return f"{s}s"
async def get_user(user_id: int) -> UserData:
"""Fetch user data from PocketBase, creating a default record if first seen."""
uid = str(user_id)
try:
record = await pb_client.get_record(uid)
if record is None:
default = _default_user()
default["user_id"] = uid # type: ignore[typeddict-unknown-key]
record = await pb_client.create_record(default)
except (aiohttp.ClientError, asyncio.TimeoutError, RuntimeError) as exc:
_log.error("PocketBase unreachable for user %s: %s", user_id, exc)
raise DatabaseError(f"Database unavailable: {exc}") from exc
user = _default_user()
for key in list(user.keys()):
if key in record:
user[key] = record[key] # type: ignore[literal-required]
user["_pb_id"] = record["id"] # type: ignore[typeddict-unknown-key]
return user
def _prestige_mult(user: UserData) -> tuple[float, float]:
"""Return (coin_mult, exp_mult) based on prestige upgrades. Both ≥1.0."""
upgrades: dict = user.get("prestige_upgrades") or {} # type: ignore[assignment]
coin_level = upgrades.get("coin_mult", 0)
exp_level = upgrades.get("exp_mult", 0)
return (
1.0 + coin_level * PRESTIGE_SHOP["coin_mult"]["effect"],
1.0 + exp_level * PRESTIGE_SHOP["exp_mult"]["effect"],
)
# ---------------------------------------------------------------------------
# Internal write helper
# ---------------------------------------------------------------------------
async def _commit(user_id: int, user: UserData) -> dict | None:
"""Persist the full user record. Returns the record as PocketBase stored it
(fields absent from the collection schema are silently dropped by PB)."""
record_id = user.get("_pb_id") # type: ignore[typeddict-item]
clean = {k: v for k, v in user.items() if k != "_pb_id"}
clean["user_id"] = str(user_id)
try:
if record_id:
return await pb_client.update_record(record_id, clean)
else:
_log.warning("_commit for user %s had no _pb_id; creating new record", user_id)
created = await pb_client.create_record(clean)
user["_pb_id"] = created["id"] # type: ignore[typeddict-unknown-key]
return created
except (aiohttp.ClientError, asyncio.TimeoutError, RuntimeError) as exc:
_log.error("_commit failed for user %s: %s", user_id, exc)
raise DatabaseError(f"Failed to persist user {user_id}: {exc}") from exc
# ---------------------------------------------------------------------------
# /reminders
# ---------------------------------------------------------------------------
@_locked_by(0)
async def do_set_reminders(user_id: int, commands: list[str]) -> None:
"""Overwrite the user's reminder list with the given command names."""
user = await get_user(user_id)
user["reminders"] = list(commands)
await _commit(user_id, user)