Javelina/classes/wallet.py

import logging
from dataclasses import dataclass
from datetime import datetime, timezone
from logging import Logger
from typing import Any, Dict, Optional

from bson import ObjectId
from pymongo.results import InsertOneResult

from classes.errors.wallet import (
    WalletBalanceLimitExceeded,
    WalletNotFoundError,
    WalletOverdraftLimitExceeded,
    WalletInsufficientFunds,
)
from modules.database import col_wallets

logger: Logger = logging.getLogger(__name__)


@dataclass
class Wallet:
    _id: ObjectId
    owner_id: int
    guild_id: int
    balance: float
    is_frozen: bool
    created: datetime

    # TODO Write a docstring
    @classmethod
    async def from_id(
        cls, owner_id: int, guild_id: int, allow_creation: bool = True
    ) -> "Wallet":
        db_entry = await col_wallets.find_one({"owner_id": owner_id, "guild_id": guild_id})

        if db_entry is None:
            if not allow_creation:
                raise WalletNotFoundError(owner_id, guild_id)

            db_entry = Wallet.get_defaults(owner_id, guild_id)

            insert_result: InsertOneResult = await col_wallets.insert_one(db_entry)

            db_entry["_id"] = insert_result.inserted_id

        return cls(**db_entry)

    def to_dict(self, json_compatible: bool = False) -> Dict[str, Any]:
        """Convert Wallet object to a JSON representation.

        Args:
            json_compatible (bool): Whether the JSON-incompatible objects like ObjectId need to be converted

        Returns:
            Dict[str, Any]: JSON representation of Wallet
        """
        return {
            "_id": self._id if not json_compatible else str(self._id),
            "owner_id": self.owner_id,
            "guild_id": self.guild_id,
            "balance": self.balance,
            "is_frozen": self.is_frozen,
            "created": self.created,
        }

    async def _set(self, key: str, value: Any) -> None:
        if not hasattr(self, key):
            raise AttributeError()

        setattr(self, key, value)

        await col_wallets.update_one({"_id": self._id}, {"$set": {key: value}}, upsert=True)

        logger.info("Set attribute '%s' of the wallet %s to '%s'", key, str(self._id), value)

    @staticmethod
    def get_defaults(
        owner_id: Optional[int] = None, guild_id: Optional[int] = None
    ) -> Dict[str, Any]:
        return {
            "owner_id": owner_id,
            "guild_id": guild_id,
            "balance": 0.0,
            "is_frozen": False,
            "created": datetime.now(tz=timezone.utc),
        }

    @staticmethod
    def get_default_value(key: str) -> Any:
        if key not in Wallet.get_defaults():
            raise KeyError(f"There's no default value for key '{key}' in Wallet")

        return Wallet.get_defaults()[key]

    # TODO Write a docstring
    async def freeze(self) -> None:
        await self._set("is_frozen", True)

    # TODO Write a docstring
    async def unfreeze(self) -> None:
        await self._set("is_frozen", False)

    # TODO Write a docstring
    async def deposit(self, amount: float, balance_limit: Optional[float] = None) -> float:
        new_balance: float = round(self.balance + amount, 2)

        if balance_limit is not None and new_balance > balance_limit:
            raise WalletBalanceLimitExceeded(self, amount, balance_limit)

        await self._set("balance", new_balance)

        return new_balance

    # TODO Write a docstring
    async def withdraw(
        self,
        amount: float,
        allow_overdraft: bool = False,
        overdraft_limit: Optional[float] = None,
    ) -> float:
        if amount > self.balance:
            if not allow_overdraft or overdraft_limit is None:
                raise WalletInsufficientFunds(self, amount)

            if allow_overdraft and amount > overdraft_limit:
                raise WalletOverdraftLimitExceeded(self, amount, overdraft_limit)

        new_balance: float = round(self.balance - amount, 2)

        await self._set("balance", new_balance)

        return new_balance

    async def transfer(
        self,
        wallet_owner_id: int,
        wallet_guild_id: int,
        amount: float,
        balance_limit: Optional[float] = None,
        allow_overdraft: bool = False,
        overdraft_limit: Optional[float] = None,
    ) -> None:
        # TODO Replace with a concrete exception
        if amount < 0:
            raise ValueError()

        wallet: Wallet = await self.from_id(
            wallet_owner_id, wallet_guild_id, allow_creation=False
        )

        if balance_limit is not None and amount + wallet.balance > balance_limit:
            raise WalletBalanceLimitExceeded(wallet, amount, balance_limit)

        if amount > self.balance:
            if not allow_overdraft or overdraft_limit is None:
                raise WalletInsufficientFunds(self, amount)

            if allow_overdraft and amount > overdraft_limit:
                raise WalletOverdraftLimitExceeded(self, amount, overdraft_limit)

        # TODO Make a sanity check to revert the transaction if anything goes wrong
        await self.withdraw(amount, allow_overdraft, overdraft_limit)
        await wallet.deposit(amount, balance_limit)