107 lines
4.4 KiB
Python
107 lines
4.4 KiB
Python
import os
|
||
|
||
import telebot
|
||
from telebot.types import Message, ReplyKeyboardRemove
|
||
|
||
from mongo import mongo
|
||
|
||
bot = telebot.TeleBot(os.getenv("TELEGRAM_TOKEN"))
|
||
|
||
|
||
class Core:
|
||
def __init__(self, message: Message):
|
||
self.message = message
|
||
self.chat_id = message.chat.id
|
||
self.message_text = message.text
|
||
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']
|
||
|
||
def process(self):
|
||
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, **kwargs):
|
||
if reply_markup is None and remove_keyboard:
|
||
reply_markup = ReplyKeyboardRemove()
|
||
bot.send_message(chat_id or self.chat_id, text, reply_markup=reply_markup, **kwargs)
|
||
|
||
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)
|
||
mongo.create_message(self.message_text, current_dialog['_id'], self.chat_id)
|
||
if current_dialog['chat_id_1'] == self.chat_id:
|
||
self.send_message(self.message_text, current_dialog['chat_id_2'])
|
||
else:
|
||
self.send_message(self.message_text, current_dialog['chat_id_1'])
|
||
|
||
def start_new_dialog(self, chat_ids):
|
||
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']])
|
||
|
||
|
||
def run_bot():
|
||
@bot.message_handler()
|
||
def do_action(message: Message):
|
||
Core(message).process()
|
||
|
||
bot.polling()
|