# Streamlit App to perform the conversational retrieval using ConversationalResponse class
# 1. Main Title of App
# 2. PDF File Loader
# 3. Streaming Chat Window to ask questions and get answers from ConversationalResponse
# 4. Callback Handler to stream the output of the ConversationalResponse
# 5. Handle the chat interaction with the ConversationalResponse

import streamlit as st
from streamlit_chat import message
from langchain.callbacks.base import BaseCallbackHandler
from src.main import ConversationalResponse
import os

# Constants
ROLE_USER = "user"
ROLE_ASSISTANT = "assistant"

st.set_page_config(page_title="Chat with Documents", page_icon="🦜")
st.title("Chat with PDF Documents 🤖📄")
st.markdown("by [Rohan Kataria](https://www.linkedin.com/in/imrohan/) view more at [VEW.AI](https://vew.ai/)")
#streamlit message block
st.markdown("This app allows you to chat with documents. You can upload a PDF file and ask questions about it. In the backround uses the ConversationalRetrival chain from langchain and Streamlit for UI.")

class StreamHandler(BaseCallbackHandler):
    """
    StreamHandler is a callback handler that streams the output of the ConversationalResponse.
    """
    def __init__(self, container: st.delta_generator.DeltaGenerator, initial_text: str = ""):
        self.container = container
        self.text = initial_text

    def on_llm_new_token(self, token: str, **kwargs) -> None:
        self.text += token
        self.container.markdown(self.text)

@st.cache_resource(ttl="1h")
def load_agent(file_path, api_key):
    """
    Load the ConversationalResponse agent from the given file path.
    """
    with st.spinner('Loading the file...'):
        agent = ConversationalResponse(file_path, api_key)
    st.success("File Loaded Successfully")
    return agent

def handle_chat(agent):
    """
    Handle the chat interaction with the user.
    """
    if "messages" not in st.session_state or st.sidebar.button("Clear message history"):
        st.session_state["messages"] = [{"role": ROLE_ASSISTANT, "content": "How can I help you?"}]

    for msg in st.session_state.messages:
        st.chat_message(msg["role"]).write(msg["content"])

    user_query = st.chat_input(placeholder="Ask me anything!")

    if user_query:
        st.session_state.messages.append({"role": ROLE_USER, "content": user_query})
        st.chat_message(ROLE_USER).write(user_query)

        # Generate the response
        with st.spinner("Generating response"):
            response = agent(user_query)
        
        # Display the response immediately
        st.chat_message(ROLE_ASSISTANT).write(response)
        
        # Add the response to the message history
        st.session_state.messages.append({"role": ROLE_ASSISTANT, "content": response})


def main():
    """
    Main function to handle file upload and chat interaction.
    """

    # API Key Loader
    api_key = st.sidebar.text_input("Enter your OpenAI API Key", type="password")
    if api_key:
        os.environ["OPENAI_API_KEY"] = api_key
    else:
        st.sidebar.error("Please enter your OpenAI API Key.")
        return

    # PDF File Loader to upload the file in the sidebar in session state
    uploaded_file = st.sidebar.file_uploader("Choose a PDF file", type="pdf")
    if uploaded_file is None:
        st.error("Please upload a file.")
        return

    file_details = {"FileName":uploaded_file.name,"FileType":uploaded_file.type,"FileSize":uploaded_file.size}
    st.write(file_details)

    # Create a temp folder
    if not os.path.exists("temp"):
        os.mkdir("temp")
    # Save the file in temp folder
    file_path = os.path.join("temp",uploaded_file.name)
    with open(file_path,"wb") as f:
        f.write(uploaded_file.getbuffer())

    agent = load_agent(file_path, api_key)

    handle_chat(agent)

    # Delete the file from temp folder
    os.remove(file_path)


if __name__ == "__main__":
    main()