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.")

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

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

    metadata = {
        "submitted_by": submitted_by.strip(),
        "model_id": model_id.strip()
    }
    leaderboard_df = server.get_leaderboard()
    if len(leaderboard_df) > 0:
        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)

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

def get_leaderboard_df():
    return server.get_leaderboard()

# 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=get_leaderboard_df,
                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...")

            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],
            ).then(
                lambda: server.get_leaderboard(),  # Show loading
                outputs=leaderboard_output
            )

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