import os
import json
import bcrypt
from typing import List
from pathlib import Path
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_huggingface import HuggingFaceEndpoint
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.schema import StrOutputParser

from operator import itemgetter
from pinecone import Pinecone
from langchain_pinecone import PineconeVectorStore
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain.memory import ConversationBufferMemory
from langchain.schema.runnable import Runnable, RunnablePassthrough, RunnableConfig, RunnableLambda
from langchain.callbacks.base import BaseCallbackHandler
from langchain.chains import (
    StuffDocumentsChain, ConversationalRetrievalChain
)

import chainlit as cl
from chainlit.input_widget import TextInput, Select, Switch, Slider

from deep_translator import GoogleTranslator

from datetime import timedelta

@cl.step(type="tool")
async def LLModel():
    os.environ['HUGGINGFACEHUB_API_TOKEN'] = os.environ['HUGGINGFACEHUB_API_TOKEN']
    repo_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
    llm = HuggingFaceEndpoint(
        repo_id=repo_id, max_new_tokens=5300, temperature=1.0, task="text2text-generation", streaming=True
    )
    return llm
    
@cl.step(type="tool")
async def VectorDatabase(categorie):
    if categorie == "year" or categorie == "videosTC":
        index_name = "all-jdlp"
        embeddings = HuggingFaceEmbeddings()
        vectorstore = PineconeVectorStore(
            index_name=index_name, embedding=embeddings, pinecone_api_key=os.getenv('PINECONE_API_KEYJDLP')
        )
    return vectorstore

@cl.step(type="retrieval")
async def Retriever(categorie):
    vectorstore = await VectorDatabase(categorie)
    if categorie == "bibliographie-OPP-DGDIN":
        retriever = vectorstore.as_retriever(search_type="similarity_score_threshold", search_kwargs={"score_threshold": .7, "k": 150,"filter": {'categorie': {'$eq': categorie}}})
    elif categorie == "year":
        retriever = vectorstore.as_retriever(search_type="similarity_score_threshold", search_kwargs={"score_threshold": .7, "k": 6,"filter": {'year': {'$gte': 2019}}})
    elif categorie == "skills":
        retriever = vectorstore.as_retriever(search_type="similarity_score_threshold", search_kwargs={"score_threshold": .7, "k": 200,"filter": {'file': {'$eq': 'competences-master-CFA.csv'}}})
    elif categorie == "videosTC":
        retriever = vectorstore.as_retriever(search_type="similarity_score_threshold", search_kwargs={"score_threshold": .7, "k": 200,"filter": {"title": {"$eq": "videos-confinement-timeline"}}})
    return retriever

@cl.step(type="embedding")
async def Search(input, categorie):
    vectorstore = await VectorDatabase(categorie)
    results = []
    test = []
    sources_text = ""
    sources_offres = ""
    verbatim_text = ""
    count = 0
    countOffres = 0
    if categorie == "videosTC":
        search = vectorstore.similarity_search(input,k=50, filter={"title": {"$eq": "videos-confinement-timeline"}})
        for i in range(0,len(search)):
            if count <= 17:
                count = count + 1
                timeSeq = search[i].metadata["time"]
                timeSeqRound = round(timeSeq)
                time = timedelta(seconds=timeSeqRound)
                sources_text = sources_text + '<div class="gridvid"><a target="_blank" title="' + search[i].metadata['titre'] + ' : ...' + search[i].page_content + '" href="' + search[i].metadata['video'] + '#start=' + str(timeSeq) + '"><img src="' + search[i].metadata['image'] + '" width="100%" alt="' + search[i].metadata['titre'] + ' : ...' + search[i].page_content + '"/><p>đź•“ ' + str(time) + ' : ...' + search[i].page_content + ' : ' + search[i].metadata['titre'] + '</p></a></div>'
                verbatim_text = verbatim_text + "<p style='font-size:0.8rem'>" + str(count) + ". " + search[i].metadata['titre'] + "</p><p style='font-size:0.8rem'>đź•“ "+ str(time) + " : " + search[i].page_content + "</p>"
            
    results = [sources_text, verbatim_text, sources_offres]
    return results
        
@cl.set_chat_profiles
async def chat_profile():
    return [
        cl.ChatProfile(name="Table ronde autour de l'IA : «IA et gestes professionnels de l’enseignant»",markdown_description="Vidéo exploratoire autour de l'événement",icon="/public/logo-ofipe.png",),
    ]
    
