"""Gradio app to validate examples of the FoQA dataset."""

from functools import partial
import os
from typing import Generator
import gradio as gr
from datasets import Dataset, load_dataset
import logging
import pandas as pd
import os
from dotenv import load_dotenv

load_dotenv()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("foqa")


dataset = load_dataset(
    "alexandrainst/foqa", split="train", token=os.getenv("HF_HUB_TOKEN")
)
assert isinstance(dataset, Dataset)
df = pd.DataFrame(dataset.to_pandas())


def non_validated_samples() -> Generator[tuple[str, str, str], None, None]:
    """Iterate over non-validated samples in the FoQA dataset.

    Yields:
        A tuple (idx, question, answer) of a non-validated sample.
    """
    for idx, sample in df.iterrows():
        if sample.validation is None:
            yield str(idx), sample.question, sample.answers["text"][0]


itr = non_validated_samples()


def main():
    idx, question, answer = next(itr)

    with gr.Blocks(theme="monochrome", title="FoQA validation") as demo:
        gr.Markdown("""
            # FoQA Validation

            This app automatically fetches examples from the Faroese Question Answering
            dataset (FoQA), allowing you to annotate whether the question and answer
            are correct Faroese or not.
        """)
        with gr.Row():
            with gr.Column():
                gr.Markdown("### Sample ID")
                idx_box = gr.Markdown(value=idx)
                gr.Markdown("### Question")
                question_box = gr.Markdown(value=question)
                gr.Markdown("### Answer")
                answer_box = gr.Markdown(value=answer)
            with gr.Column():
                correct_btn = gr.Button(value="Correct")
                incorrect_btn = gr.Button(value="Incorrect")
                incorrect_answer_btn = gr.Button(value="Incorrect Answer")
                save_results_btn = gr.Button(value="Save results")

        correct_btn.click(
            fn=partial(assign_correct, itr=itr),
            inputs=[idx_box, question_box, answer_box],
            outputs=[idx_box, question_box, answer_box],
        )
        incorrect_btn.click(
            fn=partial(assign_incorrect, itr=itr),
            inputs=[idx_box, question_box, answer_box],
            outputs=[idx_box, question_box, answer_box],
        )
        incorrect_answer_btn.click(
            fn=partial(assign_incorrect_answer, itr=itr),
            inputs=[idx_box, question_box, answer_box],
            outputs=[idx_box, question_box, answer_box],
        )
        save_results_btn.click(fn=save_results)

    auth = [
        ("admin", os.environ["ADMIN_PASSWORD"]),
        ("annika", os.environ["ANNIKA_PASSWORD"]),
    ]
    demo.launch(auth=auth)


def save_results() -> None:
    """Update the FoQA dataset with the validation status of a sample."""
    logger.info("Saving results...")
    gr.Info(message="Saving results...")
    Dataset.from_pandas(df, preserve_index=False).push_to_hub(
        repo_id="alexandrainst/foqa", token=os.getenv("HF_HUB_TOKEN")
    )
    gr.Info(message="Saved results!")
    logger.info("Saved results.")


def assign_correct(
    idx: str, question: str, answer: str, itr: Generator
) -> tuple[gr.Markdown, gr.Markdown, gr.Markdown]:
    """Assign the question and answer as correct.

    Args:
        idx:
            The index of the sample to be assigned as correct.
        question:
            The question to be assigned as correct.
        answer:
            The answer to be assigned as correct.
        itr:
            The iterator over non-validated samples.

    Returns:
        The updated textboxes.
    """
    gr.Info(message="Assigned sample as correct")
    logger.info(f"Assigned sample as correct: {question} - {answer}")
    df.iloc[int(idx)].validation = "correct"
    idx, question, answer = next(itr)
    return (
        gr.Markdown(value=idx), gr.Markdown(value=question), gr.Markdown(value=answer)
    )


def assign_incorrect(
    idx: str, question: str, answer: str, itr: Generator
) -> tuple[gr.Markdown, gr.Markdown, gr.Markdown]:
    """Assign the question and answer as incorrect.

    Args:
        idx:
            The index of the sample to be assigned as incorrect.
        question:
            The question to be assigned as incorrect.
        answer:
            The answer to be assigned as incorrect.
        itr:
            The iterator over non-validated samples.

    Returns:
        The updated textboxes.
    """
    gr.Info(message="Assigned sample as incorrect")
    logger.info(f"Assigned sample as incorrect: {question} - {answer}")
    df.iloc[int(idx)].validation = "incorrect"
    idx, question, answer = next(itr)
    return (
        gr.Markdown(value=idx), gr.Markdown(value=question), gr.Markdown(value=answer)
    )


def assign_incorrect_answer(
    idx: str, question: str, answer: str, itr: Generator
) -> tuple[gr.Markdown, gr.Markdown, gr.Markdown]:
    """Assign the answer as incorrect.

    Args:
        idx:
            The index of the sample to be assigned as incorrect.
        question:
            The question to be assigned as incorrect.
        answer:
            The answer to be assigned as incorrect.
        itr:
            The iterator over non-validated samples.

    Returns:
        The updated textboxes.
    """
    gr.Info(message="Assigned sample answer as incorrect")
    logger.info(f"Assigned sample answer as incorrect: {answer}")
    df.iloc[int(idx)].validation = "incorrect-answer"
    idx, question, answer = next(itr)
    return (
        gr.Markdown(value=idx), gr.Markdown(value=question), gr.Markdown(value=answer)
    )


if __name__ == "__main__":
    main()