import gradio as gr
import pandas as pd
from utils import load_all_results


def filter_and_sort_leaderboard(category_filters, sort_by="Video-MME", descending=True):
    df = load_all_results()
    filtered_df = df[df["Category"].isin(category_filters)]
    filtered_df = filtered_df.sort_values(by=sort_by, ascending=not descending).reset_index(drop=True)
    filtered_df.insert(0, "Rank", filtered_df.index + 1)
    return filtered_df.fillna("-")


def render_html_leaderboard(df):
    def row_color(category):
        if category == "Proprietary Models":
            return "background-color: rgba(255, 255, 153, 0.3);"
        elif category == "Open-source Transformer-based LMMs":
            return "background-color: rgba(153, 255, 204, 0.3);"
        elif category == "Open-source Efficient LMMs":
            return "background-color: rgba(204, 229, 255, 0.3);"
        return ""

    table_html = "<style>table { border-collapse: collapse; width: 100%; } th, td { border: 1px solid #ccc; padding: 6px; text-align: center; }</style>"
    table_html += "<table><thead><tr>" + "".join(f"<th>{col}</th>" for col in df.columns) + "</tr></thead><tbody>"
    for _, row in df.iterrows():
        style = row_color(row['Category'])
        row_html = "".join(f"<td>{row[col]}</td>" for col in df.columns)
        table_html += f"<tr style='{style}'>{row_html}</tr>"
    table_html += "</tbody></table>"
    return table_html


def build_interface():
    with gr.Blocks() as demo:
        gr.Markdown("## 🎥 LVU VLM Leaderboard")
        gr.Markdown("Benchmarks: Video-MME, MLVU, LVBench, LongVideoBench")

        with gr.Row():
            sort_metric = gr.Dropdown(choices=["Video-MME", "MLVU", "LVBench", "LongVideoBench"], label="Sort by", value="Video-MME")
            descending = gr.Checkbox(label="Descending", value=True)

        with gr.Row():
            proprietary = gr.Checkbox(label="Proprietary Models", value=True)
            transformer = gr.Checkbox(label="Open-source Transformer-based LMMs", value=True)
            efficient = gr.Checkbox(label="Open-source Efficient LMMs", value=True)

        default_df = filter_and_sort_leaderboard(
            ["Proprietary Models", "Open-source Transformer-based LMMs", "Open-source Efficient LMMs"], "Video-MME", True
        )
        leaderboard = gr.HTML(value=render_html_leaderboard(default_df))

        def update_board(sort_by, descending, proprietary, transformer, efficient):
            filters = []
            if proprietary:
                filters.append("Proprietary Models")
            if transformer:
                filters.append("Open-source Transformer-based LMMs")
            if efficient:
                filters.append("Open-source Efficient LMMs")
            df = filter_and_sort_leaderboard(filters, sort_by, descending)
            return render_html_leaderboard(df)

        inputs = [sort_metric, descending, proprietary, transformer, efficient]
        for inp in inputs:
            inp.change(fn=update_board, inputs=inputs, outputs=leaderboard)

    return demo


if __name__ == "__main__":
    demo = build_interface()
    demo.launch()