# gradio is a UI library for machine learning models
import gradio as gr

# loguru is a library for logging
from loguru import logger

# generative pre-trained transformer model
from model.model import Model

# load model
model = Model()

# These functions are responsible for defining the chatbot's behavior
#  when the user interacts with the interface. The respond function
#  receives a question and a conversation history. It defines the
#  question in the model (model.question) and calls the
#  question_answerer method to get the answer. The response
#  is added to the history and returned as a result.
def respond(question, history):
    model.question  = question
    history.append((question, model.question_answerer()))
    return "", history

# The set_context function takes a context and sets that context in
#  the model (model.context).
def set_context(context):
    model.context   = context


# In this part, the Gradio interface is created. 
#  the interface has two tabs: "Chat" and "Context".
with gr.Blocks() as interface:
    # In the "Chat" tab, there is a Chatbot component which is
    #  used to display the chatbot conversation. There is also
    #  a Textbox component called prompt_gradio_component
    #  used to receive the question from the user. The
    #  generate_gradio_component button is responsible
    #  for calling the respond function when clicked.
    #  The clear_gradio_component button is used to
    #  clear input fields and conversation.
    with gr.Tab("Chat"):
        chatbot_gradio_component  = gr.Chatbot(label="My Own Chatbot")

        prompt_gradio_component   = gr.Textbox(label="Prompt", lines=2)
        generate_gradio_component = gr.Button("Generate")
        clear_gradio_component    = gr.ClearButton([prompt_gradio_component, chatbot_gradio_component])

        generate_gradio_component.click(respond, [prompt_gradio_component, chatbot_gradio_component], [prompt_gradio_component, chatbot_gradio_component])
    
    # In the "Context" tab, there is a Textbox component called
    #  context_gradio_component used to receive the chatbot
    #  context. The set_context_gradio_component button is
    #  responsible for calling the set_context function
    #  when clicked. The clear_gradio_component button
    #  is used to clear the input field.
    with gr.Tab("Context"):
        context_gradio_component         = gr.Textbox(label="Context", info="your context must be <= 512 tokens!", lines=10)
        set_context_gradio_component     = gr.Button("Set")
        clear_gradio_component           = gr.ClearButton([context_gradio_component])
        
        set_context_gradio_component.click(set_context, [context_gradio_component])

# In this part, the interface is launched and executed. The launch()
#  function is called to launch the Gradio interface.
#  If any errors occur during runtime, they are
#  caught and logged using the loguru library.
if __name__ == "__main__":
    try:
        interface.launch()
    except Exception as error:
        logger.error(error)