import streamlit as st
from rag_components import load_documents, split_documents, create_embeddings, setup_vector_store, create_qa_chain
import os

st.set_page_config(page_title="Document Chatbot")
st.title("Chat with your Documents")

@st.cache_resource
def initialize_rag_components(file_path="src/me.txt"):
    """Initializes and caches RAG components."""
    if not os.path.exists(file_path):
        st.error(f"Error: Document file not found at {file_path}")
        return None, None

    documents = load_documents(file_path)
    docs = split_documents(documents)
    embeddings = create_embeddings()
    retriever = setup_vector_store(docs, embeddings)
    qa_chain = create_qa_chain(retriever)
    return qa_chain, retriever

qa_chain, retriever = initialize_rag_components()

if qa_chain is not None:
    # Initialize chat history
    if "messages" not in st.session_state:
        st.session_state.messages = []

    # Display chat messages from history on app rerun
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    # React to user input
    if prompt := st.chat_input("Ask me a question about the document"):
        # Display user message in chat message container
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)

        # Display assistant response in chat message container
        with st.chat_message("assistant"):
            message_placeholder = st.empty()
            full_response = ""
            try:
                # Assuming qa_chain.stream() yields dictionaries with a 'result' key
                for chunk in qa_chain.stream(prompt):
                    if 'result' in chunk:
                        full_response += chunk['result']
                        message_placeholder.markdown(full_response + "▌")
                message_placeholder.markdown(full_response)
            except Exception as e:
                st.error(f"An error occurred: {e}")
                full_response = "Sorry, I could not process your request."

        # Add assistant response to chat history
        st.session_state.messages.append({"role": "assistant", "content": full_response})
else:
    st.warning("RAG components could not be initialized. Please check the document file path.")