From b83e649349472dd22665a971459d985644490804 Mon Sep 17 00:00:00 2001
From: BarsTiger <kit.barsic@gmail.com>
Date: Wed, 8 Mar 2023 16:10:22 +0200
Subject: [PATCH] Added `config` and `current` commands, `Close` button for
 inline menus, `pass` to all `__init__`s

---
 bot/callbacks/__init__.py                 |  1 +
 bot/callbacks/common.py                   | 15 ++++++++++++++
 bot/callbacks/config/__init__.py          |  1 +
 bot/callbacks/factories/__init__.py       |  1 +
 bot/callbacks/factories/common.py         |  4 ++++
 bot/callbacks/factories/config.py         |  6 ++++++
 bot/callbacks/image_info.py               | 14 ++++++++++---
 bot/callbacks/register.py                 |  4 +++-
 bot/handlers/config/__init__.py           |  6 ++++++
 bot/handlers/config/config_command.py     | 11 +++++++++++
 bot/handlers/help_command/help_strings.py |  2 ++
 bot/handlers/register.py                  |  3 ++-
 bot/handlers/txt2img/__init__.py          |  2 ++
 bot/handlers/txt2img/current.py           | 24 +++++++++++++++++++++++
 bot/keyboards/__init__.py                 |  1 +
 bot/keyboards/config.py                   | 22 +++++++++++++++++++++
 bot/keyboards/set_model.py                |  5 +++++
 bot/modules/api/objects/__init__.py       |  1 +
 bot/modules/get_hash/__init__.py          |  1 +
 19 files changed, 119 insertions(+), 5 deletions(-)
 create mode 100644 bot/callbacks/common.py
 create mode 100644 bot/callbacks/config/__init__.py
 create mode 100644 bot/callbacks/factories/common.py
 create mode 100644 bot/callbacks/factories/config.py
 create mode 100644 bot/handlers/config/__init__.py
 create mode 100644 bot/handlers/config/config_command.py
 create mode 100644 bot/handlers/txt2img/current.py
 create mode 100644 bot/keyboards/config.py

diff --git a/bot/callbacks/__init__.py b/bot/callbacks/__init__.py
index e69de29..2ae2839 100644
--- a/bot/callbacks/__init__.py
+++ b/bot/callbacks/__init__.py
@@ -0,0 +1 @@
+pass
diff --git a/bot/callbacks/common.py b/bot/callbacks/common.py
new file mode 100644
index 0000000..9fba9e1
--- /dev/null
+++ b/bot/callbacks/common.py
@@ -0,0 +1,15 @@
+from bot.common import dp
+from bot.callbacks.factories.common import close_keyboard_data
+from aiogram import types
+from bot.utils.private_keyboard import other_user
+
+
+async def on_close_keyboard(call: types.CallbackQuery):
+    if await other_user(call):
+        return
+
+    await call.message.delete()
+
+
+def register():
+    dp.register_callback_query_handler(on_close_keyboard, close_keyboard_data.filter())
diff --git a/bot/callbacks/config/__init__.py b/bot/callbacks/config/__init__.py
new file mode 100644
index 0000000..2ae2839
--- /dev/null
+++ b/bot/callbacks/config/__init__.py
@@ -0,0 +1 @@
+pass
diff --git a/bot/callbacks/factories/__init__.py b/bot/callbacks/factories/__init__.py
index e69de29..2ae2839 100644
--- a/bot/callbacks/factories/__init__.py
+++ b/bot/callbacks/factories/__init__.py
@@ -0,0 +1 @@
+pass
diff --git a/bot/callbacks/factories/common.py b/bot/callbacks/factories/common.py
new file mode 100644
index 0000000..d399fcb
--- /dev/null
+++ b/bot/callbacks/factories/common.py
@@ -0,0 +1,4 @@
+from aiogram.utils.callback_data import CallbackData
+
+
+close_keyboard_data = CallbackData("close_keyboard")
diff --git a/bot/callbacks/factories/config.py b/bot/callbacks/factories/config.py
new file mode 100644
index 0000000..dd8014e
--- /dev/null
+++ b/bot/callbacks/factories/config.py
@@ -0,0 +1,6 @@
+from aiogram.utils.callback_data import CallbackData
+
+
+prompt_settings_data = CallbackData("prompt_settings")
+global_settings_data = CallbackData("global_settings")
+admin_settings_data = CallbackData("admin_settings")
diff --git a/bot/callbacks/image_info.py b/bot/callbacks/image_info.py
index ae0c317..1865a8e 100644
--- a/bot/callbacks/image_info.py
+++ b/bot/callbacks/image_info.py
@@ -6,8 +6,10 @@ from bot.keyboards.image_info import get_img_info_keyboard, get_img_back_keyboar
 from bot.utils.cooldown import throttle
 from bot.utils.private_keyboard import other_user
 from bot.modules.api.objects.prompt_request import Generated
+from bot.utils.errorable_command import wrap_exception
 
 
+@wrap_exception()
 async def on_back(call: types.CallbackQuery, callback_data: dict):
     p_id = callback_data['p_id']
     if await other_user(call):
