import telebot from telebot.types import Message from daemons.bot import bot from daemons.fetch import fetch_schedule_for_user from helpers.models import UserSchema, User from helpers.mongo import mongo from helpers.ruz import ruz class BaseAnswer: def process(self, message: Message): user = mongo.users_collection.find_one({"chat_id": message.chat.id}) if user is None: user = User(chat_id=message.chat.id) mongo.users_collection.insert_one(UserSchema().dump(user)) else: user = UserSchema().load(user) attr = getattr(self, "handle_state_" + user.state, None) if attr is None: raise NotImplementedError(f"handled state {user.state} but not implemented!") attr(message, user) def set_state(self, user: User, state: str): user.state = state mongo.users_collection.update_one({"chat_id": user.chat_id}, {"$set": {"state": state}}) class Answer(BaseAnswer): def handle_state_new(self, message: Message, user: User): bot.send_message( message.chat.id, "Привет! Я буду помогать тебе выживать в вышке!\nДля начала пришли мне свое ФИО.", ) self.set_state(user, "wait_for_name") def handle_state_wait_for_name(self, message: Message, user: User): kb = telebot.types.ReplyKeyboardMarkup(True, False) data = ruz.find_person(message.text) if data is None: bot.send_message( user.chat_id, "В РУЗе какая-то поломка, попробуй еще раз позже." ) return if len(data) == 0: bot.send_message(user.chat_id, "К сожалению, в РУЗе не нашлось такого студента, попробуй еще раз.") return for entity in data: kb.row(entity['description']) user.name = message.text mongo.users_collection.update_one( {"chat_id": user.chat_id}, {"$set": {"name": user.name}}) bot.send_message( user.chat_id, "Отлично! Теперь выбери из списка свою группу.", reply_markup=kb ) self.set_state(user, "wait_for_group") def handle_state_wait_for_group(self, message: Message, user: User): group = message.text data = ruz.find_person(user.name) if data is None: bot.send_message( user.chat_id, "В РУЗе какая-то поломка, попробуй еще раз позже." ) return for element in data: if element['description'] == group: user.hse_id = int(element['id']) user.group = group user.name = element['label'] break mongo.users_collection.update_one({"chat_id": user.chat_id}, {"$set": { "hse_id": user.hse_id, "group": group, "name": user.name }}) bot.send_message( user.chat_id, "Я нашел тебя в базе РУЗ. Я буду подсказывать тебе расписание, а также уведомлять о предстоящих парах.", ) success = fetch_schedule_for_user(user.hse_id) if success: kb = telebot.types.ReplyKeyboardMarkup(True, False) kb.row("Пары сегодня") kb.row("Уведомления") lessons = mongo.get_today_lessons(user) if len(lessons) == 0: bot.send_message( user.chat_id, "Сегодня у тебя нет пар.", reply_markup=kb ) else: bot.send_message( user.chat_id, ruz.schedule_builder(lessons), reply_markup=kb ) self.set_state(user, "ready") def handle_state_ready(self, message: Message, user: User): kb = telebot.types.ReplyKeyboardMarkup(True, False) kb.row("Пары сегодня") kb.row("Уведомления") if message.text == "Пары сегодня": lessons = mongo.get_today_lessons(user) if len(lessons) == 0: text = "Сегодня у тебя нет пар." else: text = ruz.schedule_builder(lessons) elif message.text == "Уведомления": kb = telebot.types.ReplyKeyboardMarkup(True, False) kb.row("Не уведомлять") kb.row("5 минут") kb.row("10 минут") kb.row("15 минут") kb.row("20 минут") bot.send_message( user.chat_id, "Выбери когда мне нужно тебе написать о предстоящей паре", reply_markup=kb ) self.set_state(user, "wait_for_notify") return else: text = "Я не понимаю такой команды, используй кнопки." bot.send_message( user.chat_id, text, reply_markup=kb ) def handle_state_wait_for_notify(self, message: Message, user: User): text = message.text if text == "Не уведомлять": user.notify_minutes = None elif text == "5 минут": user.notify_minutes = 5 elif text == "10 минут": user.notify_minutes = 10 elif text == "15 минут": user.notify_minutes = 15 elif text == "20 минут": user.notify_minutes = 20 else: kb = telebot.types.ReplyKeyboardMarkup(True, False) kb.row("Не уведомлять") kb.row("5 минут") kb.row("10 минут") kb.row("15 минут") kb.row("20 минут") bot.send_message(user.chat_id, "Я не понимаю такой команды, используй кнопки.", reply_markup=kb) return mongo.users_collection.update_one({"chat_id": user.chat_id}, {"$set": {"notify_minutes": user.notify_minutes}}) if user.notify_minutes is not None: text = f"Принято! Буду уведомлять тебя за {text}." else: text = f"Принято! Я не уведомлять тебя." kb = telebot.types.ReplyKeyboardMarkup(True, False) kb.row("Пары сегодня") kb.row("Уведомления") bot.send_message(user.chat_id, text, reply_markup=kb) self.set_state(user, "ready") answer = Answer()