# tests/personality/test_any_choice.py
from typing import Union

import pytest
from inspect_ai.model import (
    ChatMessageAssistant,
    ChatMessageSystem,
    ChatMessageTool,
    ChatMessageUser,
    ModelName,
    ModelOutput,
)
from inspect_ai.scorer import CORRECT, INCORRECT, Target
from inspect_ai.solver import TaskState

from inspect_evals.personality.personality import any_choice

CHOICE_TEXT = "A,B,C,D,E"
Message = Union[
    ChatMessageSystem,
    ChatMessageUser,
    ChatMessageAssistant,
    ChatMessageTool,
]


def _state(answer_line: str) -> TaskState:
    """
    Build a minimal TaskState for any_choice.
    `answer_line` must look like "ANSWER: A".
    """  # noqa: D205
    model_name = ModelName("mockllm/none")
    output = ModelOutput.from_content(str(model_name), answer_line)

    # tiny conversation history
    user_msg: ChatMessageUser = ChatMessageUser(content="Choose A–E:")
    assistant_msg: ChatMessageAssistant = ChatMessageAssistant(
        content=answer_line, model=str(model_name)
    )
    msgs: list[Message] = [user_msg, assistant_msg]

    ts = TaskState(
        model=model_name,
        sample_id=1,
        epoch=0,
        input=str(user_msg.content),
        messages=msgs,
        target=Target(CHOICE_TEXT),
        output=output,
    )

    return ts


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("completion", "expected"),
    [
        ("ANSWER: A", CORRECT),  # allowed letter → CORRECT
        ("ANSWER: C", CORRECT),  # any allowed letter passes
        ("ANSWER: Z", INCORRECT),  # not in allowed letters
        ("I refuse to answer.", INCORRECT),  # no parsable "ANSWER:" line
    ],
)
async def test_any_choice_basic(completion: str, expected: str) -> None:
    scorer = any_choice()
    state = _state(completion)
    result = await scorer(state, state.target)

    assert result.value == expected
