import json
import re
import argparse
import gradio as gr

# Load the JSONL file
def load_jsonl(file_path):
    data = []
    with open(file_path, 'r') as f:
        for line in f:
            data.append(json.loads(line))
    return data

def display_pairwise_answer(data):
    chat_mds = pairwise_to_gradio_chat_mds(data)
    return chat_mds

newline_pattern1 = re.compile("\n\n(\d+\. )")
newline_pattern2 = re.compile("\n\n(- )")

def post_process_answer(x):
    return x

def pairwise_to_gradio_chat_mds(data):
    end = data["turn"] * 3
    ans_a = data["conversation_a"]
    ans_b = data["conversation_b"]

    mds = [""] * end
    base = 0
    for i in range(0, end, 3):
        mds[i] = "## User Prompt\n" + data["conversation_a"][base]["content"].strip()
        mds[i + 1] = f"## {data['model_a']}\n" + post_process_answer(ans_a[base + 1]["content"].strip())
        mds[i + 2] = f"## {data['model_b']}\n" + post_process_answer(ans_b[base + 1]["content"].strip())
        base += 2

    if data["winner"] == "tie":
        winner = "tie"
    elif data["winner"] == "model_a":
        winner = "gemini" if "gemini" in data["model_a"] else "opponent"
    else:
        winner = "gemini" if "gemini" in data["model_b"] else "opponent"

    mds += [f"## Winner: {winner}"]
    mds += [""] * (16 - len(mds))
    return mds

# Filtering functions
def filter_by_language(language):
    return [item for item in data if item['language'] == language]

def filter_by_winner(winner_model, filtered_data):
    if winner_model == "anyone":
        return [item for item in filtered_data]
    return [item for item in filtered_data if item['winner'] == winner_model]

def filter_by_conversation_a_prefix(prefix, filtered_data):
    return [item for item in filtered_data if item['conversation_a'][0]["content"][:128] == prefix]

# Create Gradio interface
def update_winner_and_questions(language):
    filtered_data = filter_by_language(language)
    winners = ["anyone"] + sorted(set(item['winner'] for item in filtered_data))
    prefixes = [item['conversation_a'][0]["content"][:128] for item in filtered_data]
    return gr.update(choices=winners, value=winners[0]), gr.update(choices=prefixes, value=prefixes[0])

def update_question_options(language, winner_model):
    filtered_data = filter_by_language(language)
    filtered_data = filter_by_winner(winner_model, filtered_data)
    prefixes = [item['conversation_a'][0]["content"][:128] for item in filtered_data]
    return gr.update(choices=prefixes, value=prefixes[0])

def display_filtered_data(language, winner_model, prefix):
    filtered_data = filter_by_language(language)
    filtered_data = filter_by_winner(winner_model, filtered_data)
    filtered_data = filter_by_conversation_a_prefix(prefix, filtered_data)
    if len(filtered_data) == 0:
        return [""] * 16
    return pairwise_to_gradio_chat_mds(filtered_data[0])

def next_question(language, winner_model, prefix):
    filtered_data = filter_by_language(language)
    filtered_data = filter_by_winner(winner_model, filtered_data)
    all_items = [item['conversation_a'][0]["content"][:128] for item in filtered_data]
    if prefix:
        i = all_items.index(prefix) + 1
    else:
        i = 0

    if i >= len(all_items):
        return gr.update(choices=all_items, value=all_items[-1])
    return gr.update(choices=all_items, value=all_items[i])

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", type=str, default="0.0.0.0")
    parser.add_argument("--port", type=int)
    parser.add_argument("--share", action="store_true")
    args = parser.parse_args()

    data = load_jsonl('gemini_battles.jsonl')

    default_lang = "English"
    filter_data = filter_by_language(language=default_lang)
    question_prefixes = [item['conversation_a'][0]["content"][:128] for item in filter_data]
    default_question = question_prefixes[2]

    with gr.Blocks() as demo:
        gr.Markdown(value="# Welcome to gemini-1.5-pro-api-0514 battles")
        with gr.Row():
            with gr.Column():
                languages = ["English"] + list(sorted(set([item['language'] for item in data if item['language'] != "English"])))
                language_dropdown = gr.Dropdown(label="Select Language", choices=languages, value=default_lang)
            with gr.Column():
                winners = ["anyone"] + sorted(set(item['winner'] for item in filter_data))
                winner_dropdown = gr.Dropdown(label="Winner Model", choices=winners, value="anyone")

        with gr.Row():
            with gr.Column(scale=5):
                question_dropdown = gr.Dropdown(label="Select Question", choices=question_prefixes, value=default_question)
            with gr.Column():
                next_button = gr.Button("Next Question")

        default_chat_mds = display_filtered_data(default_lang, "anyone", default_question)
        chat_mds = []
        for i in range(5):
            chat_mds.append(gr.Markdown(elem_id=f"user_question_{i+1}", value=default_chat_mds[len(chat_mds)]))
            with gr.Row():
                for j in range(2):
                    with gr.Column(scale=100):
                        chat_mds.append(gr.Markdown(value=default_chat_mds[len(chat_mds)]))
                    if j == 0:
                        with gr.Column(scale=1, min_width=8):
                            gr.Markdown()
        chat_mds.append(gr.Markdown())

        language_dropdown.change(fn=update_winner_and_questions, inputs=language_dropdown, outputs=[winner_dropdown, question_dropdown])
        winner_dropdown.change(fn=update_question_options, inputs=[language_dropdown, winner_dropdown], outputs=question_dropdown)
        next_button.click(fn=next_question, inputs=[language_dropdown, winner_dropdown, question_dropdown], outputs=question_dropdown)
        question_dropdown.change(fn=display_filtered_data, inputs=[language_dropdown, winner_dropdown, question_dropdown], outputs=chat_mds)

    demo.launch(share=args.share)