import gradio as gr
from utils import MEGABenchEvalDataLoader
import os
from constants import *
from urllib.parse import quote

# ------------------------------------------------------------------------------
# Load CSS files
# ------------------------------------------------------------------------------
current_dir = os.path.dirname(os.path.abspath(__file__))
base_css_file = os.path.join(current_dir, "static", "css", "style.css")
table_css_file = os.path.join(current_dir, "static", "css", "table.css")

with open(base_css_file, "r") as f:
    base_css = f.read()
with open(table_css_file, "r") as f:
    table_css = f.read()

# ------------------------------------------------------------------------------
# Initialize data loaders
# ------------------------------------------------------------------------------
default_loader = MEGABenchEvalDataLoader("./static/eval_results/Default")
si_loader = MEGABenchEvalDataLoader("./static/eval_results/SI")

# ------------------------------------------------------------------------------
# Helper Functions
# ------------------------------------------------------------------------------

def generate_search_links(model_name):
    """For a given model name, generate a set of search links as HTML."""
    search_urls = {
        "📚📖ArXiv": lambda k: f"https://arxiv.org/search/?query={quote(k)}&searchtype=all",
        "🔮Google": lambda k: f"https://www.google.com/search?q={quote(k)}",
        "📺Youtube": lambda k: f"https://www.youtube.com/results?search_query={quote(k)}",
        "🔭Bing": lambda k: f"https://www.bing.com/search?q={quote(k)}",
        "💡Truth": lambda k: f"https://truthsocial.com/search?q={quote(k)}",
        "📱X": lambda k: f"https://twitter.com/search?q={quote(k)}",
    }
    # Build a set of inline HTML links (each opens in a new tab)
    links = " ".join(
        [f'<a href="{url(model_name)}" target="_blank">{emoji}</a>' for emoji, url in search_urls.items()]
    )
    return links

def add_search_links_to_table(headers, data):
    """
    Append a "Search Links" column to the table.
    (Assumes that each row’s column index 1 holds the model name.)
    """
    new_headers = headers.copy()
    new_headers.append("Search Links")
    new_data = []
    for row in data:
        new_row = row.copy()
        # Assume the model name is in the second column (index 1)
        model_name = new_row[1] if len(new_row) > 1 else ""
        new_row.append(generate_search_links(model_name))
        new_data.append(new_row)
    return new_headers, new_data

def clean_choice(choice):
    """
    Remove a leading emoji and space (if present) from a choice string.
    For example, "📊 Default" becomes "Default".
    """
    parts = choice.split(" ", 1)
    return parts[1] if len(parts) > 1 else choice

def update_table_and_caption(table_type, super_group, model_group):
    """
    Called when any selector changes. Cleans the emoji‐prefixed values, loads new data,
    appends a Search Links column, and returns a new Dataframe component (with a smaller max height),
    an updated caption, and the CSS style.
    """
    table_type_clean = clean_choice(table_type)
    super_group_clean = clean_choice(super_group)
    model_group_clean = clean_choice(model_group)
    
    if table_type_clean == "Default":
        headers, data = default_loader.get_leaderboard_data(super_group_clean, model_group_clean)
        caption = default_caption
    else:  # "Single Image"
        headers, data = si_loader.get_leaderboard_data(super_group_clean, model_group_clean)
        caption = single_image_caption
        
    # Append search links column to the table data
    headers, data = add_search_links_to_table(headers, data)
    n = len(headers)
    # Assume first column is a number, second (model name) is HTML, the intermediate ones are numbers, and the last (links) is HTML
    datatype = ["number", "html"] + ["number"] * (n - 3) + ["html"]
    # Adjust column widths as needed (example: first two fixed, last a bit wider)
    column_widths = ["100px", "240px"] + ["160px"] * (n - 3) + ["210px"]
    
    dataframe_component = gr.Dataframe(
        value=data,
        headers=headers,
        datatype=datatype,
        interactive=False,
        column_widths=column_widths,
        max_height=600,  # smaller height to show the full table in a compact area
        elem_classes="custom-dataframe"
    )
    
    return [dataframe_component, caption, f"<style>{base_css}\n{table_css}</style>"]

def update_selectors(table_type):
    """
    When the table selector changes, update the other radio choices.
    (Also adds an emoji prefix to each choice.)
    """
    table_type_clean = clean_choice(table_type)
    loader = default_loader if table_type_clean == "Default" else si_loader
    super_group_choices = [f"🔍 {group}" for group in list(loader.SUPER_GROUPS.keys())]
    model_group_choices = [f"🤖 {group}" for group in list(loader.MODEL_GROUPS.keys())]
    return [super_group_choices, model_group_choices]

