waterfundbot/bot/app/bot.py

# Telebot imports
from telebot.async_telebot import AsyncTeleBot
from telebot.asyncio_storage import StateMemoryStorage, StatePickleStorage
from telebot.asyncio_handler_backends import State, StatesGroup
from telebot.asyncio_filters import StateFilter
from telebot import types
from telebot.callback_data import CallbackData, CallbackDataFilter
from telebot.types import Message
from telebot.util import user_link

# Async things imports
import asyncio

# Other modules imports
from datetime import datetime
import math
from typing import Union

# Local imports
from config import token
import textbook
import keyboards

# DB
from db.base import Session, engine, Base
from db.models import User, Group, GroupMember, Fund, FundMember


bot = AsyncTeleBot(token, state_storage=StatePickleStorage())
session: Session = None


class States(StatesGroup):
    default = State()
    newfund_amount = State()


# Utils


def get_fund_text(fund: Fund):
    count = len(fund.fund_members)
    contributors = len(fund.fund_members.filter(FundMember.contributed == True))
    personal_amount = math.ceil(fund.amount / count)
    return (
        "🟢 {name}\n\n",
        "💵 Сумма: {amount}р\n\n",
        "<b>Каждый скидывает по {personal_amount}</b>\n",
        "Уже собрано: {collected_amount}\n\n",
        "👥 Скинули: {contributors}/{count} чел.",
    ).format(
        name=fund.name,
        amount=fund.amount,
        personal_amount=personal_amount,
        collected_amount=personal_amount * contributors,
        contributors=contributors,
        count=count,
    )


@bot.message_handler(commands=["start"])
async def start(msg: Message):
    if msg.chat.type == "private":
        if user := session.query(User).filter(User.id == msg.chat.id).first():
            await bot.send_message(
                chat_id=msg.chat.id,
                text=textbook.private_info.format(count=len(user.fund_members)),
            )
        else:
            user = User(
                id=msg.chat.id,
                name=msg.from_user.first_name,
                username=msg.from_user.username,
            )
            session.add(user)
            session.commit()
            await bot.send_message(chat_id=msg.chat.id, text=textbook.start_private)

    elif msg.chat.type in ("group", "supergroup"):
        if not session.query(Group).filter(Group.id == msg.chat.id).first():
            group = Group(id=msg.chat.id)
            session.add(group)
            session.commit()
            await bot.send_message(chat_id=msg.chat.id, text=textbook.start_group)


@bot.message_handler(commands=["setup"])
async def setup(msg: Message):
    if msg.chat.type in ("group", "supergroup"):
        await bot.send_message(
            chat_id=msg.chat.id, text=textbook.setup, reply_markup=keyboards.setup()
        )


@bot.callback_query_handler(func=lambda c: c.data == "register_group_member")
async def register_group_member(call: types.CallbackQuery):
    new = False
    if not session.query(User).filter(User.id == call.from_user.id).first():
        user = User(
            id=call.from_user.id,
            name=call.from_user.first_name,
            username=call.from_user.username,
        )
        session.add(user)
    if (
        group_member := session.query(GroupMember)
        .filter(
            GroupMember.group_id == call.message.chat.id,
            GroupMember.user_id == call.from_user.id,
        )
        .first()
    ):
        session.delete(group_member)
    else:
        group_member = GroupMember(
            group_id=call.message.chat.id,
            user_id=call.from_user.id,
        )
        session.add(group_member)
        new = True
    session.commit()
    await bot.answer_callback_query(
        callback_query_id=call.id,
        text=textbook.user_parted if new else textbook.user_left,
        show_alert=True,
    )


@bot.message_handler(commands=["newfund"])
async def newfund(msg: Message):
    if msg.chat.type in ("group", "supergroup"):
        if (
            session.query(Fund)
            .filter(Fund.group_id == msg.chat.id, Fund.active == True)
            .first()
        ):
            await bot.send_message(
                chat_id=msg.chat.id,
                text=textbook.newfund_already_exists,
            )
        else:
            await bot.set_state(
                user_id=msg.from_user.id,
                chat_id=msg.chat.id,
                state=States.newfund_amount,
            )
            await bot.send_message(
                chat_id=msg.chat.id,
                text=textbook.newfund_amount.format(user=user_link(msg.from_user)),
                reply_markup=keyboards.cancel(),
                parse_mode="html",
            )


@bot.callback_query_handler(
    func=lambda c: c.data == "cancel", state=States.newfund_amount
)
async def cancel_newfund_amount(call: types.CallbackQuery):
    await bot.set_state(
        user_id=call.from_user.id, chat_id=call.message.chat.id, state=States.default
    )
    await bot.answer_callback_query(
        callback_query_id=call.id,
        text=textbook.cancel,
    )


@bot.message_handler(func=lambda message: True)
async def newfund_amount(msg: Message):
    await bot.send_message(chat_id=msg.chat.id, text="test")
    if msg.text.isdigit():
        await bot.set_state(
            user_id=msg.from_user.id, chat_id=msg.chat.id, state=States.default
        )
        fund = Fund(
            owner_id=msg.from_user.id, group_id=msg.chat.id, amount=int(msg.text)
        )
        session.add(fund)
        for group_member in session.query(GroupMember).filter(
            GroupMember.group_id == msg.chat.id
        ):
            fund_member = FundMember(user_id=group_member.user.id, fund_id=fund.id)
            session.add(fund_member)
        session.commit()
        await bot.send_message(
            chat_id=msg.chat.id,
            text=textbook.fund_created.format(fund=fund.name),
        )
        await asyncio.sleep(1)
        await bot.send_message(
            chat_id=msg.chat.id,
            text=get_fund_text(fund),
            reply_markup=keyboards.fund_markup(),
        )

    else:
        await bot.send_message(
            chat_id=msg.chat.id,
            text=textbook.not_number.format(user=user_link(msg.from_user)),
            reply_markup=keyboards.cancel(),
            parse_mode="html",
        )


@bot.message_handler(commands=["mystate"])
async def mystate(msg: Message):
    state = await bot.get_state(user_id=msg.from_user.id)
    await bot.send_message(chat_id=msg.chat.id, text=state)


@bot.message_handler(commands=["chatid"])
async def chatid(msg: Message):
    await bot.send_message(chat_id=msg.chat.id, text=msg.chat.id)


async def main():
    a = asyncio.create_task(bot.polling(non_stop=True))
    await a


if __name__ == "__main__":
    print("Bot started", flush=True)
    Base.metadata.create_all(engine)
    session = Session()
    bot.add_custom_filter(StateFilter(bot))
    bot.enable_saving_states(filename="./.state-save/states.pkl")
    asyncio.run(main())