@@ -20,6 +22,7 @@ async def on_back(call: types.CallbackQuery, callback_data: dict):
     )
 
 
+@wrap_exception()
 @throttle(5)
 async def on_prompt_only(call: types.CallbackQuery, callback_data: dict):
     p_id = callback_data['p_id']
@@ -36,6 +39,7 @@ async def on_prompt_only(call: types.CallbackQuery, callback_data: dict):
     )
 
 
+@wrap_exception()
 @throttle(5)
 async def on_full_info(call: types.CallbackQuery, callback_data: dict):
     p_id = callback_data['p_id']
@@ -59,17 +63,21 @@ async def on_full_info(call: types.CallbackQuery, callback_data: dict):
     )
 
 
+@wrap_exception()
 @throttle(5)
 async def on_import(call: types.CallbackQuery, callback_data: dict):
     p_id = callback_data['p_id']
     if await other_user(call):
         return
 
-    prompt: Generated = db[DBTables.generated].get(p_id)
+    generated: Generated = db[DBTables.generated].get(p_id)
+    prompt = generated.prompt
+    prompt.creator = call.from_user.id
+    db[DBTables.prompts][call.from_user.id] = prompt
+    await db[DBTables.config].write()
 
     await call.message.edit_text(
-        f"😶‍🌫️ Not implemented yet",
-        parse_mode='html',
+        f"️🥐 Prompt imported",
         reply_markup=get_img_back_keyboard(p_id)
     )
 
diff --git a/bot/callbacks/register.py b/bot/callbacks/register.py
index 11850dc..d5169c6 100644
--- a/bot/callbacks/register.py
+++ b/bot/callbacks/register.py
@@ -5,11 +5,13 @@ def register_callbacks():
     from bot.callbacks import (
         exception,
         image_info,
-        set_model
+        set_model,
+        common
     )
 
     exception.register()
     image_info.register()
     set_model.register()
+    common.register()
 
     print('[gray]All callbacks registered[/]')
