app.py

import os

import gradio as gr
from gradio_modal import Modal

from content import HEADER_MARKDOWN, LEADERBOARD_TAB_TITLE_MARKDOWN, SUBMISSION_TAB_TITLE_MARKDOWN
from leaderboard_server import LeaderboardServer

# Initialize server and task list
server = LeaderboardServer()
TASKS = list(server.tasks_metadata.keys())
EXPECTED_TOKEN = os.environ.get("SUBMISSION_TOKEN")


def get_datasets_for_task(task):
    path = os.path.join("references", task)
    if not os.path.exists(path):
        return []
    return [f.replace(".json", "") for f in os.listdir(path) if f.endswith(".json")]


def update_datasets(task):
    return gr.CheckboxGroup.update(choices=get_datasets_for_task(task), value=[])


def submit_model(task, datasets, hyp_file, submitted_by, model_id, token):
    if not hyp_file:
        return gr.update(visible=True, value="⚠️ Please upload a hypothesis file."), gr.update(), gr.update(selected=1)

    if not submitted_by.strip() or not model_id.strip() or not token.strip():
        return gr.update(visible=True, value="⚠️ All fields are required."), gr.update(), gr.update(selected=1)

    if token.strip() != EXPECTED_TOKEN:
        return gr.update(visible=True, value="❌ Invalid submission token."), gr.update(), gr.update(selected=1)

    metadata = {
        "submitted_by": submitted_by.strip(),
        "model_id": model_id.strip()
    }
    leaderboard_df = server.get_leaderboard()
    existing = leaderboard_df[
        (leaderboard_df["Submitted by"] == submitted_by.strip()) &
        (leaderboard_df["Model ID"] == model_id.strip())
    ]
    if not existing.empty:
        return gr.update(value="❌ A submission with this name and model ID already exists.", visible=True), \
            gr.update(), gr.update(selected=1), gr.update(visible=False)

    try:
        server.prepare_model_for_submission(hyp_file.name, metadata, task, datasets)
        server.update_leaderboard()
        leaderboard_df = server.get_leaderboard()
        return gr.update(visible=True, value="✅ Submission successful!"), leaderboard_df, gr.update(selected=0)
    except Exception as e:
        return gr.update(visible=True, value=f"❌ Error: {str(e)}"), gr.update(), gr.update(selected=1)


# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown(HEADER_MARKDOWN)

    with gr.Tabs(selected=0) as tabs:
        with gr.Tab("📈 Leaderboard"):
            gr.Markdown(LEADERBOARD_TAB_TITLE_MARKDOWN)
            leaderboard_output = gr.Dataframe(
                value=server.get_leaderboard(),
                interactive=False,
                label="Leaderboard"
            )

        with gr.Tab("📤 Submit"):
            gr.Markdown(SUBMISSION_TAB_TITLE_MARKDOWN)

            with gr.Row():
                task_dropdown = gr.Dropdown(choices=TASKS, value=TASKS[0], label="Select Task")
                dataset_checkboxes = gr.CheckboxGroup(choices=get_datasets_for_task(TASKS[0]), label="Select Datasets")
            task_dropdown.change(fn=update_datasets, inputs=task_dropdown, outputs=dataset_checkboxes)

            with gr.Row():
                submitted_by_input = gr.Text(label="Submitted by")
                model_id_input = gr.Text(label="Model Identifier")
                token_input = gr.Text(label="Submission Token", type="password")

            hyp_file_upload = gr.File(label="Upload Hypothesis JSON", file_types=[".json"])
            submit_btn = gr.Button("Submit")
            with Modal("Submission Feedback", visible=False) as loading_msg:
                feedback_text = gr.Text(visible=True, label="⏳ Processing your submission...")

            with Modal("Submission Feedback", visible=False) as modal:
                feedback_text = gr.Text(visible=True, label="")

            submit_btn.click(
                lambda: gr.update(visible=True),  # Show loading
                outputs=loading_msg
            ).then(
                fn=submit_model,
                inputs=[task_dropdown, dataset_checkboxes, hyp_file_upload, submitted_by_input, model_id_input,
                        token_input],
                outputs=[feedback_text, leaderboard_output, tabs],
                show_progress=True
            ).then(
                lambda: gr.update(visible=False),
                outputs=loading_msg
            ).then(
                lambda: gr.update(visible=True),
                outputs=modal
            )
            # submit_btn.click(
            #     fn=submit_model,
            #     inputs=[
            #         task_dropdown, dataset_checkboxes, hyp_file_upload,
            #         submitted_by_input, model_id_input, token_input
            #     ],
            #     outputs=[feedback_text, leaderboard_output, tabs],
            #     show_progress=True
            # ).then(
            #     lambda: gr.update(visible=True),
            #     outputs=modal
            # )

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