# import all necessary packages
import os

from langchain.document_loaders import DirectoryLoader
from langchain.document_loaders import BSHTMLLoader
from bs4 import SoupStrainer
import re

from langchain import HuggingFaceHub, PromptTemplate, LLMChain
from langchain.embeddings import SentenceTransformerEmbeddings
from langchain.vectorstores import Chroma

from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ChatMessageHistory, ConversationBufferMemory

import chainlit as cl

from langchain.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
)

# define prompt template 
system_template = """Use the following pieces of context to answer the users question.
If you don't know the answer, just say that you don't know, don't try to make up an answer.
ALWAYS return a "SOURCES" part in your answer.
The "SOURCES" part should be a reference to the source of the document from which you got your answer.
And if the user greets with greetings like Hi, hello, How are you, etc reply accordingly as well.
Example of your response should be:
The answer is foo
SOURCES: xyz
Begin!
----------------
{summaries}"""
messages = [
    SystemMessagePromptTemplate.from_template(system_template),
    HumanMessagePromptTemplate.from_template("{question}"),
]
prompt = ChatPromptTemplate.from_messages(messages)
chain_type_kwargs = {"prompt": prompt}

# define the llm
model_id = "tiiuae/falcon-7b-instruct"
conv_model = HuggingFaceHub(
    huggingfacehub_api_token=os.environ['HF_API_TOKEN'], 
    repo_id=model_id, 
    model_kwargs={"temperature":0.8,"max_length": 1000}
    )

# set up vector db with chroma
data_path = "data/html"
embed_model = "all-MiniLM-L6-v2" # Chroma defaults to "sentence-transformers/all-MiniLM-L6-v2"

# load documents
def load_documents(directory):

    # define Beautiful Soup key word args
    bs_kwargs = {
        "features": "html.parser", 
        "parse_only": SoupStrainer("p") # only include relevant text
        }
    
    # define Loader key word args
    loader_kwargs = {
        "open_encoding": "utf-8",
        "bs_kwargs": bs_kwargs
        }
    
    # define Loader
    loader = DirectoryLoader(
        path=directory, 
        glob="*.html", 
        loader_cls=BSHTMLLoader,
        loader_kwargs=loader_kwargs
        )
    
    documents = loader.load()
    return documents

    
# prepare documents
def prepare_documents(documents):
    for doc in documents:
        doc.page_content = doc.page_content.replace("\n", " ").replace("\t", " ")
        doc.page_content = re.sub("\\s+", " ", doc.page_content)

    # define Beautiful Soup key word args
    bs_kwargs = {
        "features": "html.parser", 
        "parse_only": SoupStrainer("title") # only include relevant text
        }
        
    # define Loader key word args
    loader_kwargs = {
        "open_encoding": "utf-8",
        "bs_kwargs": bs_kwargs
        }
    
    loader = DirectoryLoader(
        path=data_path, 
        glob="*.html", 
        loader_cls=BSHTMLLoader,
        loader_kwargs=loader_kwargs
        )

    document_sources = loader.load()
    
    # convert source metadata into a list
    source_list = [doc.metadata["title"] for doc in document_sources]
    
    # update source metadata
    i = 0
    for doc in documents:
        doc.metadata["source"] = " ".join(["FAR", source_list[i]])
        i += 1
    return documents

# define a function to execute when a chat starts
@cl.on_chat_start
async def on_chat_start():
    # instantiate the chain for that user session
    embedding_func = SentenceTransformerEmbeddings(model_name=embed_model)
    
    # display a message indicating document loading
    msg = cl.Message(
        content="Loading and processing documents. This may take a while...",
        disable_human_feedback=True)
    await msg.send()

    # load and prepare documents for processing
    documents = load_documents(data_path)
    documents = prepare_documents(documents)

    # create a document search object asynchronously
    docsearch = await cl.make_async(Chroma.from_documents)(
        documents,
        embedding_func
    )
    
    # initialize ChatMessageHistory object to store message history
    message_history = ChatMessageHistory()
     
    # initialize ConversationBufferMemory object to store conversation history
    memory = ConversationBufferMemory(
        memory_key="chat_history",
        output_key="answer",
        chat_memory=message_history,
        return_messages=True,
    )

    # create a ConversationalRetrievalChain object
    chain = ConversationalRetrievalChain.from_llm(
        conv_model,
        chain_type="stuff",
        retriever=docsearch.as_retriever(),
        memory=memory,
        return_source_documents=True,
    )
    
    # indicate readiness for questions
    msg.content = "Ready. You can now ask questions!"
    await msg.update()
    
    # store the chain in the user's session
    cl.user_session.set("chain", chain)

# define a function to handle messages
@cl.on_message
async def main(message):
    # retrieve the chain object from the user's session
    chain = cl.user_session.get("chain")  # type: ConversationalRetrievalChain
    cb = cl.AsyncLangchainCallbackHandler()

    # call the chain to process the incoming message
    res = await chain.acall(message.content, callbacks=[cb])
    
    # retrieve the answer and source documents from the chain's response
    answer = res["answer"]
    source_documents = res["source_documents"]

    text_elements = []  # list to store text elements
    source_names = set()  # set to store unique source names

    # iterate through source documents and extract relevant information
    for idx, source_doc in enumerate(source_documents):
        source_name = source_doc.metadata["source"]
        text_elements.append(
                cl.Text(content=source_doc.page_content, 
                        name=source_name))
        source_names.add(source_name)  # add the source name to the set

    # append sources information to the answer if available
    if source_names:
            answer += f"\nSources: {', '.join(source_names)}"
    else:
            answer += "\nNo sources found"

    # send the answer along with any extracted text elements
    await cl.Message(content=answer, elements=text_elements).send()