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', data=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()