# Import necessary modules and define env variables

from langchain.chains import RetrievalQA
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.vectorstores import FAISS
from langchain_community.chat_models.huggingface import ChatHuggingFace
from langchain.prompts.chat import (
    ChatPromptTemplate
)
from langchain_community.llms import HuggingFaceHub
from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
import tempfile
from langchain_community.embeddings import HuggingFaceEmbeddings
import os
import io
from langchain.text_splitter import RecursiveCharacterTextSplitter
import chainlit as cl
import PyPDF2
import sys
from dotenv import load_dotenv
import huggingface_hub
load_dotenv()


from langchain.globals import set_verbose

set_verbose(False)

os.environ["HUGGINGFACEHUB_API_TOKEN"] = os.getenv("HUGGINGFACEHUB_API_TOKEN")

huggingface_hub.login(token=str(os.getenv("HUGGINGFACEHUB_API_TOKEN")))

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

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

text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=250)

@cl.on_chat_start
async def on_chat_start():
    await cl.Message(content="Hello there, Welcome to Laws of Power chat app!").send()
    msg = cl.Message(content=f"Processing Laws of Power...")
    await msg.send()
    loader = PyPDFLoader('./48lawsofpower.pdf')
    pages = loader.load_and_split()

# Create a Chroma vector store
    embeddings = HuggingFaceEmbeddings()
    faiss_index = FAISS.from_documents(pages, embeddings)

# Clean up the temporary file   
    pdf = PyPDF2.PdfReader('./48lawsofpower.pdf')
    pdf_text = ""
    for page in pdf.pages:
        pdf_text += page.extract_text()

    # Split the text into chunks
    texts = text_splitter.split_text(pdf_text)

    repo_id = "HuggingFaceH4/zephyr-7b-beta"

    chain_type_kwargs = {"prompt": prompt}

    llm = HuggingFaceHub(
        repo_id=repo_id,
        verbose = False
    )
    chat_model = ChatHuggingFace(llm=llm)
    # Create a chain that uses the Chroma vector store
    chain = RetrievalQA.from_chain_type(
        chat_model,
        chain_type="stuff",
        retriever=faiss_index.as_retriever(),
        chain_type_kwargs=chain_type_kwargs
    )

    # Let the user know that the system is ready
    msg.content = f"Processing Laws of Power done. You can now ask questions!"
    await msg.update()

    cl.user_session.set("chain", chain)


@cl.on_message
async def main(message:str):
    message = message.content
    chain = cl.user_session.get("chain")
    cb = cl.AsyncLangchainCallbackHandler()
    res = await chain.ainvoke(message, callbacks=[cb])
    answer = res
    await cl.Message(content=answer).send()