import os
import streamlit as st
from dotenv import load_dotenv
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.llms import llamacpp, 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 langchain.text_splitter import TokenTextSplitter, RecursiveCharacterTextSplitter
from langchain_community.document_loaders.directory import DirectoryLoader
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_experimental.chat_models import Llama2Chat
from langchain_community.chat_models.llamacpp import ChatLlamaCpp

lang_api_key = os.getenv("lang_api_key")

os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_ENDPOINT"] = "https://api.langchain.plus"
os.environ["LANGCHAIN_API_KEY"] = lang_api_key
os.environ["LANGCHAIN_PROJECT"] = "Lithuanian_Law_RAG_QA"





def create_retriever_from_chroma(vectorstore_path="./docs/chroma/", search_type='mmr', k=7, chunk_size=300, chunk_overlap=30,lambda_mult= 0.7):
    
    model_name = "Alibaba-NLP/gte-large-en-v1.5"
    model_kwargs = {'device': 'cpu',
                   "trust_remote_code" : 'False'}
    encode_kwargs = {'normalize_embeddings': True}
    embeddings = HuggingFaceEmbeddings(
        model_name=model_name,
        model_kwargs=model_kwargs,
        encode_kwargs=encode_kwargs
    )

    

    if os.path.exists(vectorstore_path) and os.listdir(vectorstore_path):
        vectorstore = Chroma(persist_directory=vectorstore_path,embedding_function=embeddings)
        
    else:
        st.write("Vector store doesnt exist and will be created now")
        loader = DirectoryLoader('./data/', glob="./*.txt", loader_cls=TextLoader)
        docs = loader.load()
        
        
        text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
        chunk_size=chunk_size, chunk_overlap=chunk_overlap,
        separators=["\n\n \n\n","\n\n\n", "\n\n", r"In \[[0-9]+\]", r"\n+", r"\s+"],
        is_separator_regex = True
    )
        split_docs = text_splitter.split_documents(docs)

        
        vectorstore = Chroma.from_documents(
            documents=split_docs, embedding=embeddings, persist_directory=vectorstore_path
        )
        
    
    retriever=vectorstore.as_retriever(search_type = search_type, search_kwargs={"k": k})

    
    

    return retriever






def main():


    

    st.set_page_config(page_title="Chat with multiple Lithuanian Law Documents: ",
                       page_icon=":books:")
    

    st.header("Chat with multiple Lithuanian Law Documents:" ":books:")
    
    st.markdown("###### Hi, I am Birute (Powered by gemma-2-2b-it-Q8 model), chat assistant, based on republic of Lithuania law documents. You can choose below information retrieval type and how many documents you want to be retrieved.")
    st.markdown("Available Documents: LR_Civil_Code_2022, LR_Constitution_2022, LR_Criminal_Code_2018, LR_Criminal_Procedure_code_2022,LR_Labour_code_2010. P.S it's a shame that there are no newest documents translations into English... ")

    if "messages" not in st.session_state:
        st.session_state["messages"] = [
        {"role": "assistant", "content": "Hi, I'm a chatbot who is  based on respublic of Lithuania law documents. How can I help you?"}
    ]


    search_type = st.selectbox(
        "Choose search type. Options are [Max marginal relevance search (similarity) , Similarity search (similarity). Default value (similarity)]", 
        options=["mmr", "similarity"], 
        index=1  
    )

    k = st.select_slider(
        "Select amount of documents to be retrieved. Default value (5): ", 
        options=list(range(2, 16)), 
        value=4  
    )
    retriever = create_retriever_from_chroma(vectorstore_path="docs/chroma/", search_type=search_type, k=k, chunk_size=350, chunk_overlap=30)

    
    
    rag_chain = create_conversational_rag_chain(retriever)

    
    if user_question := st.text_input("Ask a question about your documents:"):
        handle_userinput(user_question,retriever,rag_chain)
        
    
    
    
    
 
    
    

    


def handle_userinput(user_question,retriever,rag_chain):
    st.session_state.messages.append({"role": "user", "content": user_question})
    st.chat_message("user").write(user_question)
    docs = retriever.get_relevant_documents(user_question)

    with st.sidebar:
        st.subheader("Your documents")
        with st.spinner("Processing"):
            for doc in docs:
                st.write(f"Document: {doc}")
    
    doc_txt = [doc.page_content for doc in docs]
    
    
    response = rag_chain.invoke({"context": doc_txt, "question": user_question})
    st.session_state.messages.append({"role": "assistant", "content": response})
    st.chat_message("assistant").write(response)
    

                
            



def create_conversational_rag_chain(retriever):
    
    

    callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])

    llm = ChatLlamaCpp(
        model_path = "gemma-2-2b-it-Q8_0.gguf",
        seed = 41,
        n_gpu_layers=0,
        temperature=0.0,
        n_ctx=25000,
        n_batch=2000,
        max_tokens=250,
        repeat_penalty=1.7,
        last_n_tokens_size = 250,
        callback_manager=callback_manager,
        verbose=False,
    )

   
    
    template = """Answer the question, based only on the following context:
    {context}. Be consise. "Avoid naming. Contextualize your answer.
    Question: {question}
    """
    
    prompt = ChatPromptTemplate.from_template(template)
    



    rag_chain = prompt | llm | StrOutputParser()


    return rag_chain


 

if __name__ == "__main__":
    main()