# -*- coding: utf-8 -*-
import time
import gradio as gr
from pipeline import Llama3
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('-m', '--model_path', type=str, required=True, help='path to the bmodel file')
parser.add_argument('-t', '--tokenizer_path', type=str, default="../support/token_config", help='path to the tokenizer file')
parser.add_argument('-d', '--devid', type=str, default='0', help='device ID to use')
parser.add_argument('--enable_history', type=bool, default=True, help="if set, enables storing of history memory.")
parser.add_argument('--temperature', type=float, default=1.0, help='temperature scaling factor for the likelihood distribution')
parser.add_argument('--top_p', type=float, default=1.0, help='cumulative probability of token words to consider as a set of candidates')
parser.add_argument('--repeat_penalty', type=float, default=1.0, help='penalty for repeated tokens')
parser.add_argument('--repeat_last_n', type=int, default=32, help='repeat penalty for recent n tokens')
parser.add_argument('--max_new_tokens', type=int, default=1024, help='max new token length to generate')
parser.add_argument('--generation_mode', type=str, choices=["greedy", "penalty_sample"], default="greedy", help='mode for generating next token')
parser.add_argument('--prompt_mode', type=str, choices=["prompted", "unprompted"], default="prompted", help='use prompt format or original input')
parser.add_argument('--decode_mode', type=str, default="basic", choices=["basic", "jacobi"], help='mode for decoding')
args = parser.parse_args()

model = Llama3(args)


def gr_update_history():
    if model.model.token_length >= model.SEQLEN:
        # print("... (reach the maximal length)", flush=True, end='')
        gr.Warning("reach the maximal length, Llama3 would clear all history record")
        model.history = [model.system]
    else:
        model.history.append({"role": "assistant", "content": model.answer_cur})


def gr_user(user_input, history):
    model.input_str = user_input
    return "", history + [[user_input, None]]


def gr_chat(history):
    """
    Stream the prediction for the given query.
    """
    tokens = model.encode_tokens()

    # check tokens
    if not tokens:
        gr.Warning("Sorry: your question is empty!!")
        return
    if len(tokens) > model.SEQLEN:
        gr.Warning(
            "The maximum question length should be shorter than {} but we get {} instead.".format(
                model.SEQLEN, len(tokens)
            )
        )
        gr_update_history()

    model.answer_cur = ""
    model.answer_token = []
    token_num = 0

    first_start = time.time()
    token = model.model.forward_first(tokens)
    first_end = time.time()

    history[-1][1] = ""
    full_word_tokens = []

    while token not in model.EOS and model.model.token_length < model.SEQLEN:
        full_word_tokens.append(token)
        t_word = model.tokenizer.decode(full_word_tokens, skip_special_tokens=True)

        if "�" in t_word:
            token = model.model.forward_next()
            token_num += 1
            continue

        model.answer_token += full_word_tokens

        history[-1][1] += t_word
        full_word_tokens = []
        yield history
        token = model.model.forward_next()
        token_num += 1

    next_end = time.time()
    first_duration = first_end - first_start
    next_duration = next_end - first_end
    tps = token_num / next_duration

    print()
    print(f"FTL: {first_duration:.3f} s")
    print(f"TPS: {tps:.3f} token/s")

    model.answer_cur = model.tokenizer.decode(model.answer_token)
    gr_update_history()


def reset():
    model.clear()
    return [[None, None]]



description = """
# Llama3 TPU 🏁 
"""
with gr.Blocks() as demo:
    gr.Markdown(description)
    with gr.Row():
        with gr.Column():
            chatbot = gr.Chatbot(label="Llama3", height=1050)

            with gr.Row():
                user_input = gr.Textbox(show_label=False, placeholder="Ask Llama3", lines=1, min_width=300, scale=6)
                submitBtn = gr.Button("Submit", variant="primary", scale=1)
                emptyBtn = gr.Button(value="Clear", scale=1)

    user_input.submit(gr_user, [user_input, chatbot], [user_input, chatbot]).then(gr_chat, chatbot, chatbot)
    # clear_history.
    submitBtn.click(gr_user, [user_input, chatbot], [user_input, chatbot]).then(gr_chat, chatbot, chatbot)

    emptyBtn.click(reset, outputs=[chatbot])

demo.queue(max_size=20).launch(share=False, server_name="0.0.0.0", inbrowser=True, server_port=8003)