import os
import random

from huggingface_hub import HfApi, whoami

import gradio as gr
from datasets import load_dataset

from data_to_parquet import to_parquet

EXAM_DATASET_ID = os.getenv("EXAM_DATASET_ID") or "agents-course/unit_1_quiz"
EXAM_MAX_QUESTIONS = os.getenv("EXAM_MAX_QUESTIONS") or 10
EXAM_PASSING_SCORE = os.getenv("EXAM_PASSING_SCORE") or 0.7

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

upload_api = HfApi(token=os.getenv("HF_TOKEN"))
# Convert dataset to a list of dicts and randomly sort
quiz_data = ds.to_pandas().to_dict("records")
random.shuffle(quiz_data)

# Limit to max questions if specified
if EXAM_MAX_QUESTIONS:
    quiz_data = quiz_data[: int(EXAM_MAX_QUESTIONS)]


def on_user_logged_in(token: gr.OAuthToken | None):
    """
    If the user has a valid token, show Start button.
    Otherwise, keep the login button visible.
    """
    if token is not None:
        return [
            gr.update(visible=False),  # login button visibility
            gr.update(visible=True),  # start button visibility
            gr.update(visible=False),  # next button visibility
            gr.update(visible=False),  # submit button visibility
            "",  # question text
            [],  # radio choices (empty list = no choices)
            "Click 'Start' to begin the quiz",  # status message
            0,  # question_idx
            [],  # user_answers
            "",  # final_markdown content
            token,  # user token
        ]
    else:
        return [
            gr.update(visible=True),  # login button visibility
            gr.update(visible=False),  # start button visibility
            gr.update(visible=False),  # next button visibility
            gr.update(visible=False),  # submit button visibility
            "",  # question text
            [],  # radio choices
            "",  # status message
            0,  # question_idx
            [],  # user_answers
            "",  # final_markdown content
            None,  # no token
        ]


def push_results_to_hub(user_answers, token: gr.OAuthToken | None):
    """
    Create a new dataset from user_answers and push it to the Hub.
    Calculates grade and checks against passing threshold.
    """
    if token is None:
        gr.Warning("Please log in to Hugging Face before pushing!")
        return

    # Calculate grade
    correct_count = sum(1 for answer in user_answers if answer["is_correct"])
    total_questions = len(user_answers)
    grade = correct_count / total_questions if total_questions > 0 else 0

    if grade < float(EXAM_PASSING_SCORE):
        gr.Warning(
            f"Score {grade:.1%} below passing threshold of {float(EXAM_PASSING_SCORE):.1%}"
        )
        return  # do not continue

    gr.Info("Submitting answers to the Hub. Please wait...", duration=2)

    user_info = whoami(token=token.token)
    # TODO:
    # check if username already has "username.parquet" in the dataset and download that (or read values directly from dataset viewer if possible)
    # instead of replacing the values check if the new score is better than the old one
    to_parquet(
        upload_api,  # api
        "agents-course/students-data",  # repo_id
        user_info["name"],  # username
        grade,  # unit1 score
        0.0,  # unit2 score
        0.0,  # unit3 score
        0.0,  # unit4 score
        0,  # already certified or not
    )

    gr.Success(
        f"Your responses have been submitted to the Hub! Final grade: {grade:.1%}"
    )


def handle_quiz(question_idx, user_answers, selected_answer, is_start):
    """
    Handle quiz state transitions and store answers
    """
    if not is_start and question_idx < len(quiz_data):
        current_q = quiz_data[question_idx]
        correct_reference = current_q["correct_answer"]
        correct_reference = f"answer_{correct_reference}".lower()
        is_correct = selected_answer == current_q[correct_reference]
        user_answers.append(
            {
                "question": current_q["question"],
                "selected_answer": selected_answer,
                "correct_answer": current_q[correct_reference],
                "is_correct": is_correct,
                "correct_reference": correct_reference,
            }
        )
        question_idx += 1

    if question_idx >= len(quiz_data):
        correct_count = sum(1 for answer in user_answers if answer["is_correct"])
        grade = correct_count / len(user_answers)
        results_text = (
            f"**Quiz Complete!**\n\n"
            f"Your score: {grade:.1%}\n"
            f"Passing score: {float(EXAM_PASSING_SCORE):.1%}\n\n"
        )
        return [
            "",  # question_text
            gr.update(choices=[], visible=False),  # hide radio choices
            f"{'✅ Passed!' if grade >= float(EXAM_PASSING_SCORE) else '❌ Did not pass'}",
            question_idx,
            user_answers,
            gr.update(visible=False),  # start button visibility
            gr.update(visible=False),  # next button visibility
            gr.update(visible=True),  # submit button visibility
            results_text,  # final results text
        ]

    # Show next question
    q = quiz_data[question_idx]
    return [
        f"## Question {question_idx + 1} \n### {q['question']}",  # question text
        gr.update(  # properly update radio choices
            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,
        gr.update(visible=False),  # start button visibility
        gr.update(visible=True),  # next button visibility
        gr.update(visible=False),  # submit button visibility
        "",  # clear final markdown
    ]


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}"

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

    with gr.Row(variant="compact"):
        gr.Markdown(f"## Welcome to the {EXAM_DATASET_ID} Quiz")

    with gr.Row(variant="compact"):
        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."
        )

    with gr.Row(variant="panel"):
        question_text = gr.Markdown("")
        radio_choices = gr.Radio(
            choices=[], label="Your Answer", scale=1.5, visible=False
        )

    with gr.Row(variant="compact"):
        status_text = gr.Markdown("")
        final_markdown = gr.Markdown("")

    with gr.Row(variant="compact"):
        login_btn = gr.LoginButton(visible=True)
        start_btn = gr.Button("Start ⏭️", visible=True)
        next_btn = gr.Button("Next ⏭️", visible=False)
        submit_btn = gr.Button("Submit ✅", visible=False)

    # Wire up the event handlers
    login_btn.click(
        fn=on_user_logged_in,
        inputs=None,
        outputs=[
            login_btn,
            start_btn,
            next_btn,
            submit_btn,
            question_text,
            radio_choices,
            status_text,
            question_idx,
            user_answers,
            final_markdown,
            user_token,
        ],
    )

    start_btn.click(
        fn=handle_quiz,
        inputs=[question_idx, user_answers, gr.State(""), gr.State(True)],
        outputs=[
            question_text,
            radio_choices,
            status_text,
            question_idx,
            user_answers,
            start_btn,
            next_btn,
            submit_btn,
            final_markdown,
        ],
    )

    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,
            next_btn,
            submit_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()