from bot.States import States
from telegram.ext import MessageHandler, filters, CallbackQueryHandler

FLUX_SET_PROMPT = 31
FLUX_SET_SUBMODEL = 32
FLUX_SET_TYPE_REQUEST = 33
FLUX_SET_TEXT_PROMPT = 34
FLUX_SET_IMG_TEXT_PROMPT = 35

def set_states():
    states_arr = {
        "set_prompt": FLUX_SET_PROMPT,
        "set_submodel": FLUX_SET_SUBMODEL,
        "set_type_request": FLUX_SET_TYPE_REQUEST,
        "set_text_prompt": FLUX_SET_TEXT_PROMPT,
        "set_img_text_prompt": FLUX_SET_IMG_TEXT_PROMPT,
    }

    return states_arr


class FluxStates(States):

    def __init__(self):
        pass

    @staticmethod
    def get_state_by_key(key: str):
        states = set_states()
        if key in states:
            return states[key]

        return None