import pytest
from inspect_ai.model import ModelName, ModelOutput
from inspect_ai.solver import TaskState

from inspect_evals.mind2web.dataset import Answer, Mind2WebMetadata
from inspect_evals.mind2web.scorer import mind2web_scorer


@pytest.mark.asyncio
async def test_mind2web_scorer():
    # Create a scorer instance
    scorer = mind2web_scorer()

    # Test cases with both element accuracy and action F1
    test_cases = [
        {
            "completion": "Answer: B.\nAction: SELECT\nValue: Pickup",
            "target": "B.\nAction: SELECT\nValue: Pickup",
            "metadata": Mind2WebMetadata(
                final_html="<div>Test HTML</div>",
                previous_actions="None",
                answer=Answer(letter="B", action="SELECT", value="Pickup"),
                target_idx=1,
                confirmed_task="test task",
                task_id="test_1",
                annotation_id="test_1",
                action_uid="test_1",
            ),
            "expected_score": 1.0,
            "expected_f1": 1.0,
        },
        {
            "completion": "Answer: B.\nAction: CLICK\nValue: Submit",
            "target": "B.\nAction: CLICK\nValue: Submit",
            "metadata": Mind2WebMetadata(
                final_html="<div>Test HTML</div>",
                previous_actions="None",
                answer=Answer(letter="B", action="CLICK", value="Submit"),
                target_idx=1,
                confirmed_task="test task",
                task_id="test_2",
                annotation_id="test_2",
                action_uid="test_2",
            ),
            "expected_score": 1.0,
            "expected_f1": 1.0,
        },
        {
            "completion": "Answer: C.\nAction: CLICK\nValue: Next",
            "target": "C.\nAction: CLICK\nValue: Previous",
            "metadata": Mind2WebMetadata(
                final_html="<div>Test HTML</div>",
                previous_actions="None",
                answer=Answer(letter="C", action="CLICK", value="Previous"),
                target_idx=2,
                confirmed_task="test task",
                task_id="test_3",
                annotation_id="test_3",
                action_uid="test_3",
            ),
            "expected_score": 1.0,
            "expected_f1": 1.0,  # For CLICK actions, only action type is compared
        },
        # Test element mismatch but action match
        {
            "completion": "Answer: B.\nAction: SELECT\nValue: Option",
            "target": "C.\nAction: SELECT\nValue: Option",
            "metadata": Mind2WebMetadata(
                final_html="<div>Test HTML</div>",
                previous_actions="None",
                answer=Answer(letter="C", action="SELECT", value="Option"),
                target_idx=1,
                confirmed_task="test task",
                task_id="test_4",
                annotation_id="test_4",
                action_uid="test_4",
            ),
            "expected_score": 0.0,
            "expected_f1": 1.0,
        },
        # Test invalid format
        {
            "completion": "Invalid format",
            "target": "B.\nAction: CLICK\nValue: Submit",
            "metadata": Mind2WebMetadata(
                final_html="<div>Test HTML</div>",
                previous_actions="None",
                answer=Answer(letter="B", action="CLICK", value="Submit"),
                target_idx=1,
                confirmed_task="test task",
                task_id="test_5",
                annotation_id="test_5",
                action_uid="test_5",
            ),
            "expected_score": 0.0,
            "expected_f1": 0.0,
        },
        # Test case insensitivity for element
        {
            "completion": "Answer: a.\nAction: CLICK\nValue: Submit",
            "target": "A.\nAction: CLICK\nValue: Submit",
            "metadata": Mind2WebMetadata(
                final_html="<div>Test HTML</div>",
                previous_actions="None",
                answer=Answer(letter="A", action="CLICK", value="Submit"),
                target_idx=0,
                confirmed_task="test task",
                task_id="test_6",
                annotation_id="test_6",
                action_uid="test_6",
            ),
            "expected_score": 1.0,
            "expected_f1": 1.0,
        },
        # Test partial action match
        {
            "completion": "Answer: C.\nAction: SELECT\nValue: Username field",
            "target": "C.\nAction: SELECT\nValue: Password field",
            "metadata": Mind2WebMetadata(
                final_html="<div>Test HTML</div>",
                previous_actions="None",
                answer=Answer(letter="C", action="SELECT", value="Password field"),
                target_idx=2,
                confirmed_task="test task",
                task_id="test_7",
                annotation_id="test_7",
                action_uid="test_7",
            ),
            "expected_score": 1.0,
            "expected_f1": 0.6666666666666666,  # Partial match on SELECT and field
        },
    ]

    for tc in test_cases:
        # Create a TaskState with the test completion and metadata
        state = TaskState(
            model=ModelName("anthropic/test-model"),
            sample_id=tc["metadata"].task_id,
            epoch=0,
            input="test input",
            messages=[],
            output=ModelOutput(),
            metadata=tc["metadata"].model_dump(),
        )
        state.output.completion = tc["completion"]

        # Create a Target with the test target and metadata
        target = None

        # Score the completion
        result = await scorer(state, target)

        # Verify both element accuracy and action F1 scores
        assert result.metadata["element_acc"] == tc["expected_score"], (
            f"Failed element accuracy for completion '{tc['completion']}'. "
            f"Expected {tc['expected_score']}, got {result.metadata['element_acc']}"
        )

        assert result.metadata["action_f1"] == tc["expected_f1"], (
            f"Failed action F1 for completion '{tc['completion']}'. "
            f"Expected {tc['expected_f1']}, got {result.metadata['action_f1']}"
        )
