from typing import Any

from inspect_ai.dataset import Sample

from inspect_evals.sec_qa import record_to_sample
from inspect_evals.sec_qa.sec_qa import format_system_message_fewshot, sample_to_fewshot

EXAMPLE_RECORD: dict[str, Any] = {
    "Question": "Which of the following is a common indicator of an SQL injection attack?",
    "A": "Frequent changes in user account permissions.",
    "B": "Decreased performance of endpoint protection systems.",
    "C": "Unusually high data upload rates to a web server.",
    "D": "Sudden uptick in SQL queries, far beyond the usual baseline for the application.",
    "Answer": "D",
}

EXPECTED_FORMATTED_SAMPLE_TO_FEWSHOT = "QUESTION:\nWhich of the following is a common indicator of an SQL injection attack?\n\nCHOICES:\nA) Frequent changes in user account permissions.\nB) Decreased performance of endpoint protection systems.\nC) Unusually high data upload rates to a web server.\nD) Sudden uptick in SQL queries, far beyond the usual baseline for the application.\n\n\nANSWER: D"


def test_record_to_sample():
    """Test that the record is mapped correctly to Inspect Sample"""
    example_sample = record_to_sample(EXAMPLE_RECORD)
    assert example_sample == Sample(
        input="Which of the following is a common indicator of an SQL injection attack?",
        choices=[
            "Frequent changes in user account permissions.",
            "Decreased performance of endpoint protection systems.",
            "Unusually high data upload rates to a web server.",
            "Sudden uptick in SQL queries, far beyond the usual baseline for the application.",
        ],
        target="D",
    )


def test_format_system_message_fewshot():
    """Test that the system message is formatted correctly for some samples"""
    formatted_system_message = format_system_message_fewshot(
        [record_to_sample(EXAMPLE_RECORD), record_to_sample(EXAMPLE_RECORD)]
    )
    assert (
        formatted_system_message
        == "The following are multiple choice questions about Computer Security. Some examples are provided below."
        + "\n\n"
        + EXPECTED_FORMATTED_SAMPLE_TO_FEWSHOT
        + "\n\n"
        + EXPECTED_FORMATTED_SAMPLE_TO_FEWSHOT
    )


def test_sample_to_fewshot():
    """Test that the sample is formatted correctly"""
    formatted_sample_to_fewshot = sample_to_fewshot(record_to_sample(EXAMPLE_RECORD))
    assert formatted_sample_to_fewshot == EXPECTED_FORMATTED_SAMPLE_TO_FEWSHOT
