ASF-Control-Bot/bot/services/balance.py

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)