import os
import tempfile
import streamlit as st
from streamlit_chat import message
from langchain_community.vectorstores import Chroma
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_ollama.chat_models import ChatOllama
from langchain.retrievers.multi_query import MultiQueryRetriever
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors.flashrank_rerank import FlashrankRerank
from flashrank import Ranker, RerankRequest
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_ollama import OllamaEmbeddings
from langchain_community.embeddings import FastEmbedEmbeddings
from langchain_community.document_loaders import PyMuPDFLoader
from langchain.document_loaders.pdf import PyPDFDirectoryLoader
from langchain_community.document_loaders import WebBaseLoader
from langchain.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain.retrievers import BM25Retriever, EnsembleRetriever
from torch import cuda
from langchain_community.llms import LlamaCpp

device = "cuda" if cuda.is_available() else "cpu"

st.set_page_config(page_title="Chatbot", layout="wide")

class ChatPDF:
    def __init__(self):
        self.vector_db = None
        self.llm = ChatOllama(model="hf.co/bartowski/Llama-3.2-3B-Instruct-GGUF:IQ4_XS")
        self.llm = LlamaCpp(
            model_path="/home/chatbot/.cache/huggingface/hub/models--bartowski--Meta-Llama-3.1-8B-Instruct-GGUF/snapshots/bf5b95e96dac0462e2a09145ec66cae9a3f12067/Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf",
            temperature=0.75, #to maintain randomness of generated text
            n_gpu_layers=-1, #all layers of the model are loaded on the GPU if available
            max_tokens=2000, #Sets the maximum number of tokens the model will generate for each inference
            n_ctx=2048,
            top_p=1,
            # verbose=True,  # Verbose is required to pass to the callback manager
        )
        self.llm2 = LlamaCpp(
            model_path="/home/chatbot/.cache/huggingface/hub/models--bartowski--Llama-3.2-3B-Instruct-GGUF/snapshots/5ab33fa94d1d04e903623ae72c95d1696f09f9e8/Llama-3.2-3B-Instruct-IQ4_XS.gguf",
            temperature=0.75,
            n_gpu_layers=-1,
            max_tokens=2000,
            n_ctx=2048,
            top_p=1,
        )
        # self.llm2 = LlamaCpp(
        #     model_path="/home/kja/project/models--bartowski--Llama-3.2-3B-Instruct-GGUF/snapshots/5ab33fa94d1d04e903623ae72c95d1696f09f9e8/Llama-3.2-3B-Instruct-IQ4_XS.gguf",
        #     temperature=0.75,
        #     n_gpu_layers=-1,
        #     max_tokens=2000,
        #     n_ctx=2048,
        #     top_p=1,
        # )
        self.chain = None
        self.processed_files = []

    def ingest(self, file_path=None, file_name=None, webpage_url=None):
        self.vector_db = None
        self.processed_files.clear()

        if webpage_url:
            loader = WebBaseLoader(webpage_url)
            data = loader.load()
            st.success(f"Data from {webpage_url} loaded successfully!")
        else:
            self.processed_files.append(file_name)
            loader = PyMuPDFLoader(file_path=file_path)
            data = loader.load()
            st.success(f"{file_name} loaded successfully!")
        
        #breaks the document in chunks of 1000 characters.the last 200 characters of last chunk is repeated in current chunk to maintain context continuity
        #Overlapping preserves continuity, ensuring the chatbot understands full context when retrieving information.
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
        chunks = text_splitter.split_documents(data)

        self.vector_db = Chroma.from_documents(
            documents=chunks,
            embedding=FastEmbedEmbeddings(),
            collection_metadata={'hnsw': "cosine"},
            persist_directory='chromadbtest'
        )
        st.success("Vector database created successfully!")

        QUERY_PROMPT = PromptTemplate(
            input_variables=["question"],
            template="""You are an AI language model assistant. Your task is to generate three \
            different versions of the given user question to retrieve relevant documents from a vector \
            database. By generating multiple perspectives on the user question, your goal is to help\
            the user overcome some of the limitations of the distance-based similarity search. \
            Provide these alternative questions separated by newlines. Original question: {question}""",
        )

        retriever = MultiQueryRetriever.from_llm(
            self.vector_db.as_retriever(), 
            self.llm2,
            #self.llm,
            prompt=QUERY_PROMPT
        ) 

        keyword_retriever = BM25Retriever.from_documents(documents=chunks)
        main_retriever = EnsembleRetriever(retrievers=[retriever, keyword_retriever], weights=[0.5, 0.5])
        FlashrankRerank.model_rebuild()
        compressor = FlashrankRerank()
        compression_retriever = ContextualCompressionRetriever(
            base_compressor=compressor, base_retriever=main_retriever
        )

        chat_prompt = ChatPromptTemplate.from_template(
            """
            You are an expert assistant designed to answer questions based on the provided information.
            Use the context below to respond accurately and concisely to the query.
            While giving response, don't explicitly mention the document name or metadata unless otherwise asked.
            If the context does not contain the necessary information, state, 'The provided context does not contain enough information to answer the question'.

            Context:
            {context}

            Answer the question based on the above context:
            Question: {question}

            """
            )

        self.chain = (
            {"context": compression_retriever, "question": RunnablePassthrough() }
            | chat_prompt
            | self.llm2
            #|self.llm
            | StrOutputParser()
        )

    def ask(self, question):
        if not self.chain:
            return "Please upload your files first."
        return self.chain.invoke(question)

    def clear(self):
        self.vector_db = None
        self.chain = None
        self.processed_files.clear()


