from ai_bot.keyboards.CreateImageKeyboard import CreateImageKeyboard
from ai_bot.msg.CreateImageMsg import CreateImageMsg
from bot.Handler import Handler
from bot.halpers.Instance import Instance
from bot import SECOND
from telegram.constants import ParseMode


class ImageHandler(Handler):

    @staticmethod
    async def create_img(update, context):
        models_arr = ['Midjourney', 'Flux']

        reply_markup = CreateImageKeyboard()
        reply_markup.add_option("models_arr", models_arr)

        msg = CreateImageMsg()

        await context.bot.send_message(chat_id=update.effective_chat.id, text=msg.get_msg(),
                                       parse_mode=ParseMode.HTML, reply_markup=reply_markup.create_keyboard())

        return SECOND

    @staticmethod
    async def generate_img(update, context):
        query = update.callback_query
        command, params = Handler.load_callback_query(query.data)

    @staticmethod
    async def set_prompt(update, context):
        query = update.callback_query
        command, params = Handler.load_callback_query(query.data)

        model = Instance.get_instance("ai_bot.states.{model}States".format(model=params['model']), "{model}States".format(model=params['model']))
        state = model.get_state_by_key("set_prompt")

        await context.bot.send_message(chat_id=update.effective_chat.id, text="Напишите запрос.",
                                       parse_mode=ParseMode.HTML)

        return state

    @staticmethod
    async def start_model_scenario(update, context):
        query = update.callback_query
        command, params = Handler.load_callback_query(query.data)

        model = Instance.get_instance("ai_bot.handlers.{model}Handler".format(model=params['model']),
                                      "{model}Handler".format(model=params['model']))

        state = await model.start_scenario(update, context)

        return state