95 lines
3.2 KiB
Python
95 lines
3.2 KiB
Python
# 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()
|
|
|
|
|
|
@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.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())
|