from __future__ import annotations import asyncio import logging import time import uuid from pathlib import Path from telegram import InputFile, Update from telegram.constants import ChatAction, ChatType from telegram.error import BadRequest, NetworkError, TelegramError, TimedOut from telegram.ext import ( Application, CommandHandler, ContextTypes, MessageHandler, filters, ) from comfy_client import ComfyClient from config import Settings, setup_logging from db import Database, Request, UserAlreadyQueued from queue_manager import QueueEvent, QueueFull, QueueManager from workflows import WorkflowBuilder logger = logging.getLogger(__name__) START_TEXT = ( "Пришли текст — сгенерирую картинку.\n" "Пришли фото с подписью — отредактирую его по описанию.\n\n" "/status — узнать состояние текущей задачи." ) class MpakoBot: def __init__(self, settings: Settings) -> None: self._settings = settings self._db = Database(settings.db_path) self._workflows: WorkflowBuilder | None = None self._comfy: ComfyClient | None = None self._queue: QueueManager | None = None self._bot = None self._status_messages: dict[int, tuple[int, int]] = {} self._last_edit: dict[int, tuple[float, int]] = {} def run(self) -> None: builder = ( Application.builder() .token(self._settings.bot_token) .connect_timeout(20) .read_timeout(60) .write_timeout(60) .get_updates_connect_timeout(20) .get_updates_read_timeout(60) .post_init(self._post_init) .post_shutdown(self._post_shutdown) ) if self._settings.telegram_proxy: builder = builder.proxy(self._settings.telegram_proxy).get_updates_proxy( self._settings.telegram_proxy ) application = builder.build() application.add_handler(CommandHandler("start", self._start)) application.add_handler(CommandHandler("help", self._help)) application.add_handler(CommandHandler("status", self._status)) application.add_handler(MessageHandler(filters.PHOTO, self._photo)) application.add_handler( MessageHandler(filters.TEXT & ~filters.COMMAND, self._text) ) application.add_error_handler(self._on_error) application.run_polling(allowed_updates=Update.ALL_TYPES) async def _post_init(self, application: Application) -> None: self._bot = application.bot self._settings.storage_dir.mkdir(parents=True, exist_ok=True) self._workflows = WorkflowBuilder( self._settings.workflow_t2i, self._settings.workflow_i2i, self._settings.workflow_mapping, ) self._comfy = ComfyClient( self._settings.comfy_base_url, self._settings.comfy_ws_url, ) await self._db.connect() self._queue = QueueManager( db=self._db, comfy=self._comfy, workflows=self._workflows, settings=self._settings, on_event=self._on_queue_event, ) await self._queue.start() logger.info("Бот запущен, ComfyUI: %s", self._settings.comfy_base_url) async def _post_shutdown(self, application: Application) -> None: if self._queue is not None: await self._queue.stop() if self._comfy is not None: await self._comfy.close() await self._db.close() logger.info("Бот остановлен") async def _start( self, update: Update, context: ContextTypes.DEFAULT_TYPE ) -> None: if not await self._check_access(update): return message = update.effective_message if message is not None: await message.reply_text(START_TEXT) async def _help( self, update: Update, context: ContextTypes.DEFAULT_TYPE ) -> None: if not await self._check_access(update): return message = update.effective_message if message is not None: await message.reply_text(START_TEXT) async def _status( self, update: Update, context: ContextTypes.DEFAULT_TYPE ) -> None: if not await self._check_access(update): return message = update.effective_message user = update.effective_user if message is None or user is None: return active = await self._db.get_active_by_user(user.id) if active is None: await message.reply_text("Активных задач нет.") return if active.status == "running": await message.reply_text("Твоя картинка сейчас генерируется.") return position = await self._queue.position(active.id) if self._queue else 0 await message.reply_text( f"Задача в очереди. Перед тобой: {position} {self._people(position)}." ) async def _text( self, update: Update, context: ContextTypes.DEFAULT_TYPE ) -> None: if not await self._check_access(update): return message = update.effective_message user = update.effective_user if message is None or user is None or not message.text: return prompt = message.text.strip() if not prompt: await message.reply_text("Промпт пустой. Напиши, что нарисовать.") return if len(prompt) > self._settings.max_prompt_length: await message.reply_text( f"Промпт слишком длинный: максимум {self._settings.max_prompt_length} символов." ) return await self._submit(update, mode="t2i", prompt=prompt, input_path=None) async def _photo( self, update: Update, context: ContextTypes.DEFAULT_TYPE ) -> None: if not await self._check_access(update): return message = update.effective_message user = update.effective_user if message is None or user is None or not message.photo: return caption = (message.caption or "").strip() if not caption: await message.reply_text( "Пришли фото с подписью — что именно изменить на изображении." ) return if len(caption) > self._settings.max_prompt_length: await message.reply_text( f"Описание слишком длинное: максимум {self._settings.max_prompt_length} символов." ) return photo = message.photo[-1] if photo.file_size and photo.file_size > self._settings.max_photo_bytes: limit_mb = self._settings.max_photo_bytes / (1024 * 1024) await message.reply_text(f"Фото слишком большое. Максимум {limit_mb:.0f} МБ.") return user_dir = self._settings.storage_dir / str(user.id) user_dir.mkdir(parents=True, exist_ok=True) input_path = user_dir / f"{uuid.uuid4().hex}_input.jpg" try: await message.chat.send_action(ChatAction.TYPING) file = await photo.get_file() await file.download_to_drive(custom_path=input_path) except TelegramError: logger.exception("Не удалось скачать фото пользователя %s", user.id) input_path.unlink(missing_ok=True) await message.reply_text("Не удалось скачать фото, попробуй ещё раз.") return accepted = await self._submit( update, mode="i2i", prompt=caption, input_path=str(input_path), ) if not accepted: input_path.unlink(missing_ok=True) async def _submit( self, update: Update, *, mode: str, prompt: str, input_path: str | None, ) -> bool: message = update.effective_message user = update.effective_user if message is None or user is None or self._queue is None: return False try: request, position = await self._queue.enqueue( user_id=user.id, mode=mode, prompt=prompt, input_file_path=input_path, ) except UserAlreadyQueued: active = await self._db.get_active_by_user(user.id) text = "Ты уже в очереди." if active is not None: if active.status == "running": text += " Твоя картинка уже генерируется." else: ahead = await self._queue.position(active.id) text += f" Перед тобой: {ahead} {self._people(ahead)}." await message.reply_text(text) return False except QueueFull as exc: await message.reply_text(exc.user_message) return False except Exception: logger.exception("Не удалось поставить задачу пользователю %s", user.id) await message.reply_text("Не удалось поставить задачу. Попробуй позже.") return False try: status_message = await self._retry( message.reply_text, self._ack_text(position) ) except TelegramError: logger.exception( "Не удалось отправить подтверждение по запросу %s", request.id ) return True self._status_messages[request.id] = ( status_message.chat_id, status_message.message_id, ) # Воркер мог стартовать раньше, чем мы сохранили message_id. current = await self._db.get_request(request.id) if current is None: return True if current.status == "running": await self._edit_status(request.id, "Генерация…", percent=0) elif current.status == "done": await self._clear_status(request.id, delete=True) elif current.status == "failed": await self._finish_failed( current, current.error_text or "Не удалось сгенерировать изображение.", ) return True def _ack_text(self, position: int) -> str: if position <= 0: return "Принято. Ты первый, начинаю…" return f"Принято. Перед тобой: {position} {self._people(position)}." @staticmethod def _people(count: int) -> str: if count % 10 == 1 and count % 100 != 11: return "человек" if count % 10 in (2, 3, 4) and count % 100 not in (12, 13, 14): return "человека" return "человек" async def _on_queue_event(self, event: QueueEvent) -> None: if event.kind == "started": await self._edit_status(event.request.id, "Генерация…", percent=0) elif event.kind == "progress": percent = event.percent if event.percent is not None else 0 await self._edit_status( event.request.id, f"Генерация… {percent}%", percent=percent, ) elif event.kind == "done": await self._finish_done(event.request) elif event.kind == "failed": await self._finish_failed( event.request, event.error or "Не удалось сгенерировать изображение.", ) async def _retry(self, func, *args, attempts: int = 3, **kwargs): delay = 1.0 for attempt in range(1, attempts + 1): try: return await func(*args, **kwargs) except (TimedOut, NetworkError): if attempt >= attempts: raise logger.warning( "Telegram timeout (%s/%s), повтор через %.0f с", attempt, attempts, delay, ) await asyncio.sleep(delay) delay *= 2 async def _edit_status( self, request_id: int, text: str, *, percent: int | None = None, ) -> None: if self._bot is None: return target = self._status_messages.get(request_id) if target is None: return now = time.monotonic() last = self._last_edit.get(request_id) if ( last is not None and percent is not None and percent - last[1] < 3 and now - last[0] < 2.5 ): return chat_id, message_id = target try: await self._retry( self._bot.edit_message_text, chat_id=chat_id, message_id=message_id, text=text, ) except BadRequest as exc: if "not modified" not in str(exc).lower(): logger.debug("Не удалось обновить статус %s: %s", request_id, exc) return except TelegramError as exc: logger.debug("Не удалось обновить статус %s: %s", request_id, exc) return self._last_edit[request_id] = ( now, percent if percent is not None else -1, ) async def _finish_done(self, request: Request) -> None: target = self._status_messages.get(request.id) chat_id = target[0] if target is not None else request.user_id output_path = request.output_file_path if not output_path or not Path(output_path).exists(): await self._finish_failed(request, "Файл результата не найден.") return if self._bot is None: return caption = f"Готово. Seed: {request.seed}" if request.seed is not None else "Готово." filename = Path(output_path).name try: content = await asyncio.to_thread(Path(output_path).read_bytes) except OSError: logger.exception("Не удалось прочитать результат запроса %s", request.id) await self._finish_failed(request, "Файл результата не читается.") return try: await self._retry( self._bot.send_photo, chat_id=chat_id, photo=InputFile(content, filename=filename), caption=caption, read_timeout=120, write_timeout=120, attempts=2, ) except TelegramError: logger.exception("Не удалось отправить фото по запросу %s", request.id) try: await self._retry( self._bot.send_document, chat_id=chat_id, document=InputFile(content, filename=filename), caption=caption, read_timeout=120, write_timeout=120, attempts=2, ) except TelegramError: await self._finish_failed( request, "Не удалось отправить результат в Telegram." ) return await self._clear_status(request.id, delete=True) async def _finish_failed(self, request: Request, error: str) -> None: target = self._status_messages.get(request.id) chat_id = target[0] if target is not None else request.user_id text = f"Не получилось: {error}" if target is not None and self._bot is not None: try: await self._retry( self._bot.edit_message_text, chat_id=target[0], message_id=target[1], text=text, ) except TelegramError: logger.debug("Не удалось показать ошибку по запросу %s", request.id) if self._bot is not None: try: await self._retry( self._bot.send_message, chat_id=chat_id, text=text ) except TelegramError: logger.exception( "Не удалось отправить ошибку по запросу %s", request.id ) await self._clear_status(request.id, delete=False) return if self._bot is not None: try: await self._retry( self._bot.send_message, chat_id=chat_id, text=text ) except TelegramError: logger.exception("Не удалось отправить ошибку по запросу %s", request.id) async def _clear_status(self, request_id: int, *, delete: bool) -> None: target = self._status_messages.pop(request_id, None) self._last_edit.pop(request_id, None) if delete and target is not None and self._bot is not None: try: await self._retry( self._bot.delete_message, chat_id=target[0], message_id=target[1], attempts=2, ) except TelegramError: logger.debug("Не удалось удалить статус %s", request_id) async def _check_access(self, update: Update) -> bool: message = update.effective_message chat = update.effective_chat user = update.effective_user if message is None or chat is None or user is None: return False if chat.type != ChatType.PRIVATE: await message.reply_text("Напиши боту в личку.") return False if ( self._settings.allowed_user_ids and user.id not in self._settings.allowed_user_ids ): await message.reply_text("Доступ к боту ограничен.") return False return True async def _on_error( self, update: object, context: ContextTypes.DEFAULT_TYPE ) -> None: logger.error("Ошибка обработчика", exc_info=context.error) if isinstance(update, Update) and update.effective_message is not None: try: await update.effective_message.reply_text( "Внутренняя ошибка. Попробуй ещё раз." ) except TelegramError: pass def main() -> None: settings = Settings.from_env() setup_logging(settings.log_level) MpakoBot(settings).run() if __name__ == "__main__": main()