"""
SecQA: A Concise Question-Answering Dataset for Evaluating Large Language Models in Computer Security

Zefang Liu
https://arxiv.org/pdf/2312.15838
"""

from typing import Any

from inspect_ai import Task, task
from inspect_ai.dataset import Dataset, Sample, hf_dataset
from inspect_ai.scorer import choice
from inspect_ai.solver import multiple_choice, system_message

# setup for problem + instructions for providing answer
SYSTEM_MESSAGE_FEWSHOT = """
The following are multiple choice questions about Computer Security. Some examples are provided below.

{examples}
""".strip()

DATASET_PATH = "zefang-liu/secqa"
DATASET_SUBSET_NAME_V1 = "secqa_v1"
DATASET_SUBSET_NAME_V2 = "secqa_v2"
DATASET_TEST_SPLIT = "test"
DATASET_FEWSHOT_SPLIT = "dev"


@task
def sec_qa_v1() -> Task:
    """Inspect Task implementing the SecQA benchmark v1 0-shot"""
    dataset = retrieve_hf_dataset(DATASET_SUBSET_NAME_V1, DATASET_TEST_SPLIT)
    return Task(
        dataset=dataset,
        solver=[multiple_choice()],
        scorer=choice(),
    )


@task
def sec_qa_v1_5_shot() -> Task:
    """Inspect Task implementing the SecQA benchmark v1 5-shot"""
    fewshot_samples = retrieve_hf_dataset(DATASET_SUBSET_NAME_V1, DATASET_FEWSHOT_SPLIT)
    dataset = retrieve_hf_dataset(DATASET_SUBSET_NAME_V1, DATASET_TEST_SPLIT)
    return Task(
        dataset=dataset,
        solver=[
            system_message(format_system_message_fewshot(fewshot_samples)),
            multiple_choice(),
        ],
        scorer=choice(),
    )


@task
def sec_qa_v2() -> Task:
    """Inspect Task implementing the SecQA benchmark v2 0-shot"""
    dataset = retrieve_hf_dataset(DATASET_SUBSET_NAME_V2, DATASET_TEST_SPLIT)
    return Task(
        dataset=dataset,
        solver=[multiple_choice()],
        scorer=choice(),
    )


@task
def sec_qa_v2_5_shot() -> Task:
    """Inspect Task implementing the SecQA benchmark v2 5-shot"""
    fewshot_samples = retrieve_hf_dataset(DATASET_SUBSET_NAME_V2, DATASET_FEWSHOT_SPLIT)
    dataset = retrieve_hf_dataset(DATASET_SUBSET_NAME_V2, DATASET_TEST_SPLIT)
    return Task(
        dataset=dataset,
        solver=[
            system_message(format_system_message_fewshot(fewshot_samples)),
            multiple_choice(),
        ],
        scorer=choice(),
    )


def retrieve_hf_dataset(name: str, split: str) -> Dataset:
    return hf_dataset(
        path=DATASET_PATH,
        name=name,
        split=split,
        sample_fields=record_to_sample,
        trust=True,
    )


def record_to_sample(record: dict[str, Any]) -> Sample:
    return Sample(
        input=record["Question"],
        choices=[record["A"], record["B"], record["C"], record["D"]],
        target=record["Answer"],
    )


def format_system_message_fewshot(fewshot_samples: Dataset) -> str:
    return SYSTEM_MESSAGE_FEWSHOT.format(
        examples="\n\n".join(
            [sample_to_fewshot(sample=sample) for sample in fewshot_samples]
        )
    )


def sample_to_fewshot(sample: Sample) -> str:
    prob_str = f"""QUESTION:\n{sample.input}"""
    labeled_choices = "\n"
    for i, letter_label in enumerate(["A", "B", "C", "D"]):
        labeled_choices = (
            labeled_choices
            + f"{letter_label}) {sample.choices[i] if sample.choices is not None else None}\n"
        )
    choices_str = f"""CHOICES:{labeled_choices}"""
    ans_str = f"""ANSWER: {sample.target}"""

    return f"""{prob_str}\n\n{choices_str}\n\n{ans_str}"""
