queuebot/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

# Async things imports
import asyncio

# Other modules imports
import sqlite3
import json
from datetime import datetime
import math
import socket
import os

# Local imports
from config import token, admins
import textbook
import keyboards

# DB
from db.base import Session, engine, Base
from db.models import User, Queue, QueueUser


bot = AsyncTeleBot(token, state_storage=StatePickleStorage())


class States(StatesGroup):
    default = State()


def get_queue_stats_text(queue: Queue) -> str:
    s = f"Название: {queue.name}" f"Количество участников: {len(queue.users)}"
    return s


@bot.message_handler(commands=["start"])
async def start(msg: Message):
    if msg.chat.type == "private":
        user = session.query(User).filter_by(id=msg.from_user.id).first()
        if not user:
            new_user = User(
                id=msg.from_user.id,
                name=msg.from_user.first_name,
                username=msg.from_user.username,
            )
            session.add(new_user)
            session.commit()
            await bot.send_message(chat_id=msg.chat.id, text=textbook.start)
        await bot.set_state(user_id=msg.from_user.id, state=States.default)
        await bot.send_message(
            chat_id=msg.chat.id, text=textbook.menu, reply_markup=keyboards.menu()
        )

@bot.callback_query_handler(func=lambda c: c.data == "to_menu")
async def to_menu_handler(call: types.CallbackQuery):
    await bot.set_state(user_id=call.from_user.id, state=States.default)
    await bot.send_message(
        chat_id=call.message.chat.id, text=textbook.menu, reply_markup=keyboards.menu()
    )

@bot.callback_query_handler(func=lambda c: c.data == "new_queue")
async def new_queue_handler(call: types.CallbackQuery):
    user = session.query(User).filter_by(id=call.from_user.id).first()
    if user:
        if len(user.owns_queues) < 4:
            queue = Queue(owner_id=call.from_user.id)
            sessio.add(queue)
            session.commit()
            await bot.answer_callback_query(
                callback_query_id=call.id,
                text=textbook.new_queue_created.format(id=queue.id),
                show_alert=True,
            )
        else:
            await bot.answer_callback_query(
                callback_query_id=call.id, text=textbook.queue_limit, show_alert=True
            )
    await bot.answer_callback_query(callback_query_id=call.id)


@bot.callback_query_handler(func=lambda c: c.data == "my_queues")
async def my_queues_handler(call: types.CallbackQuery):
    user = session.query(User).filter_by(id=call.from_user.id).first()
    queues = user.owns_queues
    await bot.edit_message_text(
        chat_id=call.message.chat.id,
        message_id=call.message.id,
        text=textbook.my_queues_list.format(count=len(queues)),
        reply_markup=keyboards.my_queues(queues),
    )
    await bot.answer_callback_query(callback_query_id=call.id)


@bot.callback_query_handler(func=lambda c: c.data[:2] == "q:")
async def edit_queue_handler(call: types.CallbackQuery):
    queue_id = call.data[:2]
    queue = session.query(Queue).filter_by(id=queue_id).first()
    if not queue:
        await bot.answer_callback_query(
                callback_query_id=call.id, text=textbook.error
            )
        return None
    async with bot.retrieve_data(
        user_id=call.from_user.id, chat_id=call.message.chat.id
    ) as state_data:
        state_data["queue"] = call.data[2:]
    await bot.edit_message_text(
        chat_id=call.message.chat.id,
        message_id=call.message.id,
        text=get_queue_stats_text(queue),
        reply_markup=keyboards.queue_menu(),
    )
    await bot.answer_callback_query(callback_query_id=call.id)


@bot.message_handler(commands=["new_queue"])
async def nq(msg: Message):
    user = session.query(User).filter_by(id=msg.from_user.id).first()
    queue = Queue(owner_id=user.id)
    session.add(queue)
    session.commit()
    await bot.send_message(
        chat_id=msg.chat.id, text=f"Создана новая очередь: {queue.id}"
    )


@bot.message_handler(commands=["take_part"])
async def tp(msg: Message):
    try:
        command, queue_id = msg.text.split()
    except:
        await bot.send_message(chat_id=msg.chat.id, text="Вы забыли указать очередь")
        return
    qu = QueueUser(user_id=msg.from_user.id, queue_id=queue_id)
    session.add(qu)
    session.commit()
    await bot.send_message(chat_id=msg.chat.id, text="Вы приняли участие в очереди!")


@bot.message_handler(commands=["queue"])
async def q(msg: Message):
    try:
        command, queue_id = msg.text.split()
    except:
        await bot.send_message(chat_id=msg.chat.id, text="Вы забыли указать очередь")
        return
    queue = session.query(Queue).filter_by(id=queue_id).first()
    if queue:
        users = [q.user.name for q in queue.users]
        await bot.send_message(
            chat_id=msg.chat.id, text=f"Пользователи этой очереди: {', '.join(users)}"
        )
    else:
        await bot.send_message(chat_id=msg.chat.id, text="Очередь не найдена!")


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