"""Обработчики настроек хранения и вспомогательные команды."""
import logging
from aiogram import Router, F
from aiogram.types import CallbackQuery, Message, InlineKeyboardMarkup, InlineKeyboardButton
from aiogram.fsm.context import FSMContext
from aiogram.fsm.state import State, StatesGroup
from aiogram.filters import Command
from config import settings
from database.database import db
from sd.sd_client import sd_client
from utils.image_manager import cleanup_expired_images
router = Router()
logger = logging.getLogger(__name__)
class StorageSettingsState(StatesGroup):
"""Состояния для настройки хранения."""
waiting_for_ttl = State()
# --- Обработчик главного меню ---
@router.callback_query(F.data == "main_menu")
async def main_menu(callback: CallbackQuery, state: FSMContext = None):
"""Возврат в главное меню."""
# Очищаем состояние, если оно есть
if state:
await state.clear()
from bot.handlers_start import get_main_keyboard
text = (
f"👋 Привет, {callback.from_user.first_name}!\n\n"
"Я бот для генерации изображений через Stable Diffusion.\n"
"Выберите действие из меню ниже:"
)
try:
await callback.message.edit_text(
text,
parse_mode="HTML",
reply_markup=get_main_keyboard(callback.from_user.id),
)
except Exception:
# Если edit_text не сработал (напр., сообщение с фото), отправляем новое
await callback.message.answer(
text,
parse_mode="HTML",
reply_markup=get_main_keyboard(callback.from_user.id),
)
await callback.answer()
# --- Проверка статуса SD API ---
@router.callback_query(F.data == "check_sd_status")
async def check_sd_status(callback: CallbackQuery):
"""Проверка соединения с SD API."""
status_msg = await callback.message.edit_text("🔍 Проверяю соединение с SD API...")
is_connected = await sd_client.check_connection()
if is_connected:
current_model = await sd_client.get_current_model()
text = (
f"✅ SD API подключено!\n\n"
f"Адрес: {settings.SD_API_URL}\n"
f"Текущая модель: {current_model}"
)
else:
text = (
f"❌ SD API недоступно!\n\n"
f"Адрес: {settings.SD_API_URL}\n\n"
"Проверьте:\n"
"1. Запущен ли Stable Diffusion WebUI\n"
"2. Доступен ли сервер по указанному адресу\n"
"3. Правильность настройки API (--api флаг)"
)
keyboard = InlineKeyboardMarkup(inline_keyboard=[
[InlineKeyboardButton(text="🔄 Повторить проверку", callback_data="check_sd_status")],
[InlineKeyboardButton(text="📋 Главное меню", callback_data="main_menu")],
])
await status_msg.edit_text(text, parse_mode="HTML", reply_markup=keyboard)
await callback.answer()
# --- Помощь ---
@router.callback_query(F.data == "help")
async def help_command(callback: CallbackQuery):
"""Справка по боту."""
text = (
"ℹ️ Справка по боту\n\n"
"Генерация изображений:\n"
"• txt2img — создание изображения по текстовому описанию\n\n"
"Профили:\n"
"Создавайте профили с предустановленными настройками:\n"
"• Размер изображения\n"
"• Количество шагов\n"
"• CFG Scale\n"
"• Сэмплер\n"
"• Модель и LoRA\n\n"
"Вы можете установить профиль по умолчанию для быстрой генерации.\n\n"
"Настройки хранения:\n"
"Настройте время хранения сгенерированных изображений.\n"
"Изображения автоматически удаляются по истечении срока.\n\n"
"Команды:\n"
"/start — Главное меню\n"
"/menu — Показать главное меню\n"
"/cancel — Отменить текущую операцию\n"
"/status — Статус SD API"
)
keyboard = InlineKeyboardMarkup(inline_keyboard=[
[InlineKeyboardButton(text="📋 Главное меню", callback_data="main_menu")],
])
await callback.message.edit_text(text, parse_mode="HTML", reply_markup=keyboard)
await callback.answer()
@router.message(Command("status"))
async def cmd_status(message: Message):
"""Команда /status."""
is_connected = await sd_client.check_connection()
if is_connected:
current_model = await sd_client.get_current_model()
text = (
f"✅ SD API подключено!\n\n"
f"Адрес: {settings.SD_API_URL}\n"
f"Текущая модель: {current_model}"
)
else:
text = (
f"❌ SD API недоступно!\n\n"
f"Адрес: {settings.SD_API_URL}"
)
await message.answer(text, parse_mode="HTML")
# --- Настройки хранения ---
@router.callback_query(F.data == "storage_settings")
async def storage_settings(callback: CallbackQuery):
"""Меню настроек хранения."""
user_settings = await db.get_user_settings(callback.from_user.id)
current_ttl = user_settings.get("image_ttl_hours", settings.DEFAULT_IMAGE_TTL_HOURS)
text = (
"🗃️ Настройки хранения изображений\n\n"
f"Текущее время хранения: {current_ttl} часов\n"
f"Максимальное время: {settings.MAX_IMAGE_TTL_HOURS} часов\n\n"
"Изображения автоматически удаляются по истечении срока.\n"
"Очистка происходит каждые {interval} минут.".format(interval=settings.CLEANUP_INTERVAL_MINUTES)
)
keyboard = InlineKeyboardMarkup(inline_keyboard=[
[InlineKeyboardButton(text="12 часов", callback_data="set_ttl_12")],
[InlineKeyboardButton(text="24 часа (1 день)", callback_data="set_ttl_24")],
[InlineKeyboardButton(text="48 часов (2 дня)", callback_data="set_ttl_48")],
[InlineKeyboardButton(text="72 часа (3 дня)", callback_data="set_ttl_72")],
[InlineKeyboardButton(text="168 часов (7 дней)", callback_data="set_ttl_168")],
[InlineKeyboardButton(text="✏️ Ввести своё значение", callback_data="set_ttl_custom")],
[InlineKeyboardButton(text="⬅️ Назад", callback_data="main_menu")],
])
await callback.message.edit_text(text, parse_mode="HTML", reply_markup=keyboard)
await callback.answer()
@router.callback_query(F.data.startswith("set_ttl_"))
async def set_ttl(callback: CallbackQuery, state: FSMContext):
"""Установка времени хранения."""
if callback.data == "set_ttl_custom":
await state.set_state(StorageSettingsState.waiting_for_ttl)
await callback.message.edit_text(
f"⏰ Введите время хранения в часах\n\n"
f"Допустимый диапазон: 1 - {settings.MAX_IMAGE_TTL_HOURS} часов.",
parse_mode="HTML",
)
await callback.answer()
return
ttl_hours = int(callback.data.replace("set_ttl_", ""))
await _save_ttl(callback, ttl_hours)
@router.message(StorageSettingsState.waiting_for_ttl)
async def save_custom_ttl(message: Message, state: FSMContext):
"""Сохранение пользовательского значения TTL."""
try:
ttl_hours = int(message.text.strip())
if ttl_hours < 1 or ttl_hours > settings.MAX_IMAGE_TTL_HOURS:
await message.answer(
f"Значение должно быть от 1 до {settings.MAX_IMAGE_TTL_HOURS} часов."
)
return
except ValueError:
await message.answer("Введите корректное целое число.")
return
await state.clear()
await _save_ttl(message, ttl_hours)
async def _save_ttl(target, ttl_hours: int):
"""Сохранить TTL и показать подтверждение."""
user_id = target.from_user.id
await db.update_user_settings(user_id, image_ttl_hours=ttl_hours)
text = (
f"✅ Настройки сохранены!\n\n"
f"Время хранения изображений: {ttl_hours} часов\n"
f"Изображения будут автоматически удаляться по истечении этого срока."
)
keyboard = InlineKeyboardMarkup(inline_keyboard=[
[InlineKeyboardButton(text="🗃️ Другие настройки хранения", callback_data="storage_settings")],
[InlineKeyboardButton(text="📋 Главное меню", callback_data="main_menu")],
])
if hasattr(target, 'message'):
await target.message.edit_text(text, parse_mode="HTML", reply_markup=keyboard)
else:
await target.answer(text, parse_mode="HTML", reply_markup=keyboard)
# --- Команда /cancel ---
@router.message(Command("cancel"))
async def cmd_cancel(message: Message, state: FSMContext):
"""Отмена текущей операции."""
current_state = await state.get_state()
if current_state is None:
await message.answer("Нет активных операций.")
return
await state.clear()
await message.answer("❌ Операция отменена.")
from bot.handlers_start import get_main_keyboard
await message.answer("📋 Главное меню:", reply_markup=get_main_keyboard(message.from_user.id))
# --- Периодическая очистка ---
async def scheduled_cleanup():
"""Функция для периодической очистки просроченных изображений."""
try:
deleted = await cleanup_expired_images()
if deleted > 0:
logger.info(f"Очистка: удалено {deleted} изображений")
except Exception as e:
logger.error(f"Ошибка при очистке изображений: {e}")