from bot.States import States

KLING_SET_PROMPT = 51
KLING_SET_TEXT_PROMPT = 52


def set_states():
    states_arr = {
        "set_prompt": KLING_SET_PROMPT,
        "set_text_prompt": KLING_SET_TEXT_PROMPT,
    }

    return states_arr


class KlingStates(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