159 lines
6.7 KiB
Python
159 lines
6.7 KiB
Python
import os
|
||
import uuid
|
||
|
||
import requests
|
||
import telebot
|
||
from telebot.types import Message, ReplyKeyboardRemove
|
||
|
||
from tools.minio import minio_client as minio
|
||
from tools.mongo import mongo
|
||
from tools.sprint_platform import platform
|
||
from tools.redis import redis_client as redis
|
||
|
||
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 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']
|
||
|
||
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)
|
||
with redis.lock('search'):
|
||
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':
|
||
with redis.lock('search'):
|
||
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 not text:
|
||
return
|
||
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)
|
||
chat_to_send = current_dialog['chat_id_2'] if current_dialog['chat_id_1'] == self.chat_id else current_dialog['chat_id_1']
|
||
saves = platform.get_config('save')
|
||
if saves['messages']:
|
||
res = mongo.create_message(self.message.content_type, self.message_text, current_dialog['_id'], self.chat_id).inserted_id
|
||
else:
|
||
res = uuid.uuid4()
|
||
if self.message.photo:
|
||
if saves['photos']:
|
||
photo = requests.get(bot.get_file_url(self.message.photo[-1].file_id)).content
|
||
minio.put_object(f"photos/{res}", photo)
|
||
bot.send_photo(chat_to_send, self.message.photo[-1].file_id)
|
||
if self.message.sticker:
|
||
if saves['stickers']:
|
||
sticker = requests.get(bot.get_file_url(self.message.sticker.file_id)).content
|
||
minio.put_object(f"stickers/{res}", sticker)
|
||
bot.send_sticker(chat_to_send, self.message.sticker.file_id)
|
||
if self.message.voice:
|
||
if saves['voices']:
|
||
voice = requests.get(bot.get_file_url(self.message.voice.file_id)).content
|
||
minio.put_object(f"voices/{res}", voice)
|
||
bot.send_voice(chat_to_send, self.message.voice.file_id)
|
||
if self.message.video_note:
|
||
if saves['video_notes']:
|
||
video_note = requests.get(bot.get_file_url(self.message.video_note.file_id)).content
|
||
minio.put_object(f"video_notes/{res}", video_note)
|
||
bot.send_video_note(chat_to_send, self.message.video_note.file_id)
|
||
if self.message.animation:
|
||
if saves['gifs']:
|
||
video_note = requests.get(bot.get_file_url(self.message.animation.file_id)).content
|
||
minio.put_object(f"gifs/{res}", video_note)
|
||
bot.send_animation(chat_to_send, self.message.animation.file_id)
|
||
self.send_message(self.message_text, chat_to_send)
|
||
|
||
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(content_types=[
|
||
# 'audio',
|
||
'photo',
|
||
'voice',
|
||
'video_note',
|
||
# 'document',
|
||
'text',
|
||
'animation',
|
||
# 'location',
|
||
# 'contact',
|
||
'sticker'
|
||
]
|
||
)
|
||
def do_action(message: Message):
|
||
try:
|
||
Core(message).process()
|
||
except Exception as e:
|
||
print(e)
|
||
|
||
print('bot is starting')
|
||
bot.polling()
|