import os
import chainlit as cl
from operator import itemgetter
from langchain_huggingface import HuggingFaceEndpoint
from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEndpointEmbeddings
from langchain_core.prompts import PromptTemplate
from langchain.schema.runnable.config import RunnableConfig

HF_LLM_ENDPOINT = os.environ["HF_LLM_ENDPOINT"]
HF_EMBED_ENDPOINT = os.environ["HF_EMBED_ENDPOINT"]
HF_TOKEN = os.environ["HF_TOKEN"]

# -- Data loading, embeddings, vector store --
document_loader = TextLoader("./data/paul_graham_essays.txt")
documents = document_loader.load()

text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=30)
split_documents = text_splitter.split_documents(documents)
print("chunked docs")

hf_embeddings = HuggingFaceEndpointEmbeddings(
    model=HF_EMBED_ENDPOINT,
    task="feature-extraction",
    huggingfacehub_api_token=HF_TOKEN,
)
print("initialized embeddings")

vectorstore_file = "./data/vectorstore"
vectorstore = None

if os.path.exists(vectorstore_file):
   vectorstore = FAISS.load_local(
        vectorstore_file, 
        hf_embeddings, 
        allow_dangerous_deserialization=True 
    )
   hf_retriever = vectorstore.as_retriever()
else:
    for i in range(0, len(split_documents), 32):
        print(f"Embedding batch {i}")
        if i == 0:
            vectorstore = FAISS.from_documents(split_documents[i:i+32], hf_embeddings)
            print("initialized vectorstore")
            continue
        vectorstore.add_documents(split_documents[i:i+32])
        hf_retriever = vectorstore.as_retriever()
print("Created vectorstore")

## -- System prompt --
RAG_PROMPT_TEMPLATE = """\
<|start_header_id|>system<|end_header_id|>
You are a helpful expert in tech and entrepreneurship. You answer user questions based on provided context. 
If you can't answer the question with the provided context, say you don't know.<|eot_id|>

<|start_header_id|>user<|end_header_id|>
User Query:
{query}

Context:
{context}<|eot_id|>

<|start_header_id|>assistant<|end_header_id|>
"""
rag_prompt = PromptTemplate.from_template(RAG_PROMPT_TEMPLATE)
print("created prompt")

## -- LLM endpoint --
hf_llm = HuggingFaceEndpoint(
    endpoint_url=HF_LLM_ENDPOINT,
    max_new_tokens=512,
    top_k=10,
    top_p=0.95,
    temperature=0.3,
    repetition_penalty=1.15,
    huggingfacehub_api_token=HF_TOKEN,
)
print("initialized endpoint")

@cl.author_rename
def rename(original_author: str):
    rename_dict = {
        "Assistant" : "Paul Graham Essay Bot"
    }
    return rename_dict.get(original_author, original_author)

@cl.on_chat_start
async def start_chat():
    lcel_rag_chain = (
        {"context": itemgetter("query") | hf_retriever, "query": itemgetter("query")}
        | rag_prompt | hf_llm
    )
    cl.user_session.set("lcel_rag_chain", lcel_rag_chain)
    await cl.Message(content="I'm ready to chat about tech and entrepreneurship!").send()

@cl.on_message  
async def main(message: cl.Message):
    lcel_rag_chain = cl.user_session.get("lcel_rag_chain")
    msg = cl.Message(content="")
    for chunk in await cl.make_async(lcel_rag_chain.stream)(
        {"query": message.content},
        config=RunnableConfig(callbacks=[cl.LangchainCallbackHandler()]),
    ):
        await msg.stream_token(chunk)

    await msg.send()