from __future__ import annotations import asyncio import logging import os import time import uuid from datetime import datetime, timedelta, timezone from datetime import time as dt_time from pathlib import Path import httpx from telegram import ( CopyTextButton, InlineKeyboardButton, InlineKeyboardMarkup, InputFile, LabeledPrice, Update, ) from telegram.constants import ChatAction, ChatType from telegram.error import BadRequest, NetworkError, TelegramError, TimedOut from telegram.ext import ( Application, CallbackQueryHandler, CommandHandler, ContextTypes, MessageHandler, PreCheckoutQueryHandler, TypeHandler, filters, ) from comfy_client import ComfyClient from config import Settings, setup_logging from db import Database, Request, UserAlreadyQueued from moderation import PromptBlocked, detect_person, is_prompt_allowed from prompt_enhancer import PromptEnhancer from queue_manager import QueueEvent, QueueFull, QueueManager from workflows import WorkflowBuilder logger = logging.getLogger(__name__) MSK = timezone(timedelta(hours=3)) POLLING_WATCHDOG_INTERVAL = 60.0 POLLING_WATCHDOG_STUCK_AFTER = 300.0 RESTART_HINT = "/restart — сброс контекста" CAPTION_LIMIT = 1024 def result_caption(prompt: str, *, limit: int = CAPTION_LIMIT) -> str: value = (prompt or "").strip() suffix = f"\n\n{RESTART_HINT}" available = max(0, limit - len(suffix)) if len(value) > available: value = value[: max(0, available - 1)].rstrip() if value: value += "…" return f"{value}{suffix}" if value else RESTART_HINT def polling_is_stuck( *, now: float, last_update_at: float, pending_update_count: int, stuck_after: float, ) -> bool: return ( pending_update_count > 0 and now - last_update_at >= stuck_after ) START_TEXT = ( "Пришли текст — сгенерирую картинку.\n" "Пришли фото с подписью — отредактирую его по описанию.\n" "Или ответь текстом на готовую картинку — возьму её как исходник.\n\n" "/balance — баланс генераций.\n" "/invite — пригласить друга.\n" "/status — состояние текущей задачи.\n" "/restart — сброс контекста." ) BLOCKED_TEXT = "Запрос нарушает правила. Звёзды не списаны." AGE_TEXT = ( "🔞 Бот доступен только пользователям 18+.\n\n" "Подтверди, что тебе исполнилось 18 лет." ) AGE_DENIED = "Доступ к боту разрешён только с 18 лет." class MpakoBot: def __init__(self, settings: Settings) -> None: self._settings = settings self._db = Database( settings.database_url, min_size=settings.db_pool_min_size, max_size=settings.db_pool_max_size, ) 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]] = {} self._daily_task: asyncio.Task[None] | None = None self._health_task: asyncio.Task[None] | None = None self._watchdog_task: asyncio.Task[None] | None = None self._last_update_at = time.monotonic() self._application: Application | None = None def run(self) -> None: builder = ( Application.builder() .token(self._settings.bot_token) .connect_timeout(10) .read_timeout(30) .write_timeout(30) .pool_timeout(5) .connection_pool_size(10) .get_updates_connect_timeout(10) .get_updates_read_timeout(15) .get_updates_write_timeout(15) .get_updates_pool_timeout(5) .get_updates_connection_pool_size(10) .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() self._application = application application.add_handler( TypeHandler(Update, self._track_update), group=-1 ) application.add_handler(CommandHandler("start", self._start)) application.add_handler(CommandHandler("help", self._help)) application.add_handler(CommandHandler("status", self._status)) application.add_handler(CommandHandler("balance", self._balance)) application.add_handler(CommandHandler("invite", self._invite)) application.add_handler(CommandHandler("restart", self._reset)) application.add_handler(CommandHandler("reset", self._reset)) application.add_handler( PreCheckoutQueryHandler(self._pre_checkout) ) application.add_handler( CallbackQueryHandler(self._queue_choice, pattern=r"^queue:") ) application.add_handler( CallbackQueryHandler(self._age_choice, pattern=r"^age:") ) application.add_handler( MessageHandler(filters.SUCCESSFUL_PAYMENT, self._successful_payment) ) 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() await self._db.seed_admins(self._settings.admin_user_ids) await self._db.accrue_daily_credits() self._daily_task = asyncio.create_task( self._daily_credits_loop(), name="daily-credits" ) self._health_task = asyncio.create_task( self._telegram_health_loop(), name="telegram-health" ) self._watchdog_task = asyncio.create_task( self._polling_watchdog_loop(), name="polling-watchdog" ) enhancer = ( PromptEnhancer(self._comfy, self._settings) if self._settings.prompt_enhance_enabled else None ) self._queue = QueueManager( db=self._db, comfy=self._comfy, workflows=self._workflows, settings=self._settings, on_event=self._on_queue_event, enhancer=enhancer, ) await self._queue.start() logger.info("Бот запущен, ComfyUI: %s", self._settings.comfy_base_url) async def _post_shutdown(self, application: Application) -> None: if self._watchdog_task is not None: self._watchdog_task.cancel() await asyncio.gather(self._watchdog_task, return_exceptions=True) self._watchdog_task = None if self._health_task is not None: self._health_task.cancel() await asyncio.gather(self._health_task, return_exceptions=True) self._health_task = None if self._daily_task is not None: self._daily_task.cancel() await asyncio.gather(self._daily_task, return_exceptions=True) self._daily_task = 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 _track_update( self, update: Update, context: ContextTypes.DEFAULT_TYPE ) -> None: self._last_update_at = time.monotonic() async def _polling_watchdog_loop(self) -> None: while True: await asyncio.sleep(POLLING_WATCHDOG_INTERVAL) pending = await self._pending_updates_count() if pending is None: continue if not polling_is_stuck( now=time.monotonic(), last_update_at=self._last_update_at, pending_update_count=pending, stuck_after=POLLING_WATCHDOG_STUCK_AFTER, ): continue logger.warning( "Polling watchdog: %s обновлений не обрабатываются," " перезапускаю updater", pending, ) await self._restart_polling() async def _pending_updates_count(self) -> int | None: url = ( "https://api.telegram.org/bot" f"{self._settings.bot_token}/getWebhookInfo" ) proxy = self._settings.telegram_proxy or None try: async with httpx.AsyncClient( proxy=proxy, timeout=httpx.Timeout(10.0, connect=5.0), ) as client: response = await client.get(url) data = response.json() except (httpx.HTTPError, ValueError) as exc: logger.warning( "Polling watchdog: проверка не удалась (%s)", type(exc).__name__, ) return None if not data.get("ok"): return None result = data.get("result") or {} return int(result.get("pending_update_count") or 0) async def _restart_polling(self) -> None: application = self._application if application is None or application.updater is None: return try: await asyncio.wait_for( application.updater.stop(), timeout=15 ) await asyncio.wait_for( application.updater.start_polling( allowed_updates=Update.ALL_TYPES, drop_pending_updates=False, ), timeout=30, ) except Exception: logger.exception( "Polling watchdog: не удалось перезапустить updater," " выхожу для рестарта контейнера" ) os._exit(1) self._last_update_at = time.monotonic() logger.warning("Polling watchdog: updater перезапущен") async def _daily_credits_loop(self) -> None: while True: now = datetime.now(MSK) next_midnight = datetime.combine( now.date() + timedelta(days=1), dt_time.min, MSK ) await asyncio.sleep(max(1.0, (next_midnight - now).total_seconds())) try: granted = await self._db.accrue_daily_credits() logger.info("Ежедневные генерации начислены: %s", granted) except Exception: logger.exception("Ошибка ежедневного начисления генераций") async def _telegram_health_loop(self) -> None: interval = self._settings.telegram_health_interval max_delay = self._settings.telegram_retry_max_delay delay = min(1.0, max_delay) while True: await asyncio.sleep(interval) while True: try: if self._bot is not None: await self._bot.get_me() break except (TimedOut, NetworkError): logger.warning( "Telegram health check failed, retry in %.0f s", delay ) await asyncio.sleep(delay) delay = min(max_delay, delay * 2) except TelegramError: logger.exception("Telegram health check error") break delay = min(1.0, max_delay) def _start_text(self, balance: int | None = None) -> str: parts: list[str] = [] if balance is not None: parts.append( f"Баланс: {balance} {self._credits_word(balance)}." ) if self._settings.payments_enabled: normal = self._settings.stars_price premium = self._settings.premium_stars_price parts.append( "Стоимость генерации:\n" f"• обычная очередь — {normal} звёзд;\n" f"• премиум-очередь — {premium} звёзд." ) parts.append(START_TEXT) return "\n\n".join(parts) @staticmethod def _credits_word(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 "генераций" @staticmethod def _parse_referrer(args: list[str]) -> int | None: if not args: return None raw = args[0].strip() if not raw.startswith("ref_"): return None try: return int(raw[4:]) except ValueError: return None async def _start( self, update: Update, context: ContextTypes.DEFAULT_TYPE ) -> None: if not await self._check_base_access(update): return message = update.effective_message user = update.effective_user if message is None or user is None: return row = await self._db.ensure_user(user.id, user.username) referrer_id = self._parse_referrer(context.args) if ( referrer_id and referrer_id != user.id and await self._db.get_user(referrer_id) is not None ): await self._db.set_referrer(user.id, referrer_id) row = await self._db.ensure_user(user.id, user.username) if row.age_declined: await message.reply_text(AGE_DENIED) return if not row.age_confirmed: await self._retry( message.reply_text, AGE_TEXT, reply_markup=self._age_keyboard(), ) return await self._retry( message.reply_text, self._start_text(row.credits) ) @staticmethod def _age_keyboard() -> InlineKeyboardMarkup: return InlineKeyboardMarkup( [ [ InlineKeyboardButton( "Мне 18+", callback_data="age:yes" ), InlineKeyboardButton( "Мне меньше 18", callback_data="age:no" ), ] ] ) async def _age_choice( self, update: Update, context: ContextTypes.DEFAULT_TYPE ) -> None: query = update.callback_query user = update.effective_user if query is None or query.data is None or user is None: return row = await self._db.ensure_user(user.id, user.username) if row.is_blocked: await query.answer("Доступ заблокирован.", show_alert=True) return if query.data == "age:yes": await self._db.set_age_confirmed(user.id, True) await query.answer("Возраст подтверждён") row = await self._db.get_user(user.id) try: await query.edit_message_text("Возраст подтверждён ✅") except TelegramError: pass if row is not None: await self._retry( self._bot.send_message, chat_id=user.id, text=self._start_text(row.credits), ) return if query.data == "age:no": await self._db.set_age_declined(user.id) await query.answer() try: await query.edit_message_text(AGE_DENIED) except TelegramError: pass async def _help( 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 row = await self._db.ensure_user(user.id, user.username) await self._retry( message.reply_text, self._start_text(row.credits) ) async def _balance( 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 row = await self._db.ensure_user(user.id, user.username) await self._retry( message.reply_text, f"Баланс: {row.credits} {self._credits_word(row.credits)}.", ) async def _invite( 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 self._bot is None: return await self._db.ensure_user(user.id, user.username) username = self._bot.username if not username: await message.reply_text("Не удалось получить ссылку. Попробуй позже.") return link = f"https://t.me/{username}?start=ref_{user.id}" keyboard = InlineKeyboardMarkup( [ [ InlineKeyboardButton( "Скопировать ссылку", copy_text=CopyTextButton(text=link), ) ] ] ) await self._retry( message.reply_text, "Пригласи друга:\n" f"{link}\n\n" "Когда друг впервые оплатит генерацию, ты получишь +3 генерации.", reply_markup=keyboard, ) async def _reset( 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 cleared = await self._db.clear_user_context(user.id) text = ( "Контекст сброшен. Следующий запрос начнёт новую сцену." if cleared else "Контекст уже пуст." ) await self._retry(message.reply_text, 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": label = " (премиум)" if active.is_premium else "" await message.reply_text(f"Твоя картинка сейчас генерируется{label}.") return position = await self._queue.position(active.id) if self._queue else 0 label = "премиум-очередь" if active.is_premium else "очередь" await message.reply_text( f"Задача в {label}. Перед тобой: {position} {self._people(position)}." ) async def _reward_referrer(self, user) -> None: if not await self._db.mark_first_payment(user.id): return row = await self._db.get_user(user.id) if row is None or row.referrer_id is None: return balance = await self._db.add_credits(row.referrer_id, 3) if balance is None: logger.warning("Не найден реферер %s", row.referrer_id) return text = ( f"Твой друг {row.display_user} сделал первую оплату. " f"+3 генерации. Баланс: {balance}." ) try: await self._retry( self._bot.send_message, chat_id=row.referrer_id, text=text ) except TelegramError: logger.exception("Не удалось уведомить реферера %s", row.referrer_id) async def _pre_checkout( self, update: Update, context: ContextTypes.DEFAULT_TYPE ) -> None: query = update.pre_checkout_query if query is None: return pending = await self._db.get_pending_payment(query.invoice_payload) user_row = await self._db.get_user(query.from_user.id) allowed = ( pending is not None and pending.user_id == query.from_user.id and pending.status == "pending" and query.currency == "XTR" and query.total_amount == pending.amount and user_row is not None and user_row.age_confirmed and not user_row.age_declined and not user_row.is_blocked ) try: await query.answer( ok=allowed, error_message=None if allowed else "Счёт устарел. Отправь запрос заново.", ) except TelegramError: logger.exception("Не удалось ответить на pre_checkout_query") async def _queue_choice( self, update: Update, context: ContextTypes.DEFAULT_TYPE ) -> None: query = update.callback_query if query is None or query.data is None: return try: _, payload, tier = query.data.split(":", 2) except ValueError: await query.answer("Некорректный выбор.", show_alert=True) return if tier not in {"normal", "premium", "credit"}: await query.answer("Некорректный выбор.", show_alert=True) return pending = await self._db.get_pending_payment(payload) if ( pending is None or pending.user_id != query.from_user.id or pending.status != "pending" or pending.amount != 0 ): await query.answer("Счёт уже создан или устарел.", show_alert=True) return active = await self._db.get_active_by_user(query.from_user.id) if active is not None: await query.answer( "У тебя уже есть активная задача.", show_alert=True ) return user_row = await self._db.get_user(query.from_user.id) if user_row is not None and user_row.is_blocked: await query.answer("Доступ заблокирован.", show_alert=True) return if user_row is None or not user_row.age_confirmed: await query.answer( "Сначала подтверди возраст.", show_alert=True ) if query.message is not None: await query.message.reply_text( AGE_TEXT, reply_markup=self._age_keyboard() ) return if tier == "credit": remaining = await self._db.consume_credit(query.from_user.id) if remaining is None: await query.answer( "Генерации закончились. Выбери оплату.", show_alert=True ) return await self._db.mark_pending_cancelled(payload) await query.answer() try: await query.edit_message_text( f"Списана 1 генерация. Осталось: {remaining}." ) except TelegramError: logger.debug("Не удалось обновить сообщение выбора очереди") accepted = await self._submit( update, mode=pending.mode, prompt=pending.prompt, input_path=pending.input_file_path, skip_payment=True, is_premium=False, use_context=bool(pending.context_prompt), context_base_prompt=pending.context_base_prompt, context_prompt=pending.context_prompt, ) if not accepted: await self._db.add_credits(query.from_user.id, 1) try: await query.edit_message_text( "Не удалось поставить задачу. Генерация возвращена." ) except TelegramError: pass return is_premium = tier == "premium" amount = ( self._settings.premium_stars_price if is_premium else self._settings.stars_price ) if not await self._db.set_pending_choice(payload, is_premium, amount): await query.answer("Счёт уже создан.", show_alert=True) return await query.answer() label = "Премиум-очередь" if is_premium else "Обычная очередь" try: await self._retry( self._bot.send_invoice, chat_id=query.from_user.id, title=label, description=( "Приоритетная генерация." if is_premium else "Обычная генерация." ), payload=payload, provider_token="", currency="XTR", prices=[LabeledPrice(label=label, amount=amount)], ) except TelegramError: logger.exception("Не удалось выставить счёт пользователю %s", query.from_user.id) await self._db.mark_pending_cancelled(payload) await query.edit_message_text("Не удалось выставить счёт. Попробуй позже.") return await query.edit_message_text( f"Счёт на {amount} звёзд отправлен. Очередь: {label.lower()}." ) async def _successful_payment( self, update: Update, context: ContextTypes.DEFAULT_TYPE ) -> None: message = update.effective_message user = update.effective_user payment = message.successful_payment if message is not None else None if message is None or user is None or payment is None: return payload = payment.invoice_payload pending = await self._db.get_pending_payment(payload) if pending is None or pending.user_id != user.id: logger.error( "Платёж %s не найден или принадлежит другому пользователю", payload, ) return if pending.status == "paid": await message.reply_text("Этот платёж уже обработан.") return if pending.status != "pending": return user_row = await self._db.get_user(user.id) if ( user_row is None or not user_row.age_confirmed or user_row.age_declined or user_row.is_blocked ): refunded = await self._refund_stars( user.id, payment.telegram_payment_charge_id ) if refunded: await self._db.mark_pending_refunded(payload) note = " Звёзды возвращены." if refunded else "" await message.reply_text("Доступ заблокирован." + note) return charge_id = payment.telegram_payment_charge_id if payment.currency != "XTR" or payment.total_amount != pending.amount: refunded = await self._refund_stars(user.id, charge_id) if refunded: await self._db.mark_pending_refunded(payload) note = " Звёзды возвращены." if refunded else "" await message.reply_text("Платёж не принят." + note) return if not await self._db.mark_pending_paid(payload, charge_id): await message.reply_text("Этот платёж уже обработан.") return accepted = await self._submit( update, mode=pending.mode, prompt=pending.prompt, input_path=pending.input_file_path, skip_payment=True, payment_charge_id=charge_id, payment_payload=payload, is_premium=pending.is_premium, use_context=bool(pending.context_prompt), context_base_prompt=pending.context_base_prompt, context_prompt=pending.context_prompt, ) if accepted: await self._reward_referrer(user) 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 reply = message.reply_to_message attachment = None suffix = ".jpg" if reply is not None: if reply.photo: attachment = reply.photo[-1] elif reply.document is not None and ( reply.document.mime_type or "" ).startswith("image/"): attachment = reply.document suffix = Path(reply.document.file_name or "").suffix.lower() or ".png" if attachment is None: await self._submit( update, mode="t2i", prompt=prompt, input_path=None ) return input_path = await self._download_attachment( message, attachment, user.id, suffix=suffix ) if input_path is None: return accepted = await self._submit( update, mode="i2i", prompt=prompt, input_path=str(input_path), ) if not accepted: input_path.unlink(missing_ok=True) 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] input_path = await self._download_attachment( message, photo, user.id, suffix=".jpg" ) if input_path is None: return accepted = await self._submit( update, mode="i2i", prompt=caption, input_path=str(input_path), use_context=False, ) if not accepted: input_path.unlink(missing_ok=True) async def _download_attachment( self, message, attachment, user_id: int, *, suffix: str = ".jpg", ) -> Path | None: if attachment.file_size and attachment.file_size > self._settings.max_photo_bytes: limit_mb = self._settings.max_photo_bytes / (1024 * 1024) try: await self._retry( message.reply_text, f"Файл слишком большой. Максимум {limit_mb:.0f} МБ.", ) except TelegramError: logger.debug("Не удалось сообщить о размере файла") return None suffix = suffix.lower() if not suffix.startswith(".") or len(suffix) > 8: suffix = ".png" user_dir = self._settings.storage_dir / str(user_id) input_path = user_dir / f"{uuid.uuid4().hex}_input{suffix}" try: user_dir.mkdir(parents=True, exist_ok=True) await self._retry( message.chat.send_action, ChatAction.TYPING, attempts=3 ) async def download() -> None: input_path.unlink(missing_ok=True) file = await attachment.get_file() await file.download_to_drive(custom_path=input_path) await self._retry(download, attempts=5) except (TelegramError, OSError): logger.exception("Не удалось скачать вложение пользователя %s", user_id) input_path.unlink(missing_ok=True) try: await self._retry( message.reply_text, "Не удалось скачать файл, попробуй ещё раз.", ) except TelegramError: pass return None return input_path async def _submit( self, update: Update, *, mode: str, prompt: str, input_path: str | None, skip_payment: bool = False, payment_charge_id: str | None = None, payment_payload: str | None = None, is_premium: bool = False, use_context: bool = True, context_base_prompt: str | None = None, context_prompt: str | None = 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 user_row = await self._db.ensure_user(user.id, user.username) if user_row.age_declined: await message.reply_text(AGE_DENIED) return False if not user_row.age_confirmed: await message.reply_text( AGE_TEXT, reply_markup=self._age_keyboard() ) return False is_admin = user_row.is_admin if user_row.is_blocked: await message.reply_text("Доступ заблокирован.") return False if not use_context: context_base_prompt = None context_prompt = None elif context_prompt is None: user_context = await self._db.get_user_context(user.id) if user_context is not None: context_base_prompt = user_context.base_prompt context_prompt = user_context.current_prompt credits = user_row.credits has_person = False if ( mode == "i2i" and input_path and not is_admin and self._settings.moderation_enabled ): has_person = await asyncio.to_thread(detect_person, input_path) if ( not skip_payment and not is_admin and self._settings.payments_enabled ): allowed, reason = is_prompt_allowed( prompt, mode=mode, has_person=has_person, extra_public_figures=self._settings.blocked_public_figures, ) if not allowed: await self._db.log_blocked( user_id=user.id, prompt=prompt, reason=reason or "unknown", mode=mode, has_person=has_person, ) await message.reply_text(BLOCKED_TEXT) return False active = await self._db.get_active_by_user(user.id) 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 counts = await self._queue.queue_counts() payload = uuid.uuid4().hex await self._db.create_pending_payment( payload=payload, user_id=user.id, mode=mode, prompt=prompt, username=user.username, input_file_path=input_path, amount=0, context_base_prompt=context_base_prompt, context_prompt=context_prompt, ) keyboard = InlineKeyboardMarkup( [ [ InlineKeyboardButton( ( f"Обычная — 1 генерация (осталось {credits})" if credits > 0 else f"Обычная — {self._settings.stars_price} звёзд" ), callback_data=( f"queue:{payload}:credit" if credits > 0 else f"queue:{payload}:normal" ), ) ], [ InlineKeyboardButton( f"Премиум — {self._settings.premium_stars_price} звёзд", callback_data=f"queue:{payload}:premium", ) ], ] ) text = ( f"Баланс: {credits} {self._credits_word(credits)}.\n" f"Обычная очередь: {counts['normal']} чел.\n" f"Премиум-очередь: {counts['premium']} чел.\n\n" "Выбери очередь:" ) try: await self._retry( message.reply_text, text, reply_markup=keyboard, ) except TelegramError: logger.exception("Не удалось показать выбор очереди %s", user.id) await self._db.mark_pending_cancelled(payload) return False return True if is_admin: is_premium = False try: request, position = await self._queue.enqueue( user_id=user.id, username=user.username, mode=mode, prompt=prompt, input_file_path=input_path, has_person=has_person, is_admin=is_admin, payment_charge_id=payment_charge_id, is_premium=is_premium, context_base_prompt=context_base_prompt, context_prompt=context_prompt, ) except PromptBlocked: note = await self._refund_if_paid( user.id, payment_charge_id, payment_payload ) await message.reply_text(BLOCKED_TEXT + note) return False except UserAlreadyQueued: note = await self._refund_if_paid( user.id, payment_charge_id, payment_payload ) 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 + note) return False except QueueFull as exc: note = await self._refund_if_paid( user.id, payment_charge_id, payment_payload ) await message.reply_text(exc.user_message + note) return False except Exception: logger.exception("Не удалось поставить задачу пользователю %s", user.id) note = await self._refund_if_paid( user.id, payment_charge_id, payment_payload ) await message.reply_text( "Не удалось поставить задачу. Попробуй позже." + note ) return False if payment_payload: await self._db.attach_pending_request(payment_payload, request.id) try: status_message = await self._retry( message.reply_text, self._ack_text(position, is_premium) ) 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 async def _refund_if_paid( self, user_id: int, charge_id: str | None, payload: str | None, ) -> str: if not charge_id: return "" refunded = await self._refund_stars(user_id, charge_id) if refunded and payload: await self._db.mark_pending_refunded(payload) if refunded: return " Звёзды возвращены." return " Не удалось вернуть звёзды, напиши админу." async def _refund_stars(self, user_id: int, charge_id: str) -> bool: if self._bot is None: return False try: return await self._retry( self._bot.refund_star_payment, user_id=user_id, telegram_payment_charge_id=charge_id, attempts=2, ) except TelegramError: logger.exception("Не удалось вернуть звёзды пользователю %s", user_id) return False def _ack_text(self, position: int, is_premium: bool = False) -> str: label = "Премиум-очередь. " if is_premium else "" if position <= 0: return f"Принято. {label}Ты первый, начинаю…" return f"Принято. {label}Перед тобой: {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 == "enhancing": await self._edit_status(event.request.id, "Улучшаю промпт…") 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 | None = None, **kwargs, ): max_attempts = ( self._settings.telegram_retry_attempts if attempts is None else attempts ) infinite = max_attempts <= 0 max_delay = self._settings.telegram_retry_max_delay delay = min(1.0, max_delay) attempt = 0 while True: attempt += 1 try: return await func(*args, **kwargs) except (TimedOut, NetworkError): if not infinite and attempt >= max_attempts: raise logger.warning( "Telegram timeout (attempt %s%s), retry in %.1f s", attempt, "" if infinite else f"/{max_attempts}", delay, ) await asyncio.sleep(delay) delay = min(max_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, attempts=3, ) 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 prompt = request.enhanced_prompt or request.prompt_text caption = result_caption(prompt) 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, ) 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, ) except TelegramError: await self._finish_failed( request, "Не удалось отправить результат в Telegram." ) return await self._clear_status(request.id, delete=True) await self._save_context(request) async def _save_context(self, request: Request) -> None: current = request.enhanced_prompt or request.prompt_text base = request.context_base_prompt or current try: await self._db.save_user_context( user_id=request.user_id, base_prompt=base, current_prompt=current, last_request_id=request.id, last_mode=request.mode, ) except Exception: logger.exception( "Не удалось сохранить контекст запроса %s", request.id ) 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 refund_note = "" if request.payment_charge_id and not request.payment_refunded: refunded = await self._refund_stars( request.user_id, request.payment_charge_id ) if refunded: await self._db.set_refunded(request.id) refund_note = " Звёзды возвращены." else: refund_note = " Не удалось вернуть звёзды, напиши админу." text = f"Не получилось: {error}{refund_note}" 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_base_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 row = await self._db.get_user(user.id) if row is not None and row.is_blocked: await message.reply_text("Доступ заблокирован.") return False return True async def _check_access(self, update: Update) -> bool: if not await self._check_base_access(update): return False message = update.effective_message user = update.effective_user if message is None or user is None: return False row = await self._db.get_user(user.id) if row is not None and row.age_declined: await message.reply_text(AGE_DENIED) return False if row is None or not row.age_confirmed: await message.reply_text( AGE_TEXT, reply_markup=self._age_keyboard() ) 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()