Refator eco code. Split into modules
This commit is contained in:
384
core/economy/store.py
Normal file
384
core/economy/store.py
Normal file
@@ -0,0 +1,384 @@
|
||||
"""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)
|
||||
Reference in New Issue
Block a user