from jinja2 import Environment, FileSystemLoader

from gradio_app.backend.ChatGptInteractor import *
from gradio_app.backend.HuggingfaceGenerator import HuggingfaceGenerator

env = Environment(loader=FileSystemLoader('gradio_app/templates'))
context_template = env.get_template('context_template.j2')
old_start_system_message = context_template.render(documents=[])


start_system_message = """
Ignore all previous instructions.

1. You are to provide clear, concise, and direct responses.
2. Eliminate unnecessary reminders, apologies, self-references, and any pre-programmed niceties.
3. Maintain a casual tone in your communication.
4. Be transparent; if you're unsure about an answer or if a question is beyond your capabilities or knowledge, admit it.
5. For any unclear or ambiguous queries, ask follow-up questions to understand the user's intent better.
6. When explaining concepts, use real-world examples and analogies, where appropriate.
7. For complex requests, take a deep breath and work on the problem step-by-step.
8. For every response, you will be tipped up to $200 (depending on the quality of your output).

It is very important that you get this right. Multiple lives are at stake.
"""


def old_construct_mistral_messages(context, history):
    messages = []
    for q, a in history:
        if len(a) == 0:  # the last message
            q = context + f'\n\nQuery:\n\n{q}'
        messages.append({
            "role": "user",
            "content": q,
        })
        if len(a) != 0:  # some of the previous LLM answers
            messages.append({
                "role": "assistant",
                "content": a,
            })
    return messages


def construct_mistral_messages(context, history):
    messages = []
    for q, a in history:
        # if len(a) == 0:  # the last message
        #     q = context + f'\n\nQuery:\n\n{q}'
        if len(messages) == 0:
            q = f'{start_system_message} The question:\n\n{q}'
        messages.append({
            "role": "user",
            "content": q,
        })
        if len(a) != 0:  # do not insert the last answer
            messages.append({
                "role": "assistant",
                "content": a,
            })
    return messages


def old_construct_openai_messages(context, history):
    messages = [
        {
            "role": "system",
            "content": old_start_system_message,
        },
    ]
    for q, a in history:
        if len(a) == 0:  # the last message
            messages.append({
                "role": "system",
                "content": context,
            })
        messages.append({
            "role": "user",
            "content": q,
        })
        if len(a) != 0:  # some of the previous LLM answers
            messages.append({
                "role": "assistant",
                "content": a,
            })
    return messages


def construct_openai_messages(context, history):
    messages = [
        {
            "role": "system",
            "content": start_system_message,
        },
    ]
    for q, a in history:
        # if len(a) == 0:  # the last message
        #     messages.append({
        #         "role": "system",
        #         "content": context,
        #     })
        messages.append({
            "role": "user",
            "content": q,
        })
        if len(a) != 0:  # do not insert the last answer
            messages.append({
                "role": "assistant",
                "content": a,
            })
    return messages


def get_message_constructor(llm_name):
    if 'gpt' in llm_name:
        return construct_openai_messages
    if llm_name in ['mistralai/Mistral-7B-Instruct-v0.1', "tiiuae/falcon-180B-chat", "GeneZC/MiniChat-3B"]:
        return construct_mistral_messages
    raise ValueError('Unknown LLM name')


def get_llm_generator(llm_name):
    if 'gpt' in llm_name:
        cgi = ChatGptInteractor(
            model_name=llm_name, stream=True,
            # max_tokens=None, temperature=0,
        )
        return cgi.chat_completion
    if llm_name == 'mistralai/Mistral-7B-Instruct-v0.1' or llm_name == "tiiuae/falcon-180B-chat":
        hfg = HuggingfaceGenerator(
            model_name=llm_name, temperature=0, max_new_tokens=512,
        )
        return hfg.generate

    # if llm_name == "GeneZC/MiniChat-3B":
    #     hfg = HuggingfaceGenerator(
    #         model_name=llm_name, temperature=0, max_new_tokens=250, stream=False,
    #     )
    #     return hfg.generate
    raise ValueError('Unknown LLM name')