from datetime import date from typing import List, Optional, Union from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from database.crud import users from database.exceptions import currency as exceptions from database.schemas import Bank, BankSavings, NightlyData from database.utils.math.currency import ( capacity_upgrade_price, interest_rate, interest_upgrade_price, rob_upgrade_price, savings_cap, ) __all__ = [ "add_dinks", "apply_daily_interest", "claim_nightly", "deduct_dinks", "gamble_dinks", "get_bank", "get_nightly_data", "save", "upgrade_capacity", "upgrade_interest", "upgrade_rob", "withdraw", "NIGHTLY_AMOUNT", ] NIGHTLY_AMOUNT = 420 async def get_bank(session: AsyncSession, user_id: int) -> Bank: """Get a user's bank info""" user = await users.get_or_add_user(session, user_id) return user.bank async def get_savings(session: AsyncSession, user_id: int) -> BankSavings: """Get a user's savings info""" user = await users.get_or_add_user(session, user_id) return user.savings async def get_nightly_data(session: AsyncSession, user_id: int) -> NightlyData: """Get a user's nightly info""" user = await users.get_or_add_user(session, user_id) return user.nightly_data async def save( session: AsyncSession, user_id: int, amount: Union[str, int], *, bank: Optional[Bank] = None, savings: Optional[BankSavings] = None ) -> int: """Invest some of your Dinks""" bank = bank or await get_bank(session, user_id) savings = savings or await get_savings(session, user_id) if amount == "all": amount = bank.dinks if bank.dinks <= 0: return 0 # Don't allow investing more dinks than you own amount = min(bank.dinks, int(amount)) # Don't allow exceeding the limit limit = savings_cap(bank.capacity_level) if savings.saved >= limit: raise exceptions.SavingsCapExceeded if savings.saved + amount > limit: amount = max(0, limit - savings.saved) bank.dinks -= amount savings.saved += amount session.add(bank) session.add(savings) await session.commit() return amount async def withdraw(session: AsyncSession, user_id: int, amount: Union[str, int], *, bank: Optional[Bank] = None) -> int: """Withdraw your saved Dinks""" bank = bank or await get_bank(session, user_id) savings = await get_savings(session, user_id) if amount == "all": amount = savings.saved # Don't allow withdrawing more dinks than you own amount = min(savings.saved, int(amount)) bank.dinks += amount savings.saved -= amount savings.daily_minimum = min(savings.daily_minimum, savings.saved) session.add(bank) session.add(savings) await session.commit() return amount async def add_dinks(session: AsyncSession, user_id: int, amount: int, *, bank: Optional[Bank] = None): """Increase the Dinks counter for a user""" bank = bank or await get_bank(session, user_id) bank.dinks += amount session.add(bank) await session.commit() async def deduct_dinks(session: AsyncSession, user_id: int, amount: int, *, bank: Optional[Bank] = None) -> int: """Decrease the Dinks counter for a user""" bank = bank or await get_bank(session, user_id) deducted_amount = min(amount, bank.dinks) bank.dinks -= deducted_amount session.add(bank) await session.commit() return deducted_amount async def claim_nightly(session: AsyncSession, user_id: int, *, bank: Optional[Bank] = None): """Claim daily Dinks""" nightly_data = await get_nightly_data(session, user_id) now = date.today() if nightly_data.last_nightly is not None and nightly_data.last_nightly == now: raise exceptions.DoubleNightly bank = bank or await get_bank(session, user_id) bank.dinks += NIGHTLY_AMOUNT nightly_data.last_nightly = now session.add(bank) session.add(nightly_data) await session.commit() async def upgrade_capacity(session: AsyncSession, user_id: int, *, bank: Optional[Bank] = None) -> int: """Upgrade capacity level""" bank = bank or await get_bank(session, user_id) upgrade_price = capacity_upgrade_price(bank.capacity_level) # Can't afford this upgrade if upgrade_price > bank.dinks: raise exceptions.NotEnoughDinks bank.dinks -= upgrade_price bank.capacity_level += 1 session.add(bank) await session.commit() return bank.capacity_level async def upgrade_interest(session: AsyncSession, user_id: int, *, bank: Optional[Bank] = None) -> int: """Upgrade interest level""" bank = bank or await get_bank(session, user_id) upgrade_price = interest_upgrade_price(bank.interest_level) # Can't afford this upgrade if upgrade_price > bank.dinks: raise exceptions.NotEnoughDinks bank.dinks -= upgrade_price bank.interest_level += 1 session.add(bank) await session.commit() return bank.interest_level async def upgrade_rob(session: AsyncSession, user_id: int, *, bank: Optional[Bank] = None) -> int: """Upgrade rob level""" bank = bank or await get_bank(session, user_id) upgrade_price = rob_upgrade_price(bank.rob_level) # Can't afford this upgrade if upgrade_price > bank.dinks: raise exceptions.NotEnoughDinks bank.dinks -= upgrade_price bank.rob_level += 1 session.add(bank) await session.commit() return bank.rob_level async def gamble_dinks( session: AsyncSession, user_id: int, amount: Union[str, int], payout_factor: int, won: bool, *, bank: Optional[Bank] = None ) -> int: """Gamble some of your Didier Dinks""" bank = bank or await get_bank(session, user_id) if amount == "all": amount = bank.dinks if bank.dinks <= 0: return 0 amount = min(int(amount), bank.dinks) sign = 1 if won else -1 factor = (payout_factor - 1) if won else 1 await add_dinks(session, user_id, sign * amount * factor, bank=bank) return amount * factor async def rob( session: AsyncSession, amount: int, robber_id: int, robbed_id: int, *, robber_bank: Optional[Bank] = None, robbed_bank: Optional[Bank] = None ): """Rob another user's Didier Dinks""" robber = robber_bank or await get_bank(session, robber_id) robbed = robbed_bank or await get_bank(session, robbed_id) robber.dinks += amount robbed.dinks -= amount session.add(robber) session.add(robbed) await session.commit() async def apply_daily_interest(session: AsyncSession): """Apply daily interest rates to all accounts with saved Dinks""" statement = select(BankSavings) all_savings: List[BankSavings] = list((await session.execute(statement)).scalars().all()) for savings_account in all_savings: if savings_account.saved == 0: continue bank = await get_bank(session, savings_account.user_id) rate = interest_rate(bank.interest_level) savings_account.saved = float(savings_account.saved * rate) savings_account.daily_minimum = savings_account.saved session.add(savings_account) await session.commit()