from langchain import LLMChain
from langchain.prompts import PromptTemplate
from omegaconf import OmegaConf
import datetime

MODEL_SELECTION_PROMPT_TEMPLATE = """
Select one of the following models based on the given concept.
You must choose one model name based on the description of each model and the concept!

Cencept: {concept}

Model name and description: {model_list}

Warning: {warning}

The avilable model names: 
{model_name_list}

Selected model name:
"""

def load_model_list():
    models_config = OmegaConf.load('resources/models.yaml')
    models_dict = models_config['models']
    model_name_list_str = ''
    print(models_dict)
    model_list_str = ''
    for key, value in models_dict.items():
        model_list_str+="model name: " +key+', model description: '+value['desc']+'\n'
        model_name_list_str += key + ' '
    model_name_list_str += '\n'
    return model_list_str, models_dict, model_name_list_str

def model_selection_chain(llm, class_concept=None):
    chain = None
    memory = None
    if llm:
        print("class_concept", class_concept)
        if class_concept is None:
            class_concept = 'AI assistant'


        template = PromptTemplate(
            input_variables=["model_list", "concept", "warning", "model_name_list"],
            template=MODEL_SELECTION_PROMPT_TEMPLATE,
        )
        model_list_str, models_dict, model_name_list_str = load_model_list()

        personality_chain = LLMChain(
            llm=llm, prompt=template, verbose=True)
        selected_model = None
        while (selected_model is None) or not (selected_model in models_dict):
            if (selected_model is not None) and not (selected_model in models_dict):
                warning_str = '{} is not in Model list! \n'.format(selected_model)
            else:
                warning_str = ''
            selected_model = personality_chain.run({'concept': class_concept, 'model_list':model_list_str, 'warning': warning_str, 'model_name_list': model_name_list_str})
        print("Selected model name: ", selected_model)

    return models_dict[selected_model]