app.py

from __future__ import annotations

import os

import gradio as gr
import pandas as pd

from gaia_workbench import (
    API_URL,
    answer_payload,
    download_files,
    fetch_questions,
    load_answers,
    load_questions,
    status_rows,
    submit_answers,
    write_answer,
)


def refresh_questions() -> tuple[str, pd.DataFrame]:
    questions = fetch_questions()
    return f"Fetched {len(questions)} questions from {API_URL}", pd.DataFrame(status_rows())


def refresh_status() -> pd.DataFrame:
    return pd.DataFrame(status_rows())


def save_answer(task_id: str, answer: str) -> tuple[str, pd.DataFrame]:
    if not task_id.strip():
        return "Missing task_id", refresh_status()
    if not answer.strip():
        return "Missing answer", refresh_status()
    write_answer(task_id.strip(), answer.strip())
    return f"Saved answer for {task_id.strip()}", refresh_status()


def download_attached_files() -> str:
    downloaded, failed = download_files()
    parts = []
    if downloaded:
        parts.append("Downloaded:\n" + "\n".join(str(path) for path in downloaded))
    if failed:
        parts.append("Failed:\n" + "\n".join(failed))
    return "\n\n".join(parts) if parts else "No attached files found."


def _default_agent_code() -> str:
    space_id = os.getenv("SPACE_ID", "").strip()
    return f"https://huggingface.co/spaces/{space_id}/tree/main" if space_id else ""


def _submit_with(username: str, agent_code: str) -> tuple[str, pd.DataFrame]:
    agent_code = agent_code.strip() or _default_agent_code()

    try:
        result = submit_answers(username=username, agent_code=agent_code)
    except Exception as exc:
        return f"Submission failed: {exc}", refresh_status()

    return (
        "Submission successful\n"
        f"User: {result.get('username')}\n"
        f"Score: {result.get('score')}%\n"
        f"Correct: {result.get('correct_count')}/{result.get('total_attempted')}\n"
        f"Message: {result.get('message')}",
        refresh_status(),
    )


def submit_logged_in(profile: gr.OAuthProfile | None) -> tuple[str, pd.DataFrame]:
    if not profile:
        return "Please login with Hugging Face first, or use manual submit.", refresh_status()
    return _submit_with(profile.username, _default_agent_code())


def submit_manual(manual_username: str, manual_agent_code: str) -> tuple[str, pd.DataFrame]:
    return _submit_with(manual_username.strip(), manual_agent_code.strip())


with gr.Blocks(title="GAIA Unit 4 Workbench") as demo:
    gr.Markdown("# GAIA Unit 4 Workbench")
    gr.Markdown(
        "Fetch the course questions, fill exact-match answers, then submit. "
        "On Hugging Face Spaces, login with the button before submission."
    )

    is_space_runtime = bool(os.getenv("SPACE_ID", "").strip())
    if is_space_runtime:
        gr.LoginButton()

    with gr.Row():
        fetch_button = gr.Button("Fetch Questions")
        download_button = gr.Button("Download Attached Files")
        refresh_button = gr.Button("Refresh Status")

    status = gr.Textbox(label="Status", lines=5)
    table = gr.DataFrame(label="Questions / Answers", wrap=True, value=pd.DataFrame(status_rows()))

    with gr.Row():
        task_id_input = gr.Textbox(label="Task ID")
        answer_input = gr.Textbox(label="Exact Answer")
    save_button = gr.Button("Save Answer")

    gr.Markdown("## Submit")
    manual_username = gr.Textbox(label="HF Username fallback", placeholder="Used only when not logged in via HF OAuth")
    manual_agent_code = gr.Textbox(
        label="Agent Code URL fallback",
        placeholder="https://huggingface.co/spaces/YOUR_USERNAME/YOUR_SPACE/tree/main",
    )
    with gr.Row():
        submit_login_button = gr.Button("Submit With HF Login", visible=is_space_runtime)
        submit_manual_button = gr.Button("Submit Manual")

    fetch_button.click(fn=refresh_questions, outputs=[status, table])
    download_button.click(fn=download_attached_files, outputs=status)
    refresh_button.click(fn=refresh_status, outputs=table)
    save_button.click(fn=save_answer, inputs=[task_id_input, answer_input], outputs=[status, table])
    submit_login_button.click(fn=submit_logged_in, outputs=[status, table])
    submit_manual_button.click(fn=submit_manual, inputs=[manual_username, manual_agent_code], outputs=[status, table])


if __name__ == "__main__":
    print(f"Loaded {len(load_questions())} questions and {len(load_answers())} local answers.")
    print(f"Prepared {len(answer_payload())} answers for submission.")
    demo.launch()