didier/database/crud/currency.py

138 lines
3.4 KiB
Python

from datetime import date
from typing import Union
from sqlalchemy.ext.asyncio import AsyncSession
from database.crud import users
from database.exceptions import currency as exceptions
from database.schemas import Bank, NightlyData
from database.utils.math.currency import (
capacity_upgrade_price,
interest_upgrade_price,
rob_upgrade_price,
)
__all__ = [
"add_dinks",
"claim_nightly",
"get_bank",
"get_nightly_data",
"invest",
"upgrade_capacity",
"upgrade_interest",
"upgrade_rob",
"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_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 invest(session: AsyncSession, user_id: int, amount: Union[str, int]) -> int:
"""Invest all your Dinks"""
bank = await get_bank(session, user_id)
if amount == "all":
amount = bank.dinks
# Don't allow investing more dinks than you own
amount = min(bank.dinks, int(amount))
bank.dinks -= amount
bank.invested += amount
session.add(bank)
await session.commit()
return amount
async def add_dinks(session: AsyncSession, user_id: int, amount: int):
"""Increase the Dinks counter for a user"""
bank = await get_bank(session, user_id)
bank.dinks += amount
session.add(bank)
await session.commit()
async def claim_nightly(session: AsyncSession, user_id: int):
"""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 = 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) -> int:
"""Upgrade capacity level"""
bank = 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) -> int:
"""Upgrade interest level"""
bank = 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) -> int:
"""Upgrade rob level"""
bank = 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