from modules.lightning import sdxl_lightning
from modules.audio_ldm import audio_generation
from modules.graph_plotter import plot_equation
from modules.search_web import search
from modules.vqa import moondream
from modules.none_function import none 
from modules.calculator import math_solver
function_dict = {
    "robot": "none",
    "image_gen": "sdxl_model.infer",
    "code": "none",
    "music_gen": "audio_gen.infer",
    "search": "search_engine.search_text",
    "search_img": "search_engine.search_image",
    "save": "none",
    "plot_graph": "plot_graph",
    "math": "math_solver",
    "change_volume": "none",
    "current_news": "search_engine.search_news",
    "weather": "search_engine.search_text",
    "none": "none",
    "listen": "none",
    "screen": "none",
    "VisualQ": "moondream_model.infer"
}

sdxl_model = sdxl_lightning()
plot_graph = plot_equation()
moondream_model = moondream()
search_engine = search()
audio_gen = audio_generation()

class get_functions:
    def __init__(self):
        print("Initializing get functions")

    def split_variables(self, h, prefixes_to_remove):
    # Remove specified prefixes from the string
        for prefix in prefixes_to_remove:
            h = h.replace(prefix, '')
        h_split = h.split(',')
        variables = [item.strip() for item in h_split]
        return variables
        
    def execute_function(self, func_str, args_list=None, **kwargs):
        try:
            module_name, func_name = func_str.rsplit('.', 1)
            module = importlib.import_module(module_name)
            func = getattr(module, func_name)
            if args_list is not None:
                result = func(*args_list, **kwargs)
            else:
                result = func(**kwargs)
            return str(result)
        except (AttributeError, ImportError) as e:
            print(f"Error: {e}")
            
    def get_function(self, action, param):
        chosen_function = function_dict[str(action)]
        function_args = ["num_images=", "description=", "prompt=", "code=", "length=", "query=", "input=", "graph=", "equation=", "volume=", "place=", "percentage="]
        params_list = self.split_variables(str(param), function_args)
        result = self.execute_function(str(chosen_function), arg_list=params_list)
        return str(result['llm_output'])