roulette-bot/bot.py

import json
from telebot.types import Message, ReplyKeyboardRemove

from tools.mongo import mongo
from tools.queues import TasksHandlerMixin, set_task


class Core(TasksHandlerMixin):
    @property
    def queue_name(self):
        return 'roulette_bot_worker'

    def process(self, payload):
        message: Message = Message.de_json(json.dumps(payload))
        self.message = message
        self.chat_id = message.chat.id
        self.message_text = message.text or message.caption or ""
        print(f'Handled message from {self.chat_id}: {self.message_text}')
        user = mongo.chats_collection.find_one({"chat_id": message.chat.id})
        if user is None:
            doc = {
                "state": "new",
                "chat_id": message.chat.id,
            }
            mongo.chats_collection.insert_one(doc)
        else:
            doc = user
        self.doc = doc
        self.state = doc['state']
        if self.message_text.startswith('/'):
            self.exec_command()
            return
        getattr(self, "handle_state_" + self.state, self.handle_default)()

    def exec_command(self):
        if self.message_text == '/pause':
            self.set_state('pause')
            if self.state == 'dialog':
                current_dialog = mongo.get_current_dialog(self.chat_id)
                mongo.finish_dialog(current_dialog['_id'])
                if self.chat_id == current_dialog['chat_id_1']:
                    another_chat_id = current_dialog['chat_id_2']
                else:
                    another_chat_id = current_dialog['chat_id_1']
                self.send_message('🤖 Диалог окончен, жду тебя снова!')
                self.start_new_dialog([another_chat_id])
                return
            if self.state == 'pause':
                self.send_message('🤖 Сейчас твой аккаунт не активен. Активируй его с помощью команды /start')
                return
            if self.state == 'search':
                self.send_message('🤖 Поиск собеседника окончен, жду тебя снова!')
                return
        if self.message_text == '/next' or self.message_text == '/start':
            if self.state == 'dialog':
                dialog = mongo.get_current_dialog(self.chat_id)
                self.start_new_dialog([dialog['chat_id_1'], dialog['chat_id_2']])
                return
            else:
                self.start_new_dialog([self.chat_id])
                return

    def handle_state_search(self):
        self.send_message('🤖 Поиски собеседника продолжаются')

    def send_message(self, text, chat_id=None, reply_markup=None, remove_keyboard=True):
        if not text:
            return
        if reply_markup is None and remove_keyboard:
            reply_markup = ReplyKeyboardRemove()
        body = {
            'chat_id': chat_id or self.chat_id,
            'text': text,
        }
        if reply_markup:
            body['reply_markup'] = reply_markup.to_json()
        set_task(
            'botalka_mailbox', 
            {
                'project': 'roulette-bot', 
                'name': 'telegram-bot',
                'method': 'send_message',
                'body': body,
            },
            1,
        )

    def send(self, chat_id, method, **kwargs):
        set_task(
            'botalka_mailbox', 
            {
                'project': 'roulette-bot', 
                'name': 'telegram-bot',
                'method': method,
                'body': {
                    'chat_id': chat_id,
                    **kwargs,
                },
            },
            1,
        )

    def set_state(self, state, chat_ids=None):
        mongo.chats_collection.update_many({"chat_id": {"$in": chat_ids or [self.chat_id]}}, {"$set": {"state": state}})

    def handle_default(self):
        raise NotImplementedError(f"handler for {self.state} is not implemented")

    def handle_state_new(self):
        self.start_new_dialog([self.chat_id])

    def handle_state_dialog(self):
        current_dialog = mongo.get_current_dialog(self.chat_id)
        chat_to_send = current_dialog['chat_id_2'] if current_dialog['chat_id_1'] == self.chat_id else current_dialog['chat_id_1']
        if self.message.photo:
            self.send(chat_to_send, 'send_photo', photo=self.message.photo[-1].file_id)
        if self.message.sticker:
            self.send(chat_to_send, 'send_data', data=self.message.sticker.file_id, data_type='sticker')
        if self.message.voice:
            self.send(chat_to_send, 'send_voice', voice=self.message.voice.file_id)
        if self.message.video_note:
            self.send(chat_to_send, 'send_video_note', data=self.message.video_note.file_id)
        if self.message.animation:
            self.send(chat_to_send, 'send_animation', animation=self.message.animation.file_id)
        self.send_message(self.message_text, chat_to_send)

    def start_new_dialog(self, chat_ids):
        for chat in chat_ids:
            current_dialog = mongo.get_current_dialog(chat)
            if current_dialog:
                mongo.finish_dialog(current_dialog['_id'])
        self.set_state('search', chat_ids)
        for chat in chat_ids:
            self.send_message("🤖 Начинаю искать собеседника. Сообщу тебе, когда найду его.", chat)
            next_chat = mongo.find_searching(chat_ids)
            if not next_chat:
                continue
            self.send_message('🤖 Собеседник найден! Можешь начинать общаться', chat)
            self.send_message('🤖 Собеседник найден! Можешь начинать общаться', next_chat['chat_id'])
            mongo.create_dialog(chat, next_chat['chat_id'])
            self.set_state('dialog', [chat, next_chat['chat_id']])


if __name__ == '__main__':
    Core().poll()