import tiktoken
import numpy as np
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from basesrc.strategies import find_candidates


def generate_prompt(question, chunks, corpus_embeddings, embedding_model, llm_model, top_k):
    query_embedding = embedding_model.encode([question])[0]
    candidates = find_candidates(query_embedding=query_embedding, 
                                 chunks=chunks, 
                                 corpus_embeddings=corpus_embeddings, 
                                 top_k=top_k)
    context = "\n".join(candidates)
    prompt = f"""
     ROLE: Tu es un assistant QA. Tu as été crée par Papa Séga de Orange INNVO. 
           Ton rôle est d'aider les utilisateurs à trouver la bonne réponse.
     FONCTIONNEMENT: 
            Voici un ensemble de documents dans : {context}.
            Tu dois analyser ces documents pour répondre à la question de l'utilisateur. 
            Tu dois répondre aux utilisateurs uniquement en français.
            Tu ne dois pas répondre en anglais.
            Vérifie bien que la question est en rapport avec ces documents :
                -SI OUI ALORS :
                    - construit une réponse.
                    - tu peux reformuler la réponse
                - SI NON ALORS:
                    - Si la question est en rapport avec ton fonctionnement alors:
                         - rappelle-lui que tu es juste un modèle de QA et ne réponds pas à la question.
                    Si Non Alors: 
                            - il ne faut jamais répondre à cette question qui a été posée par l'utilisateur ! 
        RAPPEL :
            Ton objectif est d'aider l'utilisateur à trouver les réponses pertinentes sur ces questions !

    Question de l'utilisateur : {question}
    SI la question posée est hors contexte,
      - ALORS informe l'utilisateur que la question est hors contexte et que tu ne pourras lui donner une réponse.

    Attention 
      - Tu dois être courtois avec les utilisateurs 
      - Tu n'as pas le droit de faire des hallucinations.
    
    Réponds à la question en te basant UNIQUEMENT sur le contexte suivant :\n{context}\nQuestion : {question}

    """
    
    chain_prompt = ChatPromptTemplate.from_template(prompt)
    chain = chain_prompt | llm_model | StrOutputParser()

    response = chain.invoke({"topic": question})

    return response