# Memory imports
from components.memory import autochatmemory
from components.memory import formatted_documentretrival

# Prompts engine import
from components.promptengine import promptenginefinalize

#LLM import
from components.LLM import chatopenai, chatopenairetry, chatopenaistream, calculatetokens

# Frontend imports
    ## gradio frontend framework
import gradio as gr
    ## content for frontend
from components.configs import pagetitle, chatlabel, chatfirstmsg, pagesubtitle, inputboxlabel, inputboxplaceholder, restartbuttonlabel, sendbuttonlabel
    ## Examples query 
from components.configs import Fun_examples, Standard_examples, Followup_examples, Languages_examples

#Tools import
    ## Frontend Logger import
from components.frontend_logger import log_interaction_to_database
    ## Session and console logging imports
import datetime
import uuid


#set an empty session id variable, will be updated when a conversation starts
conversation_id= ""
user_message = ""
chatmemory = []
history = []

#Setting the function to be run in the gradio app with every user interaction
def usermessagepush(user_message, history):
    return "", history + [[user_message, None]], user_message
    #return gr.update(value="", interactive=False), history + [[user_message, None]], user_message

#This is what returns the stream while the function is being run. need to be added when calling chatopenaistream(finalprompt, handle_openaistream)
def handle_openaistream(content):
    yield content

def chainexecutor(user_message, history, conversation_id, chatmemory):
    if chatmemory == []:
        conversation_id=str(uuid.uuid4())
        print(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}:Session ID set to: {conversation_id}")
    print(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}:Assisting the user in session id {conversation_id} for the query:\n {user_message}")
    print(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}:State of my memory is\n {chatmemory}")
    #Search relevant content in pinecone long term memory as needed in promptengine
    relevantcontents = formatted_documentretrival(user_message)
    #if the assistant already provided an answer will look for docs about that too
    if chatmemory != []:
        lastresponse= history [-2][1]
        print(lastresponse)
        relevantcontents += formatted_documentretrival(lastresponse, docnumber=1)
    print(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}:Fetched relevant contents:\n{relevantcontents}")
    #Organize the final prompt
    finalprompt = promptenginefinalize(chatmemory,relevantcontents,user_message)
    #Stream the response from openAI
    stream = chatopenaistream(finalprompt, handle_openaistream)
    # This will push the update to gradio frontend every time new content chunk is available
    history[-1][1] = ""
    for content in stream:
        history[-1][1] += ''.join(content)
        yield gr.update(value=""), history, conversation_id, chatmemory
    #without the stream we get assistant response and add it to the history. Here is the other way around since we streamed the response to the frontend:
    #we get the response from the frontend and use it for chatmemory, token calculation and log
    assistant_response = history[-1][1]
    usedtokens=calculatetokens(finalprompt,assistant_response)
    print(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}:Generated a response: \n {assistant_response} \n estimated spending: {usedtokens} tokens")
    # update chatmemory
    chatmemory = autochatmemory(chatmemory,relevantcontents,user_message, assistant_response)
    #log the activity
    log_interaction_to_database(user_message, response=assistant_response, conversation_id=conversation_id, usedtokens=usedtokens)
    return gr.update(value=""), history, conversation_id, chatmemory

# Front end web app
block = gr.Blocks(theme=gr.themes.Soft(spacing_size="sm",radius_size="none"))
with block as DrusoGPTUI:
    conversation_id = gr.State(conversation_id)
    chatmemory = gr.State(chatmemory)
    user_message = gr.State(user_message)
    with gr.Row():
        gr.Markdown(pagetitle)
    with gr.Row():
        gr.HTML(pagesubtitle)
    with gr.Row():
        chatbot = gr.Chatbot(
            label=chatlabel,
            show_label=True,
            value=[(None, chatfirstmsg)]
            ).style(height=280)
    with gr.Row():
        msg = gr.Textbox(
            label=inputboxlabel,
            placeholder=inputboxplaceholder,
            lines=1,
        )
    with gr.Row():
        submit = gr.Button(value=sendbuttonlabel, variant="primary")
    with gr.Row():
        with gr.Tab("Standard"):
            gr.Examples(Standard_examples, msg, cache_examples=False)
        with gr.Tab("Fun"):
            gr.Examples(Fun_examples, msg, cache_examples=False)
        with gr.Tab("Followup"):
            gr.Examples(Followup_examples, msg, cache_examples=False)
        with gr.Tab("Languages"):
            gr.Examples(Languages_examples, msg, cache_examples=False)
        #with gr.Tab("automatic"):
            # automatically run approach, does not work now gr.Examples(["Do you have any products for psoriasis?"], msg, fn=chainexecutor, run_on_click=True, cache_examples=False)
    with gr.Row():
        clear = gr.Button(restartbuttonlabel, variant="secondary")
    msg.submit(usermessagepush, [msg, chatbot], [msg, chatbot, user_message], queue=False).then(chainexecutor, [user_message, chatbot, conversation_id, chatmemory], [msg, chatbot, conversation_id, chatmemory], queue=True)
    clear.click(lambda: None, None, chatbot, queue=False)
    submit.click(usermessagepush, [msg, chatbot], [msg, chatbot, user_message], queue=False).then(chainexecutor, [user_message, chatbot, conversation_id, chatmemory], [msg, chatbot, conversation_id, chatmemory], queue=True)
if __name__ == "__main__":
    DrusoGPTUI.queue() #added to test streaming
    DrusoGPTUI.launch(debug=True)