diff --git a/bot/handlers/config/__init__.py b/bot/handlers/config/__init__.py
new file mode 100644
index 0000000..2c59a40
--- /dev/null
+++ b/bot/handlers/config/__init__.py
@@ -0,0 +1,6 @@
+from bot.common import dp
+from .config_command import *
+
+
+def register():
+    dp.register_message_handler(config_command, commands='config')
diff --git a/bot/handlers/config/config_command.py b/bot/handlers/config/config_command.py
new file mode 100644
index 0000000..f2aa0fc
--- /dev/null
+++ b/bot/handlers/config/config_command.py
@@ -0,0 +1,11 @@
+from aiogram import types
+from bot.db import db, DBTables
+from bot.utils.cooldown import throttle
+from bot.utils.errorable_command import wrap_exception
+from bot.keyboards.config import get_config_keyboard
+
+
+@wrap_exception()
+@throttle(cooldown=60*60, admin_ids=db[DBTables.config].get('admins'), by_id=False)
+async def config_command(message: types.Message):
+    await message.reply("⚙️ Configuration:", reply_markup=get_config_keyboard(message.from_id))
diff --git a/bot/handlers/help_command/help_strings.py b/bot/handlers/help_command/help_strings.py
index 2218f39..686cf72 100644
--- a/bot/handlers/help_command/help_strings.py
+++ b/bot/handlers/help_command/help_strings.py
@@ -2,6 +2,7 @@ help_data = {
     'generate': 'Generate picture using configuration set by user. You can pass prompt also in command arguments or '
                 'use it without arguments to generate picture with prompt, that was used for last generation',
     'imginfo': 'Get information about image, that was generated using this bot',
+    'current': 'Get your current prompt',
     'setprompt': 'Set default prompt for images, will be overwritten if you specify prompt in generate command',
     'setnegative': 'Set negative prompt',
     'setsize': 'Set size for image (in hxw format)',
@@ -12,6 +13,7 @@ help_data = {
     'setscale': 'Set CFG Scale (prompt stringency)',
     'setfaces': 'Set restore faces mode',
     'status': 'Ping API endpoint host',
+    'config': 'Edit config using inline buttons',
     'setmodel': '(global) Sets StableDiffusion model for all users. Can be used only once an hour',
     'setendpoint': '(admin) Set StableDiffusion API endpoint',
     'addadmin': '(admin) Add new admin - reply to message or type user ID',
diff --git a/bot/handlers/register.py b/bot/handlers/register.py
index 75ded6d..de3a159 100644
--- a/bot/handlers/register.py
+++ b/bot/handlers/register.py
@@ -3,7 +3,7 @@ from rich import print
 
 def register_handlers():
     from bot.handlers import (
-        initialize, admin, help_command, txt2img, image_info
+        initialize, admin, help_command, txt2img, image_info, config
     )
 
     initialize.register()
@@ -11,5 +11,6 @@ def register_handlers():
     help_command.register()
     txt2img.register()
     image_info.register()
+    config.register()
 
     print('[gray]All handlers registered[/]')
diff --git a/bot/handlers/txt2img/__init__.py b/bot/handlers/txt2img/__init__.py
index 081769e..414a54c 100644
--- a/bot/handlers/txt2img/__init__.py
+++ b/bot/handlers/txt2img/__init__.py
@@ -2,6 +2,7 @@ from bot.common import dp
 from .txt2img import generate_command
 from .set_model import set_model_command
 from .status import get_status
+from .current import get_current
 from .set_settings import (
     set_height_command, set_negative_prompt_command, set_size_command, set_steps_command, set_width_command,
     set_prompt_command, set_sampler_command, set_cfg_scale_command, set_restore_faces_command
@@ -21,3 +22,4 @@ def register():
     dp.register_message_handler(set_settings.set_restore_faces_command, commands='setfaces')
     dp.register_message_handler(set_model.set_model_command, commands='setmodel')
     dp.register_message_handler(status.get_status, commands='status')
+    dp.register_message_handler(current.get_current, commands='current')
diff --git a/bot/handlers/txt2img/current.py b/bot/handlers/txt2img/current.py
new file mode 100644
index 0000000..d726a11
--- /dev/null
+++ b/bot/handlers/txt2img/current.py
@@ -0,0 +1,24 @@
+from aiogram import types
+from bot.db import db, DBTables
+from bot.utils.cooldown import throttle
+from bot.modules.api.objects.prompt_request import Prompt
+from bot.utils.errorable_command import wrap_exception
+
+
+@wrap_exception()
+@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins'))
+async def get_current(message: types.Message):
+    prompt: Prompt = db[DBTables.prompts].get(message.from_id)
+    if prompt is None:
+        await message.reply('❌ No prompts found for you')
+
+    await message.reply(
+        f"🖤 Prompt: {prompt.prompt} \n"
+        f"🐊 Negative: {prompt.negative_prompt} \n"
+        f"🪜 Steps: {prompt.steps} \n"
+        f"🧑‍🎨 CFG Scale: {prompt.cfg_scale} \n"
+        f"🖥️ Size: {prompt.width}x{prompt.height} \n"
+        f"😀 Restore faces: {'on' if prompt.restore_faces else 'off'} \n"
+        f"⚒️ Sampler: {prompt.sampler} \n",
+        parse_mode='html'
+    )
diff --git a/bot/keyboards/__init__.py b/bot/keyboards/__init__.py
index e69de29..2ae2839 100644
--- a/bot/keyboards/__init__.py
+++ b/bot/keyboards/__init__.py
@@ -0,0 +1 @@
+pass
diff --git a/bot/keyboards/config.py b/bot/keyboards/config.py
new file mode 100644
index 0000000..eaf3da9
--- /dev/null
+++ b/bot/keyboards/config.py
@@ -0,0 +1,22 @@
+from aiogram import types
+from bot.db import db, DBTables
+from bot.callbacks.factories.config import (prompt_settings_data, global_settings_data, admin_settings_data)
+from bot.callbacks.factories.common import close_keyboard_data
+
+
+def get_config_keyboard(user_id: int) -> types.InlineKeyboardMarkup:
+    buttons = [
+        types.InlineKeyboardButton("Prompt settings", callback_data=prompt_settings_data.new()),
+        types.InlineKeyboardButton("Global settings", callback_data=global_settings_data.new())
+    ]
+    if user_id in db[DBTables.config].get('admins'):
+        buttons.append(
+            types.InlineKeyboardButton("Admin settings", callback_data=admin_settings_data.new())
+        )
+
+    buttons.append(
+        types.InlineKeyboardButton("🔻 Close", callback_data=close_keyboard_data.new())
+    )
+    keyboard = types.InlineKeyboardMarkup(row_width=1)
+    keyboard.add(*buttons)
+    return keyboard
diff --git a/bot/keyboards/set_model.py b/bot/keyboards/set_model.py
index 8db0b89..e43a355 100644
--- a/bot/keyboards/set_model.py
+++ b/bot/keyboards/set_model.py
@@ -1,6 +1,7 @@
 from aiogram import types
 from bot.db import db, DBTables
 from bot.callbacks.factories.set_model import set_model_page, set_model
+from bot.callbacks.factories.common import close_keyboard_data
 
 
 def get_set_model_keyboard(page: int) -> types.InlineKeyboardMarkup:
@@ -28,4 +29,8 @@ def get_set_model_keyboard(page: int) -> types.InlineKeyboardMarkup:
     else:
         keyboard.add(*models_buttons)
 
+    keyboard.add(
+        types.InlineKeyboardButton("🔻 Cancel", callback_data=close_keyboard_data.new())
+    )
+
     return keyboard
diff --git a/bot/modules/api/objects/__init__.py b/bot/modules/api/objects/__init__.py
index e69de29..2ae2839 100644
--- a/bot/modules/api/objects/__init__.py
+++ b/bot/modules/api/objects/__init__.py
@@ -0,0 +1 @@
+pass
diff --git a/bot/modules/get_hash/__init__.py b/bot/modules/get_hash/__init__.py
index e69de29..2ae2839 100644
--- a/bot/modules/get_hash/__init__.py
+++ b/bot/modules/get_hash/__init__.py
@@ -0,0 +1 @@
+pass