133 lines
5.7 KiB
Python
133 lines
5.7 KiB
Python
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, method='send_message', **kwargs):
|
||
if text is None and not kwargs:
|
||
return
|
||
if reply_markup is None and remove_keyboard:
|
||
reply_markup = ReplyKeyboardRemove()
|
||
body = {
|
||
'chat_id': chat_id or self.chat_id,
|
||
}
|
||
if text:
|
||
body['text'] = text
|
||
if reply_markup:
|
||
body['reply_markup'] = reply_markup.to_json()
|
||
body.update(kwargs)
|
||
set_task(
|
||
'botalka_mailbox',
|
||
{
|
||
'project': 'roulette-bot',
|
||
'name': 'telegram-bot',
|
||
'method': method,
|
||
'body': body,
|
||
},
|
||
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_message(None, chat_to_send, method='send_photo', photo=self.message.photo[-1].file_id)
|
||
if self.message.sticker:
|
||
self.send_message(None, chat_to_send, method='send_data', data=self.message.sticker.file_id, data_type='sticker')
|
||
if self.message.voice:
|
||
self.send_message(None, chat_to_send, method='send_voice', voice=self.message.voice.file_id)
|
||
if self.message.video_note:
|
||
self.send_message(None, chat_to_send, method='send_video_note', data=self.message.video_note.file_id)
|
||
if self.message.animation:
|
||
self.send_message(None, chat_to_send, method='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()
|