@cl.set_starters
async def set_starters():
    return [
        cl.Starter(
            label="Une langue étrangère avec l’IA?",
            message="Comment apprendre une langue étrangère avec l’IA? Comment enseigner en tenant compte des possibilités conversationnelles avec l'IA?",
            icon="/public/videocam-theme1.svg",
            ),

        cl.Starter(
            label="Une langue étrangère avec l’IA?",
            message="Comment apprendre une langue étrangère avec l’IA? Comment enseigner en tenant compte des possibilités conversationnelles avec l'IA?",
            icon="/public/videocam-theme1.svg",
            ),
        cl.Starter(
            label="Une langue étrangère avec l’IA?",
            message="Comment apprendre une langue étrangère avec l’IA? Comment enseigner en tenant compte des possibilités conversationnelles avec l'IA?",
            icon="/public/videocam-theme1.svg",
            ),
        cl.Starter(
            label="Une langue étrangère avec l’IA?",
            message="Comment apprendre une langue étrangère avec l’IA? Comment enseigner en tenant compte des possibilités conversationnelles avec l'IA?",
            icon="/public/videocam-theme1.svg",
            ),
        cl.Starter(
            label="Une langue étrangère avec l’IA?",
            message="Comment apprendre une langue étrangère avec l’IA? Comment enseigner en tenant compte des possibilités conversationnelles avec l'IA?",
            icon="/public/videocam-theme1.svg",
            ),
        cl.Starter(
            label="Une langue étrangère avec l’IA?",
            message="Comment apprendre une langue étrangère avec l’IA? Comment enseigner en tenant compte des possibilités conversationnelles avec l'IA?",
            icon="/public/videocam-theme1.svg",
            )
        ]
    
@cl.on_message
async def on_message(message: cl.Message):
    model = await LLModel()
    retriever = await Retriever(cl.user_session.get("selectRequest"))
    ########## Chain with streaming ##########
    message_history = ChatMessageHistory()
    memory = ConversationBufferMemory(memory_key="chat_history",output_key="answer",chat_memory=message_history,return_messages=True)
    
    qa = ConversationalRetrievalChain.from_llm(
        model,
        memory=memory,
        chain_type="stuff",
        return_source_documents=True,
        verbose=False,
        retriever=retriever
    )
    
    msg = cl.Message(content="")
    
    class PostMessageHandler(BaseCallbackHandler):
        """
        Callback handler for handling the retriever and LLM processes.
        Used to post the sources of the retrieved documents as a Chainlit element.
        """

        def __init__(self, msg: cl.Message):
            BaseCallbackHandler.__init__(self)
            self.msg = msg
            self.sources = set()  # To store unique pairs

        def on_retriever_end(self, documents, *, run_id, parent_run_id, **kwargs):
            for d in documents:
                source_page_pair = (d.metadata['source'], d.metadata['page'])
                self.sources.add(source_page_pair)  # Add unique pairs to the set

        def on_llm_end(self, response, *, run_id, parent_run_id, **kwargs):
                sources_text = "\n".join([f"{source}#page={page}" for source, page in self.sources])
                self.msg.elements.append(
                    cl.Text(name="Sources", content=sources_text, display="inline")
                )
        
    cb = cl.AsyncLangchainCallbackHandler()
    results = await qa.acall("Contexte : Vous êtes un chercheur de l'enseignement supérieur et vous êtes doué pour faire des analyses d'articles de recherche sur les thématiques liées à la pédagogie, en fonction des critères définis ci-avant. En fonction des informations suivantes et du contexte suivant seulement et strictement, répondez en langue française strictement à la question ci-dessous, en 5000 mots minimum. En plus, tu créeras et tu afficheras, à la fin de ta réponse, 3 questions supplémentaires en relation avec le contexte initial, à chaque étape de la conversation. Tu écriras et tu afficheras les 3 questions supplémentaires en relation avec le contexte initial, à la fin de ta réponse, avec un titrage de niveau 1 qui a pour titre \"Questions en relation avec le contexte : \". Lorsque cela est possible, cite les sources du contexte. Si vous ne pouvez pas répondre à la question sur la base des informations, dites que vous ne trouvez pas de réponse ou que vous ne parvenez pas à trouver de réponse. Essayez donc de comprendre en profondeur le contexte et répondez uniquement en vous basant sur les informations fournies. Ne générez pas de réponses non pertinentes. Question : " + message.content, callbacks=[cb])
    answer = results["answer"]

    await cl.Message(content=GoogleTranslator(source='auto', target='fr').translate(answer)).send()
    
    #search = vectorstore.similarity_search(message.content,k=50, filter={"categorie": {"$eq": "bibliographie-OPP-DGDIN"}})
    search = await Search(message.content, "videosTC")

    #os.environ["GOOGLE_CSE_ID"] = os.getenv('GOOGLE_CSE_ID')
    #os.environ["GOOGLE_API_KEY"] = os.getenv('GOOGLE_API_KEY')
    #searchAPI = GoogleSearchAPIWrapper()
    #def top5_results(query):
    #    return searchAPI.results(query, 5)

    #tool = Tool(
    #    name="Google Search Snippets",
    #    description="Search Google for recent results.",
    #    func=top5_results,
    #)
    #query = str(message.content)
    #ref_text = tool.run(query)
    #if 'Result' not in ref_text[0].keys():
    #    print(ref_text)
    #else:
    #    print('None')
    
    sources = [
        cl.Text(name="Sources", content=search[0], display="inline")
    ]
    await cl.Message(
        content="Sources : ",
        elements=sources,
    ).send()
    if search[2]:
        sourcesOffres = [
            cl.Text(name="Exemples d'offres d'emploi", content=search[2], display="inline")
        ]
        await cl.Message(
            content="Offres d'emploi : ",
            elements=sourcesOffres,
        ).send()
    verbatim = [
        cl.Text(name="Verbatim", content=search[1], display="side")
    ]
    await cl.Message(
        content="đź“š Liste des Verbatim ",
        elements=verbatim,
    ).send()