from __future__ import annotations

import gradio as gr
import pandas as pd

from .leaderboard_category import (
    filter_leaderboard,
    style_leaderboard,
)


TASK_DISPLAY_COLUMNS = ["Org", "Model", "Size", "Overall", "MCQ", "Y/N", "Open"]


def render_task_leaderboard(
    leaderboard_df: pd.DataFrame,
    search_query: str,
    selected_split: str,
    selected_types: list[str],
    selected_size_ranges: list[str],
    selected_families: list[str],
):
    filtered = filter_leaderboard(
        leaderboard_df,
        TASK_DISPLAY_COLUMNS,
        search_query,
        selected_split,
        selected_types,
        selected_size_ranges,
        selected_families,
    )
    return style_leaderboard(filtered, TASK_DISPLAY_COLUMNS)


def update_split_task_leaderboards(
    leaderboard_df: pd.DataFrame,
    search_query: str,
    selected_types: list[str],
    selected_size_ranges: list[str],
    selected_families: list[str],
):
    return (
        render_task_leaderboard(
            leaderboard_df,
            search_query,
            "Test",
            selected_types,
            selected_size_ranges,
            selected_families,
        ),
        render_task_leaderboard(
            leaderboard_df,
            search_query,
            "Validation",
            selected_types,
            selected_size_ranges,
            selected_families,
        ),
    )


def render_leaderboard_task_tab(
    leaderboard_df: pd.DataFrame,
    search_box,
    type_filter,
    size_filter,
    family_filter,
) -> None:
    initial_test_table_value = render_task_leaderboard(
        leaderboard_df,
        "",
        "Test",
        [],
        [],
        [],
    )
    initial_validation_table_value = render_task_leaderboard(
        leaderboard_df,
        "",
        "Validation",
        [],
        [],
        [],
    )
    with gr.Tabs(elem_id="leaderboard-task-split-tabs", elem_classes=["leaderboard-split-tabs"]):
        with gr.Tab("Test"):
            test_leaderboard_table = gr.Dataframe(
                value=initial_test_table_value,
                max_height=760,
                interactive=False,
                wrap=True,
                buttons=[],
                elem_id="leaderboard-table-test",
                elem_classes=["leaderboard-table"],
            )
        with gr.Tab("Validation"):
            validation_leaderboard_table = gr.Dataframe(
                value=initial_validation_table_value,
                max_height=760,
                interactive=False,
                wrap=True,
                buttons=[],
                elem_id="leaderboard-table-validation",
                elem_classes=["leaderboard-table"],
            )
    search_box.submit(
        fn=lambda search_query, selected_type, selected_size, selected_family: update_split_task_leaderboards(
            leaderboard_df,
            search_query,
            selected_type,
            selected_size,
            selected_family,
        ),
        inputs=[search_box, type_filter, size_filter, family_filter],
        outputs=[test_leaderboard_table, validation_leaderboard_table],
    )

    for component in (type_filter, size_filter, family_filter):
        component.change(
            fn=lambda search_query, selected_type, selected_size, selected_family: update_split_task_leaderboards(
                leaderboard_df,
                search_query,
                selected_type,
                selected_size,
                selected_family,
            ),
            inputs=[search_box, type_filter, size_filter, family_filter],
            outputs=[test_leaderboard_table, validation_leaderboard_table],
        )