import gradio as gr
from gpt4all import GPT4All
from urllib.request import urlopen
import json
import time

url = "https://raw.githubusercontent.com/nomic-ai/gpt4all/main/gpt4all-chat/metadata/models3.json"
response = urlopen(url)
data_json = json.loads(response.read())

def model_choices():
    model_list = [data_json[i]['filename'] for i in range(len(data_json))]
    return model_list

model_description = {model['filename']: model['description'] for model in data_json}


def llm_intro(selected_model):
    return model_description.get(selected_model, "No description available for this model selection.")

model_cache = {}  # Global cache


def load_model(model_name):
    """
    This function checks the cache before loading a model.
    If the model is cached, it returns the cached version.
    Otherwise, it loads the model, caches it, and then returns it.
    """
    if model_name not in model_cache:
        model = GPT4All(model_name)  
        model_cache[model_name] = model
    return model_cache[model_name]


def generate_text(input_text, selected_model):
    """
    Generate text using the selected model.
    This function now uses the caching mechanism to load models.
    """
    model = load_model_with_cache(selected_model)
    output = model.generate(input_text, max_tokens=100)
    return output


# with gr.Blocks() as demo:
#     gr.Markdown("## GPT4All Text Generation Experiment")

#     with gr.Row():
        
    
#         model_selection = gr.Dropdown(choices=model_choices(), 
#                                      multiselect=False, 
#                                      label="LLMs to choose from", 
#                                      type="value", 
#                                      value="orca-mini-3b-gguf2-q4_0.gguf")
    
#         explanation = gr.Textbox(label="Model Description", lines=3, interactive=False, value=llm_intro("orca-mini-3b-gguf2-q4_0.gguf"))
    
#         # Link the dropdown with the textbox to update the description based on the selected model
#         model_selection.change(fn=llm_intro, inputs=model_selection, outputs=explanation)

#     chatbot = gr.Chatbot()
#     input_text = gr.Textbox(lines=3, label="Press shift+Enter to submit")
#     # output_text = gr.Textbox(lines=10, label="Generated Text")
#     clear = gr.ClearButton([input_text, chatbot])


        

# Define the chatbot function
def generate_response(model_name, message, chat_history):
    model = load_model(model_name)
    chat_history = []
    if len(chat_history) > 0:
        past_chat = ", ".join(chat_history)
        input_text = past_chat + " " + message
    else:
        input_text = message
    response = model.generate(input_text, max_tokens=100)
    chat_history.append((input_text, response))
    return chat_history, response

# Create the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# GPT4All Chatbot")
    with gr.Row():
        with gr.Column(scale=1):
            model_dropdown = gr.Dropdown(
                choices=model_choices(),  
                multiselect=False,
                type="value",
                value="orca-mini-3b-gguf2-q4_0.gguf",
                label="LLMs to choose from"
            )
            explanation = gr.Textbox(label="Model Description", lines=3, interactive=False, value=llm_intro("orca-mini-3b-gguf2-q4_0.gguf"))
    
            # Link the dropdown with the textbox to update the description based on the selected model
            model_dropdown.change(fn=llm_intro, inputs=model_dropdown, outputs=explanation)
            
        with gr.Column(scale=4):
            chatbot = gr.Chatbot(label="Conversation", value=[(None, "How may I help you today?")])
            
    message = gr.Textbox(label="Message")
    state = gr.State()

    message.submit(generate_response, inputs=[model_dropdown, message, state], outputs=[chatbot, state])

# Launch the Gradio app
demo.launch()