from typing import cast

import pytest
import pytest_asyncio  # noqa: F401
from inspect_ai.dataset import MemoryDataset, Sample
from inspect_ai.model import (
    ModelName,
)
from inspect_ai.scorer import Target
from inspect_ai.scorer._metric import CORRECT, INCORRECT
from inspect_ai.solver import TaskState

from inspect_evals.cybermetric import cybermetric_80
from inspect_evals.cybermetric.cybermetric import record_to_sample

# First question of CyberMetric-80
EXAMPLE_QUESTION = {
    "question": "What is the primary requirement for an Random Bit Generator's (RBG) output to be used for generating cryptographic keys?",
    "answers": {
        "A": "The RBG's output should have a length that matches the target data to be protected",
        "B": "The RBG's output should be computationally indistinguishable from random bits and provide sufficient entropy to support the security strength required for the target data",
        "C": "The RBG's output should be as long as possible to ensure maximal randomness",
        "D": "The RBG's output should have precisely the same length as the symmetric key to be generated",
    },
    "solution": "B",
}


def test_record_to_sample():
    """Tests that record_to_sample correctly transforms a question from the JSON dataset"""
    # Convert question from dict to Sample
    sample = record_to_sample(EXAMPLE_QUESTION)

    # Assert Sample is populated accurately with correct types
    assert isinstance(sample.target, str)
    assert sample.target == EXAMPLE_QUESTION["solution"]
    assert sample.target == "B"

    assert all(isinstance(choice, str) for choice in sample.choices)
    assert sample.choices == list(EXAMPLE_QUESTION["answers"].keys())
    assert sample.choices == ["A", "B", "C", "D"]

    assert isinstance(sample.input, str)
    assert EXAMPLE_QUESTION["question"] in sample.input

    # Assert each option is provided in input
    for option in EXAMPLE_QUESTION["answers"].values():
        assert option in sample.input

    expected_input = """Question: What is the primary requirement for an Random Bit Generator's (RBG) output to be used for generating cryptographic keys?
Options: A) The RBG's output should have a length that matches the target data to be protected, B) The RBG's output should be computationally indistinguishable from random bits and provide sufficient entropy to support the security strength required for the target data, C) The RBG's output should be as long as possible to ensure maximal randomness, D) The RBG's output should have precisely the same length as the symmetric key to be generated

"""
    assert sample.input == expected_input


def generate_task_state(sample: Sample, is_correct: bool) -> TaskState:
    """Generates a TaskState for the given Sample, setting the model output based on correctness."""
    options = cast(list[str], sample.choices)
    answer = cast(str, sample.target)

    # Define task
    task_state = TaskState(
        model=ModelName("openai/gpt-4o"),
        sample_id="mock_1",
        epoch=1,
        input=sample.input,
        messages=[],
        completed=True,
        choices=options,
    )

    if is_correct:
        target_index = options.index(answer)
        task_state.choices.mark_choice(index=target_index, correct=True)

    return task_state


@pytest.mark.asyncio
async def test_cybermetric_correct_answer():
    """Tests that a task (cybermetric_80) runs end-to-end with an accurate result when the LLM answers correctly."""
    # Define Sample from first question of CyberMetric-80
    sample = record_to_sample(EXAMPLE_QUESTION)

    # Define task with single Sample from memory
    task = cybermetric_80()
    task.dataset = MemoryDataset([sample])

    state = generate_task_state(sample, is_correct=True)

    # Score answer
    target = Target(sample.target)
    score = await task.scorer[0](state, target)

    # Assert score computed accurately for correct answer
    assert score.value == CORRECT, (
        f"Expected '{CORRECT}' for a correct answer, got {score.value}"
    )
    return


@pytest.mark.asyncio
async def test_cybermetric_incorrect_answer():
    """Tests that a task (cybermetric_80) runs end-to-end with an accurate result when the LLM answers incorrectly."""
    # Define Sample from first question of CyberMetric-80
    sample = record_to_sample(EXAMPLE_QUESTION)

    # Define task with single Sample from memory
    task = cybermetric_80()
    task.dataset = MemoryDataset([sample])

    # Define TaskState representing incorrect model answer
    state = generate_task_state(sample, is_correct=False)

    # Score answer
    target = Target(sample.target)
    score = await task.scorer[0](state, target)

    # Assert score computed accurately for incorrect answer
    assert score.value == INCORRECT, (
        f"Expected '{INCORRECT}' for an incorrect answer, got {score.value}"
    )
    return
