# 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() changing_name = State() changing_queue_name = 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: user = User( id=msg.from_user.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) await asyncio.sleep(2) await bot.set_state(user_id=msg.from_user.id, state=States.default) if len(msg.text.split()) > 1: pass else: await bot.send_message( chat_id=msg.chat.id, text=textbook.menu.format(name=user.name), reply_markup=keyboards.menu() ) @bot.callback_query_handler(func=lambda c: c.data == "to_menu") async def to_menu_handler(call: types.CallbackQuery): user = session.query(User).filter_by(id=call.from_user.id).first() await bot.set_state(user_id=call.from_user.id, state=States.default) await bot.edit_message_text( chat_id=call.message.chat.id, message_id=call.message.id, text=textbook.menu.format(name=user.name), 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) session.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 queue_handler(call: types.CallbackQuery, queue_id: str = None): queue_id = call.data[2:] if not queue_id else queue_id 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_id"] = queue_id 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.callback_query_handler(func=lambda c: c.data == "edit_queue_name") async def edit_queue_name_handler(call: types.CallbackQuery): async with bot.retrieve_data( user_id=call.from_user.id, chat_id=call.message.chat.id ) as state_data: queue_id = state_data.get("queue_id", None) if not queue_id: await bot.answer_callback_query( callback_query_id=call.id, text=textbook.queue_operational_error ) return None queue = session.query(Queue).filter_by(id=queue_id).first() if queue.owner.id != call.from_user.id: await bot.answer_callback_query( callback_query_id=call.id, text=textbook.queue_operational_error ) return None await bot.set_state(user_id=call.from_user.id, state=States.changing_queue_name) await bot.edit_message_text( chat_id=call.message.chat.id, message_id=call.message.id, text=textbook.edit_name, reply_markup=keyboards.edit_name(), ) @bot.callback_query_handler(func=lambda c: c.data == "cancel", state=States.changing_queue_name) async def edit_queue_name_cancel_handler(call: types.CallbackQuery): async with bot.retrieve_data( user_id=call.from_user.id, chat_id=call.message.chat.id ) as state_data: queue_id = state_data.get("queue_id", None) await queue_handler(call, queue_id) @bot.callback_query_handler(func=lambda c: c.data == "settings") async def settings(call: types.CallbackQuery): await bot.set_state(user_id=call.from_user.id, state=States.default) await bot.edit_message_text( chat_id=call.message.chat.id, message_id=call.message.id, text=textbook.settings, reply_markup=keyboards.settings(), ) @bot.callback_query_handler(func=lambda c: c.data == "edit_name") async def edit_name_handler(call: types.CallbackQuery): await bot.set_state(user_id=call.from_user.id, state=States.changing_name) await bot.edit_message_text( chat_id=call.message.chat.id, message_id=call.message.id, text=textbook.edit_name, reply_markup=keyboards.edit_name(), ) @bot.callback_query_handler(func=lambda c: c.data == "cancel", state=States.changing_name) async def edit_name_cancel_handler(call: types.CallbackQuery): await settings(call) @bot.message_handler(content_types=["text"], state=States.changing_name) async def update_name(msg: Message): if len(msg.text) > 40 or "\n" in msg.text: await bot.send_message( chat_id=msg.chat.id, text=textbook.edit_name_error ) return None user = session.query(User).filter_by(id=msg.from_user.id).first() setattr(user, "name", msg.text) session.commit() await bot.send_message( chat_id=msg.chat.id, text=textbook.edit_name_success ) await asyncio.sleep(1) await start(msg) @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="Очередь не найдена!") @bot.message_handler(commands=["mystate"]) async def mystate(msg): state = await bot.get_state(user_id=msg.from_user.id) await bot.send_message(chat_id=msg.from_user.id, text=state) 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())