AutoShop-Wednesday/tgbot/database/db_refill.py

# - *- coding: utf- 8 - *-
from typing import Any, Dict, Optional, Union

from sqlalchemy import BigInteger, Float, Integer, String, text
from sqlalchemy.orm import Mapped, mapped_column

from tgbot.database.core import Base, session_factory
from tgbot.database.entities import Refill
from tgbot.database.repository import BaseRepository
from tgbot.utils.const_functions import get_unix


class RefillModel(Base):
    __tablename__ = "storage_refill"

    increment: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    user_id: Mapped[int] = mapped_column(BigInteger, nullable=False, index=True)
    refill_receipt: Mapped[int] = mapped_column(BigInteger, nullable=False, index=True)
    refill_comment: Mapped[str] = mapped_column(String(255), nullable=False, default="")
    refill_amount: Mapped[float] = mapped_column(Float, nullable=False, default=0)
    refill_method: Mapped[str] = mapped_column(String(64), nullable=False)
    refill_unix: Mapped[int] = mapped_column(Integer, nullable=False, default=get_unix)


ModelBase = Refill
BaseModel = Refill


class Refillx(BaseRepository[RefillModel, Refill]):
    # Подключение модели пополнений
    def __init__(self):
        super().__init__()
        self.table_model = RefillModel
        self.entity_model = Refill
        self.storage_name = RefillModel.__tablename__

    # Добавление записи пополнения
    async def add(
            self,
            user_id: int,
            refill_receipt: int,
            refill_comment: str,
            refill_amount: float,
            refill_method: str,
    ) -> Refill:
        return await self._insert(
            user_id=user_id,
            refill_comment=refill_comment,
            refill_amount=refill_amount,
            refill_receipt=refill_receipt,
            refill_method=refill_method,
            refill_unix=get_unix(),
        )

    # Обновление пополнения по чеку или фильтру
    async def update(self, where: Optional[Union[Dict[str, Any], int]] = None, **kwargs) -> int:
        if isinstance(where, int):
            where = {"refill_receipt": where}

        return await self._update(where=where, **kwargs)

    # Успешное зачисление пополнения
    @staticmethod
    async def success(
            user_id: int,
            pay_receipt: int,
            pay_comment: str,
            pay_amount: float,
            pay_method: str,
    ) -> str:
        async with session_factory() as session:
            try:
                await session.execute(text("BEGIN IMMEDIATE"))

                if pay_comment != "":
                    duplicate_refill = await session.execute(
                        text("SELECT increment FROM storage_refill WHERE refill_comment = :comment LIMIT 1"),
                        {"comment": pay_comment},
                    )
                else:
                    duplicate_refill = await session.execute(
                        text("SELECT increment FROM storage_refill WHERE refill_receipt = :receipt LIMIT 1"),
                        {"receipt": pay_receipt},
                    )

                if duplicate_refill.mappings().first() is not None:
                    await session.rollback()
                    return "ALREADY"

                user_result = await session.execute(
                    text(
                        """
                        SELECT user_balance, user_refill, user_referrer_id
                        FROM storage_users
                        WHERE user_id = :user_id
                        """
                    ),
                    {"user_id": user_id},
                )
                get_user = user_result.mappings().first()

                if get_user is None:
                    await session.rollback()
                    return "USER_NOT_FOUND"

                settings_result = await session.execute(
                    text(
                        "SELECT status_referral, referral_refill_percent "
                        "FROM storage_settings WHERE id = 1"
                    )
                )
                get_settings = settings_result.mappings().first()

                new_balance = round(float(get_user["user_balance"]) + float(pay_amount), 2)
                new_refill = round(float(get_user["user_refill"]) + float(pay_amount), 2)

                await session.execute(
                    text(
                        """
                        INSERT INTO storage_refill (user_id,
                                                    refill_receipt,
                                                    refill_comment,
                                                    refill_amount,
                                                    refill_method,
                                                    refill_unix)
                        VALUES (:user_id, :receipt, :comment, :amount, :method, :unix)
                        """
                    ),
                    {
                        "user_id": user_id,
                        "receipt": pay_receipt,
                        "comment": pay_comment,
                        "amount": pay_amount,
                        "method": pay_method,
                        "unix": get_unix(),
                    },
                )
                await session.execute(
                    text(
                        """
                        UPDATE storage_users
                        SET user_balance = :balance,
                            user_refill  = :refill
                        WHERE user_id = :user_id
                        """
                    ),
                    {"balance": new_balance, "refill": new_refill, "user_id": user_id},
                )

                referrer_id = get_user["user_referrer_id"]
                referral_reward = 0.0

                if (
                    get_settings is not None
                    and get_settings["status_referral"] == "True"
                    and referrer_id is not None
                    and referrer_id != user_id
                ):
                    referral_reward = round(
                        float(pay_amount)
                        * float(get_settings["referral_refill_percent"])
                        / 100,
                        2,
                    )

                if referral_reward > 0:
                    referrer_result = await session.execute(
                        text(
                            "SELECT user_referral_balance FROM storage_users "
                            "WHERE user_id = :user_id"
                        ),
                        {"user_id": referrer_id},
                    )
                    get_referrer = referrer_result.mappings().first()

                    if get_referrer is not None:
                        source_id = f"refill:{pay_comment or pay_receipt}"

                        await session.execute(
                            text(
                                "UPDATE storage_users "
                                "SET user_referral_balance = :balance "
                                "WHERE user_id = :user_id"
                            ),
                            {
                                "balance": round(
                                    float(get_referrer["user_referral_balance"])
                                    + referral_reward,
                                    2,
                                ),
                                "user_id": referrer_id,
                            },
                        )
                        await session.execute(
                            text(
                                "INSERT INTO storage_referral_transactions "
                                "(user_id, related_user_id, transaction_type, source_id, amount, transaction_unix) "
                                "VALUES (:user_id, :related_user_id, :transaction_type, :source_id, :amount, :transaction_unix)"
                            ),
                            {
                                "user_id": referrer_id,
                                "related_user_id": user_id,
                                "transaction_type": "refill_percent",
                                "source_id": source_id,
                                "amount": referral_reward,
                                "transaction_unix": get_unix(),
                            },
                        )
                await session.commit()
            except Exception:
                await session.rollback()
                raise

        return "ok"