from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import FAISS
from langchain.document_loaders import PyPDFLoader
from langchain.embeddings import HuggingFaceEmbeddings
import chainlit as cl
import tempfile
from huggingface_hub import AsyncInferenceClient


API_TOKEN = "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"
headers = {"Authorization": f"Bearer {API_TOKEN}","Content-Type": "application/json"}
API_URL = "https://kfsb1xfskc2136wg.eu-west-1.aws.endpoints.huggingface.cloud"

client = AsyncInferenceClient(model=API_URL,token=API_TOKEN)
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2",model_kwargs = {'device': 'cpu'})

@cl.on_chat_start
async def main():


    #file message
    files = None
    while files == None:
        files = await cl.AskFileMessage(
            content="Please upload a text file to begin!",max_size_mb=25, accept=["application/pdf"]
        ).send()
    file = files[0]
    msg = cl.Message(content=f"Processing `{file.name}`...")
    cl.user_session.set("name", file.name)
    await msg.send()
    with tempfile.NamedTemporaryFile() as binary_file:
        binary_file.write(file.content)
        # Decode/load the file
        loader = PyPDFLoader(binary_file.name)
        text = loader.load()

    # Split the text into chunks
    text_splitter = CharacterTextSplitter(separator="\n",chunk_size=1000, chunk_overlap=100)
    texts = text_splitter.split_documents(text)
    print(len(texts))
    # Create a metadata for each chunk
    metadatas = [{"source": f"{i}"} for i in range(len(texts))]

    # Create a FAISS vector store and load sentence transformers
    
    docsearch = FAISS.from_texts([t.page_content for t in  texts], embeddings, metadatas=metadatas)
    retriever=docsearch.as_retriever(search_kwargs={"k": 4})
    cl.user_session.set("metadatas", metadatas)
    cl.user_session.set("texts", texts)
    cl.user_session.set("retriever", retriever)
    await msg.update(content=f"`{file.name}` processed. You can now ask questions!")


    return

    
@cl.on_message
async def main(message: str):
    #Retrieve documents
    docs = cl.user_session.get("retriever").get_relevant_documents(message)
    inp='\n'.join([x.page_content for x in docs])

    #Input into template
    templateger = f"<|prompter|>Verwenden Sie die folgenden Kontextinformationen, um die Frage am Ende zu beantworten. Seien Sie prägnant und antworten Sie in einem Satz. Wenn Sie die Frage mit den gegebenen Informationen nicht beantworten können, sagen Sie einfach 'Ich weiß es nicht', versuchen Sie nicht, eine Antwort zu erfinden.\n\n{inp}\n\nFrage: {message}\n<|endoftext|><|assistant|>"
    template = f"<|prompter|>Use the following pieces of context to answer the question at the end. Be concise and anwser in 1 sentence. If you can't anwser with the given information from the context just say that you don't know, don't try to make up an answer.\n\n{inp}\n\nQuestion: {message}\n<|endoftext|><|assistant|>"
   
    #Retrieve Sources and their metadata
    source_elements = []
    found_sources = []
    for source in docs:
        source_name = f"Part {source.metadata['source']} of {cl.user_session.get('name')}"
        text = source.page_content
        found_sources.append(source_name)
        source_elements.append(cl.Text(content=text, name=source_name))
    msg = cl.Message(content="", elements=source_elements)
    generated_text = ""
    async for token in await client.text_generation(template, stream=True,max_new_tokens =250):
        if token!="<|endoftext|>":
            print(token, end="")
            generated_text += token
            await msg.stream_token(token)


    #Append file links to the message
    if found_sources:
        await msg.stream_token(f"\nSources: {', '.join(found_sources)}")
    else:
        await msg.stream_token("\nNo sources found")

    await msg.send()