from sqlalchemy import func, select, update from bot.config import ADMIN_ID from bot.db.models import AppSetting, BotAccount, User from bot.db.session import async_session BOT_SLOT_PRICE_KEY = "bot_slot_price" DEFAULT_BOT_SLOT_PRICE = 1 class BalanceError(ValueError): pass class InsufficientBalance(BalanceError): pass def _amount(value: int) -> int: if isinstance(value, bool) or not isinstance(value, int) or value <= 0: raise BalanceError("Amount must be a positive integer") return value async def _get_or_create_user(session, telegram_id: int) -> User: user = await session.scalar(select(User).where(User.telegram_id == telegram_id)) if user is None: user = User(telegram_id=telegram_id, balance_credits=0) session.add(user) await session.flush() return user async def get_user_balance(telegram_id: int) -> int: async with async_session() as session: user = await session.scalar(select(User).where(User.telegram_id == telegram_id)) return int(user.balance_credits) if user else 0 async def change_user_balance( telegram_id: int, amount: int, *, actor_telegram_id: int | None = None ) -> int: amount = _amount(amount) if ( actor_telegram_id is not None and telegram_id != actor_telegram_id and actor_telegram_id != ADMIN_ID ): raise PermissionError("Only an administrator can adjust another user") async with async_session() as session: user = await _get_or_create_user(session, telegram_id) await session.execute( update(User).where(User.id == user.id).values(balance_credits=amount) ) await session.commit() return int( ( await session.scalar( select(User.balance_credits).where(User.id == user.id) ) ) or 0 ) async def credit_user( telegram_id: int, amount: int, *, actor_telegram_id: int | None = None ) -> int: amount = _amount(amount) if ( actor_telegram_id is not None and telegram_id != actor_telegram_id and actor_telegram_id != ADMIN_ID ): raise PermissionError("Only an administrator can adjust another user") async with async_session() as session: user = await _get_or_create_user(session, telegram_id) await session.execute( update(User) .where(User.id == user.id) .values(balance_credits=User.balance_credits + amount) ) await session.commit() return int( ( await session.scalar( select(User.balance_credits).where(User.id == user.id) ) ) or 0 ) async def debit_user( telegram_id: int, amount: int, *, actor_telegram_id: int | None = None ) -> int: amount = _amount(amount) if ( actor_telegram_id is not None and telegram_id != actor_telegram_id and actor_telegram_id != ADMIN_ID ): raise PermissionError("Only an administrator can adjust another user") async with async_session() as session: user = await _get_or_create_user(session, telegram_id) result = await session.execute( update(User) .where(User.id == user.id, User.balance_credits >= amount) .values(balance_credits=User.balance_credits - amount) ) if result.rowcount != 1: await session.rollback() raise InsufficientBalance("Insufficient balance") await session.commit() return int( ( await session.scalar( select(User.balance_credits).where(User.id == user.id) ) ) or 0 ) async def get_bot_slot_price() -> int: async with async_session() as session: setting = await session.get(AppSetting, BOT_SLOT_PRICE_KEY) if setting is None: setting = AppSetting( key=BOT_SLOT_PRICE_KEY, integer_value=DEFAULT_BOT_SLOT_PRICE ) session.add(setting) await session.commit() return DEFAULT_BOT_SLOT_PRICE if setting.integer_value <= 0: return DEFAULT_BOT_SLOT_PRICE return int(setting.integer_value) async def set_bot_slot_price(price: int, *, updated_by: int) -> int: price = _amount(price) async with async_session() as session: setting = await session.get(AppSetting, BOT_SLOT_PRICE_KEY) if setting is None: setting = AppSetting( key=BOT_SLOT_PRICE_KEY, integer_value=price, updated_by=updated_by ) session.add(setting) else: setting.integer_value = price setting.updated_by = updated_by await session.commit() return price async def get_user_profile(telegram_id: int) -> tuple[int, int, int]: async with async_session() as session: user = await session.scalar(select(User).where(User.telegram_id == telegram_id)) if user is None: return telegram_id, 0, 0 count = await session.scalar( select(func.count(BotAccount.id)).where(BotAccount.owner_id == user.id) ) return telegram_id, int(user.balance_credits), int(count or 0)