def display_messages():

    for i, (msg, is_user) in enumerate(st.session_state["messages"]):
        message(msg, is_user=is_user, key=str(i))
    st.session_state["thinking_spinner"] = st.empty()


def process_input():
    if st.session_state["user_input"].strip():
        user_text = st.session_state["user_input"].strip()
        with st.session_state["thinking_spinner"], st.spinner("Thinking..."):
            agent_text = st.session_state["assistant"].ask(user_text)

        st.session_state["messages"].append((user_text, True))
        st.session_state["messages"].append((agent_text, False))
        st.session_state["user_input"] = ""


def read_and_save_file():
    st.session_state["assistant"].clear()
    st.session_state["messages"] = []
    st.session_state["user_input"] = ""

    for file in st.session_state["file_uploader"]:
        with tempfile.NamedTemporaryFile(delete=False) as tf:
            tf.write(file.getbuffer())
            file_path = tf.name

        with st.session_state["ingestion_spinner"], st.spinner(f"Ingesting {file.name}..."):
            st.session_state["assistant"].ingest(file_path=file_path, file_name=file.name)
        os.remove(file_path)


def ingest_webpage():
    st.session_state["assistant"].clear()
    st.session_state["messages"] = []
    st.session_state["user_input"] = ""

    webpage_url = st.session_state["webpage_url"]
    with st.session_state["ingestion_spinner"], st.spinner(f"Ingesting data from {webpage_url}..."):
        st.session_state["assistant"].ingest(webpage_url=webpage_url)


def page():
    if "messages" not in st.session_state:
        st.session_state["messages"] = []
        st.session_state["assistant"] = ChatPDF()

    st.markdown(
    """
    <style>
        .title-container {
            display: flex;
            justify-content: center;
            align-items: center;
            font-size: 40px;
            font-weight: bold;
            color: #1E90FF;  # Customize the title color
        }
    </style>
    <div class="title-container">
        Chatbot - AI Assistant🌟<br><br>
        
        
    </div>
    """,
    unsafe_allow_html=True
)
    
    st.markdown(
    """
    <style>
        .welcome-text {
            font-size: 20px; /* Slightly larger text for emphasis */
            font-weight: bold;
            color: #20C997; /* Vibrant teal for attention */
            margin-bottom: 20px;
        }
        .sub-heading {
            font-size: 20px; /* Smaller size for sub-heading */
            font-weight: bold;
            color: #FFD700; /* Golden yellow for contrast */
            margin-top: 30px;
            margin-bottom: 10px;
        }
        .instructions {
            font-size: 16px; /* Smaller text for instructions */
            color: #FFFFFF; /* White text for black background themes */
            line-height: 1.6;
        }
    </style>

    <div class="welcome-text">
        Welcome to your AI-Powered Document Assistant!!  
        Chat with your documents and ask questions effortlessly.
    </div>

    <div class="sub-heading">Get Started!!</div>
    <div class="instructions">
        1. <b>Upload Your Documents</b>: The system accepts multiple PDF files at once, analyzing the content to provide comprehensive insights. <br>
        2. <b>Ask a Question</b>: After processing the documents, ask any question related to the content of your uploaded documents for a precise answer. <br><br><br>
    </div>
    """,
    unsafe_allow_html=True
)


  
    


    with st.sidebar:
        st.header("Options")
        st.subheader("📄 Upload Documents")
        st.file_uploader(
            "Upload PDF documents",
            type=["pdf"],
            key="file_uploader",
            on_change=read_and_save_file,
            accept_multiple_files=True,
        )
        st.subheader("Source Status")
        if st.session_state["assistant"].processed_files:
            st.write(f"""Uploaded {len(st.session_state["assistant"].processed_files)} document(s)""")
        else:
            st.write("No documents uploaded yet.")


    st.session_state["ingestion_spinner"] = st.empty()
    display_messages()
    st.text_input("Ask a question ✍", key="user_input", on_change=process_input, placeholder="Type your question here...")
    



if __name__ == "__main__":
    page()