ruz-bot/helpers/answer.py

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