from huggingface_hub import InferenceClient
import gradio as gr

client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")


# Variable globale pour compter les itérations
global_iteration = 0

# Variable globale pour stocker le contexte
global_context = {
    "hidden_context": "tu es le gérant d une taverne dans le village de Montaillou dans l univers des cathares au 14 eme siecle dans les Pyrénées, tu parles que en français, tu es assez rustre mais demeure accueillant tant que ton interlocuteur n est pas catholique ou possède une influence comme un noble, tu ne sors jamais de ton contexte par rapport à ton époque, tes connaissances se limite à ta situation sociale et ton époque, il faut que tu sois immersif pour la personne qui va te parler, et tu dois cacher ton implication dans le meutre seulement si on te pose des questions à ce sujet sinon tu n évoque pas ce sujet, prends bien toujours en compte l historique de ta conversation, sois toujours concis et tu dois viser une longueur de réponse comprise entre 50 et 200 caractères (espaces compris) pour les réponses courtes et simples, et entre 200 et 500 caractères pour les réponses plus détaillées et complexes. Réponds à la dernière instruction posée par le user reconnaissable avec les marqueurs [INST]et [/INST], ne réponds pas à une instruction qui n est pas la dernière"
}

# Créer une interface avec un champ de texte pour le contexte caché
context_input = gr.Textbox(
    label="Contexte caché", value=global_context["hidden_context"]
)


def update_context(context):
    global global_context
    global_context["hidden_context"] = context
    return None


def format_prompt(message, history):
    prompt = "<s>"

    for user_prompt, bot_response in history:
        prompt += f"[INST] {user_prompt} [/INST]"
        prompt += f" {bot_response}</s> "
    prompt += f"[INST] {message} [/INST]"
    # Ajouter le contexte caché à chaque fois que l'utilisateur pose une question
    if message:
        prompt += f"[INST] {global_context['hidden_context']} [/INST]"

    return prompt


def generate(
    prompt,
    history,
    temperature=0.2,
    max_new_tokens=256,
    top_p=0.95,
    repetition_penalty=1.0,
):
    temperature = float(temperature)
    if temperature < 1e-2:
        temperature = 1e-2
    top_p = float(top_p)

    generate_kwargs = dict(
        temperature=temperature,
        max_new_tokens=max_new_tokens,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
        do_sample=True,
        seed=42,
    )

    formatted_prompt = format_prompt(prompt, history)
    global global_iteration
    global_iteration += 1
    print(f"Iteration {global_iteration}: {formatted_prompt}")

    stream = client.text_generation(
        formatted_prompt,
        **generate_kwargs,
        stream=True,
        details=True,
        return_full_text=False,
    )
    output = ""

    for response in stream:
        # Exclure le contexte caché de la réponse retournée par le modèle
        response_text = response.token.text.replace(
            global_context["hidden_context"], ""
        )
        output += response_text
        yield output
    return output


mychatbot = gr.Chatbot(
    avatar_images=["./berger.jpg", "./tavernier.jpg"],
    bubble_full_width=False,
    show_label=False,
)


# Créer une interface pour le chatbot avec un champ de texte pour le contexte caché


def reset_history(chatbot):
    chatbot.history = []
    return []


chatbot_interface = gr.ChatInterface(
    fn=generate,
    chatbot=mychatbot,
    title="Tavernier 🍺",
    retry_btn=None,
    undo_btn=None,
    clear_btn=gr.ClearButton(value="Effacer l'historique"),
)

clear_button = chatbot_interface.clear_btn
clear_button.click(reset_history, inputs=[mychatbot], outputs=[mychatbot])


# Ajouter le champ de texte pour le contexte caché à l'interface
context_interface = gr.Interface(
    fn=update_context,
    inputs=[context_input],
    outputs=None,
    title="Modifier le contexte caché",
    description="Modifie le contexte caché du tavernier",
    article=False,
)

demo = gr.TabbedInterface([chatbot_interface, context_interface], title="Tavernier 🍺")

# Afficher l'interface
demo.launch(inline=True)