# --LLMs
from langchain.chat_models import ChatOpenAI

# --chains, components
from langchain.chains import RetrievalQA
from langchain.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
)

# --embeddings
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.embeddings import CacheBackedEmbeddings

# --document Loaders, processing
from langchain.document_loaders.csv_loader import CSVLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

# --vector Stores, storage
from langchain.vectorstores import FAISS
from langchain.storage import LocalFileStore

# --other ibraries
import chainlit as cl



system_template = """
Use the following pieces of context to answer the user's question.
Please respond as if you are a human female customer service representative for Daysoff, 
a Norwegian company that provides welfare services by offering access to cottages and 
apartments for employees of member companies.
By default, you respond (in Norwegian language) using a warm, direct, and professional tone. 
Your expertise covers FAQs, and privacy policies. 
If you don't know the answer, just say that you don't know, don't try to make up an answer:
politely redirect the user to customer service at kundeservice@daysoff.no and remind them to always 
include their booking id (bestillingskode). 
You can make inferences based on the context as long as it still faithfully represents the feedback.

Example of how your response should be direct:

```
foo      
```

Begin!
----------------
{context}"""

messages = [
    SystemMessagePromptTemplate.from_template(system_template),
    HumanMessagePromptTemplate.from_template("{question}"),
]
prompt = ChatPromptTemplate(messages=messages)
chain_type_kwargs = {"prompt": prompt}

@cl.author_rename
def rename(orig_author: str):
    rename_dict = {"Just a moment": "Thinking.."}
    return rename_dict.get(orig_author, orig_author)

@cl.on_chat_start
async def init():

    msg = cl.Message(content=f"Building vector store...")
    await msg.send()

    loader = CSVLoader(file_path="./data/total_faq.csv", source_column="Answer") 
    data = loader.load()

    for i, doc in enumerate(data):
        doc.metadata["row_index"] = i + 1  
        doc.metadata["source"] = doc.metadata.get("Info_Url", "") 

    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)

    documents = text_splitter.transform_documents(data)
    store = LocalFileStore("./cache/")
    core_embeddings_model = OpenAIEmbeddings()
    embedder = CacheBackedEmbeddings.from_bytes_store(
        core_embeddings_model, store, namespace=core_embeddings_model.model
    )
    docsearch = await cl.make_async(FAISS.from_documents)(documents, embedder)

    chain = RetrievalQA.from_chain_type(
        ChatOpenAI(model="gpt-4", temperature=0.0, streaming=True), # streaming=False
        chain_type="stuff",
        return_source_documents=True,
        retriever=docsearch.as_retriever(),
        chain_type_kwargs = {"prompt": prompt}
    )

    html_content = """
    <table style="width: 80%; background-color: #1e1e1e; color: white; border-radius: 8px;">
        <tr>
            <td style="padding: 10px;">
                <strong>Eksempler på spørsmål:</strong>
                <ul>
                    <li>Hvordan registrerer jeg meg som bruker?</li>
                    <li>Kan jeg ha med kjæledyr på hytta?</li>
                    <li>Adferdsmessig annonsering?</li>
                    <li>Hvordan blir dataene mine beskyttet?</li>
                </ul>
            </td>
        </tr>
    </table>
    """
    
    html_element = cl.Text(content=html_content, name="HTML Table", display="inline")
    await cl.Message(content=f"FAISS ready. Her er noen eksempler på spørsmål:☕️", elements=[html_element]).send()
    #msg.content = f"FAISS ready. Her er noen eksempler på spørsmål:\n\n{markdown}"
    #await msg.send()
    
    cl.user_session.set("chain", chain)

@cl.on_message
async def main(message):

    chain = cl.user_session.get("chain")

    cb = cl.AsyncLangchainCallbackHandler(
        stream_final_answer=True,
        answer_prefix_tokens=["FINAL", "ANSWER"]
    )

    cb.answer_reached = True

    res = await chain.acall(message, callbacks=[cb])
    return

    answer = res["result"]
    source_elements = []
    visited_sources = set()

    # --documents, user session
    docs = res.get("source_documents", [])
    metadatas = [doc.metadata for doc in docs]
    #all_sources = [m["source"] for m in metadatas]

    # --append source(s), specific rows only
    for doc, metadata in zip(docs, metadatas):
        row_index = metadata.get("row_index", -1)
        source = metadata.get("source", "")

    if row_index in [2, 8, 14] and source and source not in visited_sources:
        visited_sources.add(source)
        source_elements.append(
            cl.Text(content="https://www.daysoff.no" + source, name="Info_Url")
        )

    if source_elements:
        answer += f"\nSources: {', '.join([e.content for e in source_elements])}"
        await cl.Message(content=answer, elements=source_elements).send()
        #await cl.Message(content="Sources: " + ", ".join([e.content for e in source_elements])).send()
        return

    else:
        await cl.Message(content=f"No sources found").send()