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.extra_functions import calculator, execute_code, none
from llm_modules.llm import llm_usage
from rw_modules.xtts_voice import xtts_model
# be sure to change the model paths for everything!
xtts = xtts_model()
llm_model = llm_usage()
sdxl_model = sdxl_lightning()
moondream_model = moondream()
search_engine = search()
audio_gen = audio_generation()

function_dict = {
    "robot": none,
    "image_gen": sdxl_model.infer,
    "code": execute_code,
    "music_gen": audio_gen.infer,
    "search": search_engine.search_text,
    "search_img": search_engine.search_image,
    "save": none,
    "plot_graph": plot_equation,
    "math": calculator,
    "change_volume": none,
    "current_news": search_engine.search_news,
    "weather": search_engine.search_text,
    "none": none,
    "listen": none,
    "screen": none,
    "VisualQ": moondream_model.infer
}

class FunctionExecutor:
    def __init__(self):
        pass

    def split_variables(self, h, prefixes_to_remove):
        for prefix in prefixes_to_remove:
            h = h.replace(prefix, '')
            if "'" in h:
                variables = [item.strip() for item in h.split("',")]
            elif '"' in h:
                variables = [item.strip() for item in h.split('",')]
            else:
                variables = [item.strip() for item in h.split(',')]
            for vars in variables:
                result = vars[1:-1] if vars.startswith('"') and vars.endswith('"') else vars
                result = vars[1:-1] if vars.startswith("'") and vars.endswith("'") else vars
                full_results.append(result)
        return full_results

    def execute(self, function, args):
        try:
            return function(*args)
        except (ImportError, AttributeError):
            print("Error executing the function.")

    def get_function(self, action, param):
        chosen_function = function_dict.get(action)
        if chosen_function:
            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(chosen_function, params_list)
            return result
        else:
            print("Function not found.")

function_executor = FunctionExecutor()

prompt = "Hello, how are you today? Please generate a nice image of a man."
personality = "a caring, nice assistant named hermes"
name = "Hermes"
previous_prompts = []
previous_answers = []
final_answer = False

def llm_speech(prompt, functions, function_executor, personality, name, previous_prompts, previous_answers):
    speech_output = ""
    for out in llm_model.process_llm(prompt, functions, function_executor=function_executor, personality=personality, name=name, previous_prompts=previous_prompts, previous_answers=previous_answers):
        token = out['token']
        if token:
            #print(token, end="")
            yield {'token': token, 'speech': None}
            if "answer:" in full_output:
                speech_output += token
                symbols = ['!', '.', '?']
                if any(speech_output.endswith(symbol) for symbol in symbols):
                    for speech in xtts.infer(speech_output):
                        yield {'speech': speech, 'token': None}
                    speech_output = ""
                        
for output in llm_speech(prompt, functions, function_executor=function_executor, personality=personality, name=name, previous_prompts=previous_prompts, previous_answers=previous_answers)
    if output['token']:
        print(token, end="")
    else:
        print("only in ohio")