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 User as TgUser
from telebot.types import Chat as TgChat
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.members)
    contributors = len(list(filter(lambda m: m.contributed, fund.members)))
    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,
    )


def get_user(tg_user: TgUser) -> User:
    if user := session.query(User).filter(User.id == tg_user.id).first():
        return user
    user = User(
        id=tg_user.id,
        name=tg_user.first_name,
        username=tg_user.username,
    )
    session.add(user)
    session.commit()
    return user


def get_group(tg_chat: TgChat) -> Group:
    if group := session.query(User).filter(User.id == tg_chat.id).first():
        return group
    group = Group(
        id=tg_chat.id,
    )
    session.add(group)
    session.commit()
    return group


# Bot logic

@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()
        user = get_user(msg.from_user)
        await bot.send_message(chat_id=msg.chat.id, text=textbook.start_private.format(count=len(user.fund_members)))

    elif msg.chat.type in ("group", "supergroup"):
        get_group(msg.chat)
        await bot.send_message(chat_id=msg.chat.id, text=textbook.start_group)


@bot.message_handler(commands=["setup"])
async def setup(msg: Message):
    get_group(msg.chat)
    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
    user = get_user(call.from_user.id)
    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(Group).filter(Group.id == msg.chat.id).first():
            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",
                )
        else:
            await bot.send_message(
                chat_id=msg.chat.id,
                text=textbook.not_initialized
            )


@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(content_types=["text"], state=States.newfund_amount)
async def newfund_amount(msg: Message):
    if not session.query(User).filter(User.id == msg.from_user.id).first():
        user = User(
            id=msg.from_user.id,
            name=msg.from_user.first_name,
            username=msg.from_user.username,
        )
        session.add(user)
    session.commit()
    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=["fund"])
async def fund(msg: Message):


@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())