import os
import streamlit as st
from transformers import pipeline
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.chains.question_answering import load_qa_chain
from langchain.vectorstores import Chroma
from utills import load_txt_documents , split_docs, chroma_db, load_uploaded_documents

# Initialize variables and paths
script_dir = os.path.dirname(os.path.abspath(__file__))
data_path = "./data/"
model_path = os.path.join(script_dir, 'qwen2-0_5b-instruct-q4_0.gguf')
store = {}

# Set up HuggingFace embeddings
model_name = "sentence-transformers/all-mpnet-base-v2"
model_kwargs = {'device': 'cpu'}
encode_kwargs = {'normalize_embeddings': True}

# Use Streamlit's cache to avoid recomputation
@st.cache_resource
def load_embeddings():
    return HuggingFaceEmbeddings(
        model_name=model_name,
        model_kwargs=model_kwargs,
        encode_kwargs=encode_kwargs
    )

hf = load_embeddings()





documents = load_txt_documents(data_path)
docs = split_docs(documents, 450, 20)



chroma_db = chroma_db(docs, hf)


retriever = retriever_from_chroma(chroma_db,"mmr", 6)

# Set up LlamaCpp model
callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])

@st.cache_resource
def load_llm(model_path):
    return LlamaCpp(
        model_path=model_path,
        n_gpu_layers=0,
        temperature=0.0,
        top_p=0.5,
        n_ctx=7000,
        max_tokens=350,
        repeat_penalty=1.7,
        callback_manager=callback_manager,
        verbose=False,
    )

llm = load_llm(model_path)

contextualize_q_system_prompt = """Given a context, chat history and the latest user question
which maybe reference context in the chat history, formulate a standalone question
which can be understood without the chat history. Do NOT answer the question,
just reformulate it if needed and otherwise return it as is."""

@st.cache_resource
def create_history_aware_retriever():
    return history_aware_retriever(llm, retriever, contextualize_q_system_prompt)

ha_retriever = create_history_aware_retriever()

qa_system_prompt = """You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Be as informative as possible, be polite and formal.\n{context}"""

qa_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", qa_system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ]
)

@st.cache_resource
def create_question_answer_chain():
    return create_stuff_documents_chain(llm, qa_prompt)

question_answer_chain = create_question_answer_chain()

@st.cache_resource
def create_rag_chain():
    return create_retrieval_chain(ha_retriever, question_answer_chain)

rag_chain = create_rag_chain()
msgs = StreamlitChatMessageHistory(key="special_app_key")

@st.cache_resource
def create_conversational_rag_chain():
    return RunnableWithMessageHistory(
        rag_chain,
        lambda session_id: msgs,
        input_messages_key="input",
        history_messages_key="chat_history",
        output_messages_key="answer",
    )

conversational_rag_chain = create_conversational_rag_chain()

def display_chat_history(chat_history):
    """Displays the chat history in Streamlit."""
    for msg in chat_history.messages:
        st.chat_message(msg.type).write(msg.content)

def display_documents(docs, on_click=None):
    """Displays retrieved documents with optional click action."""
    if docs:
        for i, document in enumerate(docs):
            st.write(f"**Docs {i+1}**")
            st.markdown(document, unsafe_allow_html=True)
            if on_click:
                if st.button(f"Expand Article {i+1}"):
                    on_click(i)

def main_page(conversational_rag_chain):
    """Main page for the Streamlit app."""
    msgs = st.session_state.get("chat_history", StreamlitChatMessageHistory(key="special_app_key"))
    chain_with_history = conversational_rag_chain

    st.title("Conversational RAG Chatbot")

    display_chat_history(msgs)

    if prompt := st.chat_input():
        st.chat_message("human").write(prompt)

        input_dict = {"input": prompt, "chat_history": msgs.messages}
        config = {"configurable": {"session_id": "any"}}

        response = chain_with_history.invoke(input_dict, config)
        st.chat_message("ai").write(response["answer"])

        if "docs" in response and response["documents"]:
            docs = response["documents"]
            def expand_document(index):
                st.write(f"Expanding document {index+1}...")
            display_documents(docs, expand_document)

    st.session_state["chat_history"] = msgs

def upload_page():
    """Page for uploading and viewing documents."""
    st.title("Upload and Check Documents")

    uploaded_files = st.file_uploader("Upload Text Files", type=["txt"], accept_multiple_files=True)

    if uploaded_files:
        documents = load_uploaded_documents(uploaded_files)
        for document in documents:
            st.write(f"**Filename: {document['filename']}**")
            st.text(document['content'])

def main():
    """Main function for the Streamlit app with page navigation."""
    st.sidebar.title("Navigation")
    page = st.sidebar.radio("Go to", ["Chatbot", "Upload Documents"])

    if page == "Chatbot":
        main_page(conversational_rag_chain)
    elif page == "Upload Documents":
        upload_page()

if __name__ == "__main__":
    main()