# ------------------------------------------------------------------------------
# Build Gradio App Layout
# ------------------------------------------------------------------------------
with gr.Blocks() as block:
    # Add CSS via an invisible HTML component
    css_style = gr.HTML(f"<style>{base_css}\n{table_css}</style>", visible=False)
    
    # NOTE: The original top-level introduction markdown has been removed here.
    # It is now placed at the bottom of the MEGA-Bench tab in a collapsed Accordion.
    with gr.Tabs(elem_classes="tab-buttons") as tabs:
    
        # -------------------- Tab 1: MEGA-Bench --------------------
        with gr.TabItem("📊 MEGA-Bench", elem_id="qa-tab-table1", id=1):
            # A Citation accordion (collapsed by default)
            with gr.Row():
                with gr.Accordion("Citation", open=False):
                    citation_button = gr.Textbox(
                        value=CITATION_BUTTON_TEXT,
                        label=CITATION_BUTTON_LABEL,
                        elem_id="citation-button",
                        lines=10,
                    )
            gr.Markdown(TABLE_INTRODUCTION)
            
            # ---- Top-left “button‐bar” of radio selectors with emoji labels ----
            with gr.Row():
                table_selector = gr.Radio(
                    choices=["📊 Default", "🖼️ Single Image"],
                    label="Select table to display",
                    value="📊 Default"
                )
            with gr.Row():
                super_group_selector = gr.Radio(
                    choices=[f"🔍 {group}" for group in list(default_loader.SUPER_GROUPS.keys())],
                    label="Select a dimension",
                    value=f"🔍 {list(default_loader.SUPER_GROUPS.keys())[0]}"
                )
                model_group_selector = gr.Radio(
                    choices=[f"🤖 {group}" for group in list(BASE_MODEL_GROUPS.keys())],
                    label="Select a model group",
                    value="🤖 All"
                )
            
            # A caption component for the table
            caption_component = gr.Markdown(
                value=default_caption,
                elem_classes="table-caption",
                latex_delimiters=[{"left": "$", "right": "$", "display": False}],
            )
            
            # ---- Initial Dataframe (with search links appended) ----
            initial_headers, initial_data = default_loader.get_leaderboard_data(
                list(default_loader.SUPER_GROUPS.keys())[0], "All"
            )
            initial_headers, initial_data = add_search_links_to_table(initial_headers, initial_data)
            n = len(initial_headers)
            initial_datatype = ["number", "html"] + ["number"] * (n - 3) + ["html"]
            initial_column_widths = ["100px", "240px"] + ["160px"] * (n - 3) + ["210px"]
            
            data_component = gr.Dataframe(
                value=initial_data,
                headers=initial_headers,
                datatype=initial_datatype,
                interactive=False,
                elem_classes="custom-dataframe",
                max_height=600,
                column_widths=initial_column_widths
            )
            
            # ---- Controls to update the table ----
            refresh_button = gr.Button("Refresh")
            refresh_button.click(
                fn=update_table_and_caption, 
                inputs=[table_selector, super_group_selector, model_group_selector], 
                outputs=[data_component, caption_component, css_style]
            )
            super_group_selector.change(
                fn=update_table_and_caption, 
                inputs=[table_selector, super_group_selector, model_group_selector], 
                outputs=[data_component, caption_component, css_style]
            )
            model_group_selector.change(
                fn=update_table_and_caption, 
                inputs=[table_selector, super_group_selector, model_group_selector],
                outputs=[data_component, caption_component, css_style]
            )
            table_selector.change(
                fn=update_selectors,
                inputs=[table_selector],
                outputs=[super_group_selector, model_group_selector]
            ).then(
                fn=update_table_and_caption,
                inputs=[table_selector, super_group_selector, model_group_selector],
                outputs=[data_component, caption_component, css_style]
            )
            
            # ---- Move the introductory text to a collapsed accordion at the bottom ----
            with gr.Accordion("Introduction", open=False):
                gr.Markdown(LEADERBOARD_INTRODUCTION)
                
        # -------------------- Tab 2: Data Information --------------------
        with gr.TabItem("📝 Data Information", elem_id="qa-tab-table2", id=2):
            gr.Markdown(DATA_INFO, elem_classes="markdown-text")
            
        # -------------------- Tab 3: Submit --------------------
        with gr.TabItem("🚀 Submit", elem_id="submit-tab", id=3):
            with gr.Row():
                gr.Markdown(SUBMIT_INTRODUCTION, elem_classes="markdown-text")
                
if __name__ == "__main__":
    block.launch(share=True)