# -----------------------------------------
# LIBRARY IMPORTS
# -----------------------------------------

# Streamlit is a library for creating web apps with Python.
import streamlit as st

# langchain.chat_models provides classes related to chat models.
from langchain.chat_models import ChatOpenAI

# langchain.schema provides schemas for system, human, and AI messages.
from langchain.schema import (
    AIMessage,
    HumanMessage,
    SystemMessage
)

# -----------------------------------------
# STREAMLIT UI CONFIGURATION
# -----------------------------------------

# Set the Streamlit page title and icon.
st.set_page_config(page_title="Langchain Prototype", page_icon=":robot:")

# Display a header for the application.
st.header("Hey, I'm your Chat GPT")

# Initialize the session state variable 'sessionMessages' if it doesn't exist.
# This variable will store the conversation.
if "sessionMessages" not in st.session_state:
    st.session_state.sessionMessages = [
        SystemMessage(content="You are a helpful assistant.")
    ]

# -----------------------------------------
# FUNCTION DEFINITIONS
# -----------------------------------------

# This function fetches a response to a given question from the ChatOpenAI model.
def load_answer(question):
    
    # Append the human's message to the session's message list.
    st.session_state.sessionMessages.append(HumanMessage(content=question))
    
    # Get the assistant's response using the chat function.
    assistant_answer = chat(st.session_state.sessionMessages)
    
    # Append the assistant's message to the session's message list.
    st.session_state.sessionMessages.append(AIMessage(content=assistant_answer.content))
    
    # Return the content of the assistant's answer.
    return assistant_answer.content

# Function to get text input from the user via Streamlit's interface.
def get_text():
    input_text = st.text_input("You: ", key=input)
    return input_text

# Initialize the ChatOpenAI model with temperature set to 0.
chat = ChatOpenAI(temperature=0)

# -----------------------------------------
# STREAMLIT INTERACTION HANDLING
# -----------------------------------------

# Get the user input from the Streamlit interface.
user_input = get_text()

# Display a button on the Streamlit page.
submit = st.button('Generate')  

# If the 'Generate' button is clicked:
if submit:
    # Fetch the response from the chat function.
    response = load_answer(user_input)
    
    # Display the response below the subheader "Answer:".
    st.subheader("Answer:")
    st.write(response, key=1)

# -----------------------------------------
# END OF STREAMLIT APPLICATION
# -----------------------------------------