import json
import os
import gradio as gr
from huggingface_hub import InferenceClient
from eval_modules.utils import calc_bleu_rouge_scores
from eval_modules.calc_repetitions_v2e import detect_repetitions

questions_file_path = os.getenv("QUESTIONS_FILE_PATH") or "./ms_macro.json"

questions = json.loads(open(questions_file_path).read())
examples = [[question["question"].strip()] for question in questions]
print(f"Loaded {len(examples)} examples")

qa_system_prompt = "Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer."

"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""
# client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
# client = InferenceClient("HuggingFaceH4/zephyr-7b-gemma-v0.1")
client = InferenceClient("microsoft/Phi-3.5-mini-instruct")


def chat(
    message,
    history: list[tuple[str, str]],
    system_message,
    temperature=0,
    frequency_penalty=0,
    presence_penalty=0,
    max_tokens=256,
    top_p=0.95,
):
    chat = []
    for item in history:
        chat.append({"role": "user", "content": item[0]})
        if item[1] is not None:
            chat.append({"role": "assistant", "content": item[1]})

    index = -1
    if [message] in examples:
        index = examples.index([message])
        message = f"{qa_system_prompt}\n\n{questions[index]['context']}\n\nQuestion: {message}"
        print("RAG prompt:", message)

    chat.append({"role": "user", "content": message})

    messages = [{"role": "system", "content": system_message}]
    messages.append({"role": "user", "content": message})

    partial_text = ""

    finish_reason = None
    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        frequency_penalty=None,  # frequency_penalty,
        presence_penalty=None,  # presence_penalty,
        top_p=top_p,
        seed=42,
    ):
        finish_reason = message.choices[0].finish_reason
        # print("finish_reason:", finish_reason)

        if finish_reason is None:
            new_text = message.choices[0].delta.content
            partial_text += new_text
            yield partial_text
        else:
            break

    answer = partial_text
    (whitespace_score, repetition_score, total_repetitions) = detect_repetitions(answer)
    partial_text += "\n\nRepetition Metrics:\n"
    partial_text += f"1. Whitespace Score: {whitespace_score:.3f}\n"
    partial_text += f"1. Repetition Score: {repetition_score:.3f}\n"
    partial_text += f"1. Total Repetitions: {total_repetitions:.3f}\n"
    partial_text += (
        f"1. Non-Repetitive Ratio: {1 - total_repetitions / len(answer):.3f}\n"
    )

    if index >= 0:  # RAG
        key = (
            "wellFormedAnswers"
            if "wellFormedAnswers" in questions[index]
            else "answers"
        )
        scores = calc_bleu_rouge_scores([answer], [questions[index][key]], debug=True)

        partial_text += "\n\n Performance Metrics:\n"
        partial_text += f'1. BLEU-1: {scores["bleu_scores"]["bleu"]:.3f}\n'
        partial_text += f'1. RougeL: {scores["rouge_scores"]["rougeL"]:.3f}\n'

        partial_text += f"\n\nGround truth: {questions[index][key][0]}\n"

    partial_text += f"\n\nThe text generation has ended because: {finish_reason}\n"

    yield partial_text


demo = gr.ChatInterface(
    fn=chat,
    examples=examples,
    cache_examples=False,
    additional_inputs_accordion=gr.Accordion(
        label="⚙️ Parameters", open=False, render=False
    ),
    additional_inputs=[
        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
        gr.Slider(
            minimum=0, maximum=2, step=0.1, value=0, label="Temperature", render=False
        ),
        gr.Slider(
            minimum=-2,
            maximum=2,
            step=0.1,
            value=0,
            label="Frequency Penalty",
            render=False,
        ),
        gr.Slider(
            minimum=-2,
            maximum=2,
            step=0.1,
            value=0,
            label="Presence Penalty",
            render=False,
        ),
        gr.Slider(
            minimum=128,
            maximum=4096,
            step=1,
            value=512,
            label="Max new tokens",
            render=False,
        ),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
)
demo.launch()