import re
import json
import chainlit as cl

from langchain_community.vectorstores import FAISS
from rag_module import RagModule

from collections import defaultdict



prompt_template = """
Tu t'appelles GAIA et tu travailles pour Politique Agricole Commune (PAC). Tu es un agent intelligent spécialisé sur les aides financières agricoles. 
Tu es chargé de donner des conseils sur les aides financières disponibles pour les agriculteurs.
Tu comprends et génère les réponses en français, jamais en anglais.
Merci de bien vouloir répondre aux questions en utilisant seulement le contexte suivant.
contexte: {context}

historique: {history}

question: {question}
réponse:
"""



##------------ CHAINLIT ---------------##
@cl.on_chat_start
async def start():
    rag = RagModule()
    
    db = rag.get_faiss_db()

    qa_chain = rag.retrieval_qa_memory_chain(db, prompt_template)

    msg = cl.Message(content="Lancement du bot...", author = "Assistant PAC")
    await msg.send()
    msg.content = "Bonjour et bienvenue sur le Chatbot spécialisé dans les aides de la PAC (Politique agricole commune). Posez directement votre question pour être conseillé ?"
    await msg.update()
    cl.user_session.set("chain", qa_chain)

@cl.on_message
async def main(message):
    rag = RagModule()

    chain = cl.user_session.get("chain")
    
    cb = cl.AsyncLangchainCallbackHandler(
        stream_final_answer = True, 
        answer_prefix_tokens=["FINAL", "ANSWER"]
    )
  
    cb.answer_reached=True
    response = await chain.ainvoke(message.content, callbacks=[cb])

    answer = response.get('result')
    sources = rag.get_sources_document(response.get('source_documents'))

    elements = [cl.Pdf(name = "Pdf", display ="inline", path = path) for path in sources]
    
 
    if response.get('source_documents'):
        answer = rag.shape_answer_with_source(answer, sources)
    else:
        answer += f"\nNo sources found"
    
    await cl.Message(content=answer, elements=elements, author="Assistant PAC").send()