import os

import gradio as gr
from datasets import load_dataset, Dataset
from huggingface_hub import whoami


EXAM_DATASET_ID = os.getenv("EXAM_DATASET_ID") or "burtenshaw/exam_questions"

ds = load_dataset(EXAM_DATASET_ID, split="train")

# Convert dataset to a list of dicts so we can iterate similarly to quiz_data
quiz_data = ds.to_pandas().to_dict("records")  # or use a for-loop if you prefer


def on_user_logged_in(token: gr.OAuthToken | None):
    """
    If the user has a valid token, hide the login button and show the Start button.
    Otherwise, keep the login button visible, hide Start.
    """
    if token is not None:
        return gr.update(visible=False), gr.update(visible=True)
    else:
        # Not logged in, keep the login visible, hide Start
        return gr.update(visible=True), gr.update(visible=False)


def push_results_to_hub(user_answers, token: gr.OAuthToken | None):
    """
    Create a new dataset from user_answers and push it to the Hub.
    We use the user's HF token to determine the correct repo.
    If no one is logged in, we'll return an error message.
    """
    if token is None:
        gr.Warning("Please log in to Hugging Face before pushing!")
        return
    else:
        gr.Info("Submitting answers to the Hub. Please wait...", duration=2)
    user_info = whoami(token=token.token)
    repo_id = f"{user_info['name']}/quiz-responses"  # e.g. 'myUsername/quiz-responses'

    new_ds = Dataset.from_list(user_answers)
    new_ds.push_to_hub(repo_id)
    gr.Success("Your responses have been submitted to the Hub!")


def handle_quiz(question_idx, user_answers, selected_answer, is_start):
    """
    A single function that handles both 'Start' and 'Next' logic:
      - If is_start=True, skip storing an answer and show the first question.
      - Otherwise, store the last answer and move on.
      - If we've reached the end, display results.
    """
    # Hide the start button once the first question is shown
    start_btn_update = gr.update(visible=False) if is_start else None

    # If this is the first time (start=True), begin at question_idx=0
    if is_start:
        question_idx = 0
    else:
        # If not the very first question, store the user's last selection
        if question_idx < len(quiz_data):
            current_q = quiz_data[question_idx]
            user_answers.append(
                {"question": current_q["question"], "selected_answer": selected_answer}
            )
        question_idx += 1

    # If we've reached the end, show final results
    if question_idx >= len(quiz_data):
        final_text = f"**All questions answered!**\n\nHere are your selections:\n\n{user_answers}"
        return (
            "",  # question_text becomes blank
            gr.update(choices=[], visible=False),
            "",  # status_text (can clear or reuse)
            question_idx,
            user_answers,
            start_btn_update,
            gr.update(value=final_text, visible=True),  # show final_markdown
        )
    else:
        # Otherwise, show the next question
        q = quiz_data[question_idx]
        updated_question = f"## **Question {question_idx + 1}**: {q['question']}"
        return (
            updated_question,
            gr.update(
                choices=[
                    q["answer_a"],
                    q["answer_b"],
                    q["answer_c"],
                    q["answer_d"],
                ],
                value=None,
                visible=True,
            ),
            "Select an answer and click 'Next' to continue.",
            question_idx,
            user_answers,
            start_btn_update,
            gr.update(visible=False),  # Hide final_markdown for now
        )


def success_message(response):
    # response is whatever push_results_to_hub returned
    return f"{response}\n\n**Success!**"


with gr.Blocks() as demo:
    demo.title = f"Dataset Quiz for {EXAM_DATASET_ID}"

    # Header
    gr.Markdown(f"## Welcome to the {EXAM_DATASET_ID} Quiz")
    gr.Markdown(
        "Log in first, then click 'Start' to begin. Answer each question, click 'Next', and finally click 'Submit' to publish your results to the Hugging Face Hub."
    )

    # Step 1: Login
    login_btn = gr.LoginButton()
    # We'll hide the Start button until user logs in
    start_btn = gr.Button("Start", visible=False)

    # State variables
    question_idx = gr.State(value=0)
    user_answers = gr.State(value=[])

    # We display question text with Markdown
    question_text = gr.Markdown("")
    status_text = gr.Markdown("")
    radio_choices = gr.Radio(label="Your Answer", choices=[], visible=False)

    # Final results after all questions are done
    final_markdown = gr.Markdown("", visible=False)

    next_btn = gr.Button("Next")
    submit_btn = gr.Button("Submit")

    # Use click() instead of login()
    login_btn.click(fn=on_user_logged_in, inputs=None, outputs=[login_btn, start_btn])

    # Click "Start" => show first question, hide Start button
    start_btn.click(
        fn=handle_quiz,
        inputs=[question_idx, user_answers, radio_choices, gr.State(True)],
        outputs=[
            question_text,
            radio_choices,
            status_text,
            question_idx,
            user_answers,
            start_btn,
            final_markdown,
        ],
    )

    # Click "Next" => store selection, move on
    next_btn.click(
        fn=handle_quiz,
        inputs=[question_idx, user_answers, radio_choices, gr.State(False)],
        outputs=[
            question_text,
            radio_choices,
            status_text,
            question_idx,
            user_answers,
            start_btn,
            final_markdown,
        ],
    )

    submit_btn.click(fn=push_results_to_hub, inputs=[user_answers])


if __name__ == "__main__":
    # Note: If testing locally, you'll need to run `huggingface-cli login` or set HF_TOKEN
    # environment variable for the login to work locally.
    demo.launch()