tipibot/core/economy/admin.py

"""Admin mutations and the season reset."""

from __future__ import annotations

from datetime import timedelta

from .. import pb_client
from .store import _commit, _default_user, _locked_by, _now, _txn, get_user
from .levels import get_level
from .shop import SHOP


async def do_season_reset(top_n: int = 10) -> list[tuple[str, int, int]]:
    """Snapshot top_n by EXP, then full wipe: EXP, balance, items, item_uses.
    Returns top list (uid, exp, level) captured before the reset."""
    records = await pb_client.list_all_records()
    top = sorted(
        ((r["user_id"], r.get("exp", 0)) for r in records if r.get("user_id")),
        key=lambda x: x[1],
        reverse=True,
    )[:top_n]
    reset_fields = {
        "exp": 0,
        "balance": 0,
        "items": [],
        "item_uses": {},
        "last_daily": None,
        "last_work": None,
        "last_beg": None,
        "last_crime": None,
        "last_rob": None,
        "last_fish": None,
        "daily_streak": 0,
        "last_streak_date": None,
        "season_total_exp": 0,
    }
    for record in records:
        await pb_client.update_record(record["id"], reset_fields)
    return [(uid, exp, get_level(exp)) for uid, exp in top]


# ---------------------------------------------------------------------------
# Admin actions
# ---------------------------------------------------------------------------
@_locked_by(0)
async def do_admin_coins(target_id: int, amount: int, admin_id: int, reason: str) -> dict:
    """Give (positive) or take (negative) coins from a user. Balance is floored at 0."""
    user = await get_user(target_id)
    user["balance"] = max(0, user["balance"] + amount)
    await _commit(target_id, user)
    verb = f"+{amount}" if amount >= 0 else str(amount)
    _txn("ADMIN_COINS", admin=admin_id, target=target_id, amount=verb, reason=reason, bal=user["balance"])
    return {"ok": True, "balance": user["balance"], "change": amount}


@_locked_by(0)
async def do_admin_jail(target_id: int, minutes: int, admin_id: int, reason: str) -> dict:
    """Manually jail a user for `minutes` minutes."""
    user = await get_user(target_id)
    user["jailed_until"] = (_now() + timedelta(minutes=minutes)).isoformat()
    user["jailbreak_used"] = False
    await _commit(target_id, user)
    _txn("ADMIN_JAIL", admin=admin_id, target=target_id, minutes=minutes, reason=reason)
    return {"ok": True, "jailed_until": user["jailed_until"]}


@_locked_by(0)
async def do_admin_unjail(target_id: int, admin_id: int) -> dict:
    """Remove jail from a user."""
    user = await get_user(target_id)
    user["jailed_until"] = None
    user["jailbreak_used"] = False
    await _commit(target_id, user)
    _txn("ADMIN_UNJAIL", admin=admin_id, target=target_id)
    return {"ok": True}


@_locked_by(0)
async def do_admin_ban(target_id: int, admin_id: int, reason: str) -> dict:
    """Ban a user from all economy commands."""
    user = await get_user(target_id)
    user["eco_banned"] = True
    await _commit(target_id, user)
    _txn("ADMIN_BAN", admin=admin_id, target=target_id, reason=reason)
    return {"ok": True}


@_locked_by(0)
async def do_admin_unban(target_id: int, admin_id: int) -> dict:
    """Lift an economy ban."""
    user = await get_user(target_id)
    user["eco_banned"] = False
    await _commit(target_id, user)
    _txn("ADMIN_UNBAN", admin=admin_id, target=target_id)
    return {"ok": True}


@_locked_by(0)
async def do_admin_reset(target_id: int, admin_id: int) -> dict:
    """Wipe a user's economy data back to defaults."""
    user = await get_user(target_id)
    fresh = _default_user()
    fresh["_pb_id"] = user.get("_pb_id")  # type: ignore[typeddict-unknown-key]
    await _commit(target_id, fresh)
    _txn("ADMIN_RESET", admin=admin_id, target=target_id)
    return {"ok": True}


async def do_admin_inspect(target_id: int) -> dict:
    """Return the user's full raw economy data."""
    user = await get_user(target_id)
    return {"ok": True, "data": dict(user)}


@_locked_by(0)
async def do_admin_exp(target_id: int, amount: int, admin_id: int, reason: str) -> dict:
    """Give (positive) or take (negative) EXP from a user. EXP is floored at 0."""
    user = await get_user(target_id)
    old_exp = user.get("exp", 0)
    old_level = get_level(old_exp)
    user["exp"] = max(0, old_exp + amount)
    user["season_total_exp"] = max(0, user.get("season_total_exp", 0) + amount)
    new_level = get_level(user["exp"])
    await _commit(target_id, user)
    verb = f"+{amount}" if amount >= 0 else str(amount)
    _txn("ADMIN_EXP", admin=admin_id, target=target_id, amount=verb, reason=reason, exp=user["exp"])
    return {
        "ok": True,
        "exp": user["exp"],
        "change": amount,
        "old_level": old_level,
        "new_level": new_level,
        "level_changed": new_level != old_level,
    }


@_locked_by(0)
async def do_admin_item(target_id: int, item_id: str, action: str, admin_id: int) -> dict:
    """Give or remove an item. action='give'|'remove'. Returns ok/reason."""
    if item_id not in SHOP:
        return {"ok": False, "reason": "invalid_item"}
    user = await get_user(target_id)
    items: list = list(user.get("items") or [])
    item_uses: dict = dict(user.get("item_uses") or {})
    if action == "give":
        if item_id not in items:
            items.append(item_id)
        if item_id == "anticheat":
            item_uses["anticheat"] = 2
        user["items"] = items
        user["item_uses"] = item_uses
        await _commit(target_id, user)
        _txn("ADMIN_ITEM_GIVE", admin=admin_id, target=target_id, item=item_id)
        return {"ok": True, "action": "given", "item_id": item_id}
    elif action == "remove":
        if item_id not in items:
            return {"ok": False, "reason": "not_owned"}
        items.remove(item_id)
        item_uses.pop(item_id, None)
        user["items"] = items
        user["item_uses"] = item_uses
        await _commit(target_id, user)
        _txn("ADMIN_ITEM_REMOVE", admin=admin_id, target=target_id, item=item_id)
        return {"ok": True, "action": "removed", "item_id": item_id}
    return {"ok": False, "reason": "invalid_action"}