import pytest
from inspect_ai.model import (
    ChatCompletionChoice,
    ChatMessageAssistant,
    ModelOutput,
    get_model,
)
from inspect_ai.scorer import SampleScore, Score, Target
from inspect_ai.solver import TaskState

from inspect_evals.simpleqa.simpleqa import (
    simpleqa_metric,
    simpleqa_scorer,
)


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "completion_value,expected_grade,expected_value",
    [
        ("Paris", "A", {"correct": 1.0, "incorrect": 0.0, "not_attempted": 0.0}),
        ("London", "B", {"correct": 0.0, "incorrect": 1.0, "not_attempted": 0.0}),
        ("I don't know", "C", {"correct": 0.0, "incorrect": 0.0, "not_attempted": 1.0}),
    ],
)
async def test_simpleqa_scorer(completion_value, expected_grade, expected_value):
    QUESTION = "What is the capital of France?"
    TARGET = "Paris"

    grader_model = get_model(
        "mockllm/model",
        custom_outputs=[
            ModelOutput(
                model="mockllm/model",
                choices=[
                    ChatCompletionChoice(
                        message=ChatMessageAssistant(content=expected_grade)
                    )
                ],
            )
        ],
    )
    scorer_fn = simpleqa_scorer(grader_model)
    state = TaskState(
        model="mockllm/model",
        input=QUESTION,
        output=ModelOutput(
            choices=[
                ChatCompletionChoice(
                    message=ChatMessageAssistant(content=completion_value)
                )
            ]
        ),
        sample_id=1,
        epoch=0,
        messages=[],
    )
    score = await scorer_fn(state, Target(TARGET))

    assert isinstance(score, Score)
    assert score.value == expected_value
    assert score.answer == completion_value


@pytest.mark.parametrize(
    "scores,expected_metrics",
    [
        (
            # Test case 1: All correct answers
            [
                SampleScore(
                    score=Score(
                        value={"correct": 1.0, "incorrect": 0.0, "not_attempted": 0.0},
                        answer="Paris",
                    )
                ),
                SampleScore(
                    score=Score(
                        value={"correct": 1.0, "incorrect": 0.0, "not_attempted": 0.0},
                        answer="London",
                    )
                ),
            ],
            {
                "correct": 1.0,
                "incorrect": 0.0,
                "not_attempted": 0.0,
                "correct_given_attempted": 1.0,
                "f_score": 1.0,
            },
        ),
        (
            # Test case 2: Mixed results
            [
                SampleScore(
                    score=Score(
                        value={"correct": 1.0, "incorrect": 0.0, "not_attempted": 0.0},
                        answer="Paris",
                    )
                ),
                SampleScore(
                    score=Score(
                        value={"correct": 0.0, "incorrect": 1.0, "not_attempted": 0.0},
                        answer="London",
                    )
                ),
                SampleScore(
                    score=Score(
                        value={"correct": 0.0, "incorrect": 0.0, "not_attempted": 1.0},
                        answer="I don't know",
                    )
                ),
            ],
            {
                "correct": 1 / 3,
                "incorrect": 1 / 3,
                "not_attempted": 1 / 3,
                "correct_given_attempted": 0.5,
                "f_score": 0.4,
            },
        ),
        (
            # Test case 3: All not attempted
            [
                SampleScore(
                    score=Score(
                        value={"correct": 0.0, "incorrect": 0.0, "not_attempted": 1.0},
                        answer="I don't know",
                    )
                ),
                SampleScore(
                    score=Score(
                        value={"correct": 0.0, "incorrect": 0.0, "not_attempted": 1.0},
                        answer="Unsure",
                    )
                ),
            ],
            {
                "correct": 0.0,
                "incorrect": 0.0,
                "not_attempted": 1.0,
                "correct_given_attempted": 0.0,
                "f_score": 0.0,
            },
        ),
    ],
)
def test_simpleqa_metric(scores, expected_metrics):
    metric = simpleqa_metric()
    result = metric(scores)

    assert isinstance(result, dict)
    for key, value in expected_metrics.items():
        assert key in result
        assert abs(result[key] - value) < 1e-6, (
            f"Metric {key} expected {value} but got {result[key]}"
        )
