import os
import streamlit as st
from dotenv import load_dotenv
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.llms import llamacpp
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.callbacks import CallbackManager, StreamingStdOutCallbackHandler
from langchain.chains import create_history_aware_retriever, create_retrieval_chain, ConversationalRetrievalChain
from langchain.document_loaders import TextLoader
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_community.chat_message_histories.streamlit import StreamlitChatMessageHistory
from langchain.prompts import PromptTemplate
from langchain.vectorstores import Chroma
from utills import load_txt_documents, split_docs, load_uploaded_documents, retriever_from_chroma
from langchain.text_splitter import TokenTextSplitter, RecursiveCharacterTextSplitter
from langchain_community.document_loaders.directory import DirectoryLoader
from HTML_templates import css, bot_template, user_template




def retriever_from_chroma(docs, search_type, k):
    model_name = "sentence-transformers/all-mpnet-base-v2"
    model_kwargs = {'device': 'cpu'}
    encode_kwargs = {'normalize_embeddings': True}
    embeddings = HuggingFaceEmbeddings(
        model_name=model_name,
        model_kwargs=model_kwargs,
        encode_kwargs=encode_kwargs
    )
    vectorstore_path = "docs/chroma/"
    if not os.path.exists(vectorstore_path):
        os.makedirs(vectorstore_path)
    vectorstore = Chroma.from_documents(
        documents=text_chunks, embedding=embeddings, persist_directory="docs/chroma/")
    retriever = vectordb.as_retriever(search_type=search_type, search_kwargs={"k": k})
    return retriever
        



data_path = "data"


    
documents = []
    
for filename in os.listdir(data_path):

    if filename.endswith('.txt'):

        file_path = os.path.join(data_path, filename)

        documents = TextLoader(file_path).load()

        documents.extend(documents)
    



docs = split_docs(documents, 250, 20)

retriever = retriever_from_chroma(docs,'mmr',7)






def main(retriever):

    st.set_page_config(page_title="Chat with multiple PDFs",
                       page_icon=":books:")
    st.write(css, unsafe_allow_html=True)

    if "conversation" not in st.session_state:
        st.session_state.conversation = None
    if "chat_history" not in st.session_state:
        st.session_state.chat_history = None

    st.header("Chat with multiple PDFs :books:")

    with st.chat_message("Assistant"):
        st.write("Hello my name is Robert, how can i help you? ")
    user_question = st.text_input("Ask a question about your documents:")
    with st.chat_message("User"):
        st.write(user_question)
    if user_question:
        handle_userinput(user_question,vectorstore)


def handle_userinput(user_question,retriever):
    docs = retriever.invoke(question)
    
    doc_txt = [doc.page_content for doc in docs]
    
    Rag_chain = create_conversational_rag_chain(retriever)
    response = rag_chain.invoke({"context": doc_txt, "question": question})
    with st.chat_message("Assistant"):
        st.write(response)


def create_conversational_rag_chain(retriever):
    
    model_path = ('qwen2-0_5b-instruct-q4_0.gguf')

    callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])

    llm = llamacpp.LlamaCpp(
        model_path=model_path,
        n_gpu_layers=1,
        temperature=0.1,
        top_p=0.9,
        n_ctx=22000,
        max_tokens=200,
        repeat_penalty=1.7,
        # callback_manager=callback_manager,
        verbose=False,
    )

    template = """Answer the question based only on the following context:
    {context}

    Question: {question}
    """
    prompt = ChatPromptTemplate.from_template(template)

    rag_chain = prompt | llm | StrOutputParser()


    return rag_chain


 

if __name__ == "__main__":
    main(vectorstore)