import json

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

from inspect_evals.mind2web.dataset import Mind2WebMetadata
from inspect_evals.mind2web.scorer import (
    extract_answer,
    mind2web_scorer,
    step_success_rate,
    task_success_rate,
)
from inspect_evals.mind2web.task import mind2web


def load_test_data():
    """Load sample test data from JSON files."""
    with open("tests/mind2web/data/sample_10.json") as f:
        sample_10 = json.load(f)
    with open("tests/mind2web/data/sample_100.json") as f:
        sample_100 = json.load(f)
    return sample_10, sample_100


@pytest.mark.dataset_download
@pytest.mark.asyncio
async def test_mind2web_task_creation():
    """Test mind2web task creation and configuration."""
    sample_10, _ = load_test_data()
    task = mind2web()

    # Verify task configuration
    assert task.name == "mind2web_eval"
    assert task.dataset is not None
    assert len(task.dataset.samples) > 0
    # Check solver and scorer functions
    assert len(task.solver) == 1  # solve()
    assert len(task.scorer) == 1  # mind2web_scorer()

    # Verify sample structure with real data from sample_10.json
    sample = sample_10[0]  # First sample from test data
    assert sample["id"] is not None
    assert sample["input"] is not None
    assert isinstance(sample["choices"], list)
    assert len(sample["choices"]) > 1
    assert sample["choices"][0] == "A. None of the above"
    assert all(
        c.startswith(("A.", "B.", "C.", "D.", "E.", "F.")) for c in sample["choices"]
    )


@pytest.mark.asyncio
async def test_mind2web_scoring_click():
    """Test scoring with CLICK action using real data."""
    sample_10, _ = load_test_data()

    # Use first sample from sample_10.json which has a CLICK action
    sample: Sample = Sample.model_validate(sample_10[0])  # "Quick Tools" click action

    metadata = sample.metadata_as(Mind2WebMetadata)
    scorer_fn = mind2web_scorer()

    # Get target answer from sample metadata
    target_answer = metadata.answer
    letter = target_answer.letter  # "B"
    action = target_answer.action  # "CLICK"

    # Test cases with real data - no indentation to match scorer's expectations
    test_cases = [
        # Exact match case - for CLICK, only action type matters, not value
        (
            f"Answer: {letter}\nAction: {action}\nValue: Different value",
            1.0,
            1.0,  # Should be 1.0 since CLICK only compares action type
        ),
        # Wrong action type
        (
            f"Answer: {letter}\nAction: TYPE",
            1.0,
            0.0,
        ),
        # Wrong element but correct action
        (
            f"Answer: A\nAction: {action}",
            0.0,
            0.0,  # Action is correct even though element is wrong
        ),
    ]

    for completion, expected_element_acc, expected_action_f1 in test_cases:
        state = TaskState(
            messages=[],
            output=ModelOutput(
                choices=[
                    ChatCompletionChoice(
                        message=ChatMessageAssistant(content=completion)
                    )
                ]
            ),
            choices=sample.choices,
            model="test-model",
            sample_id=sample.id,
            epoch=0,
            input=sample.input,
            metadata=sample.metadata,
        )

        score = await scorer_fn(
            state,
            target=Target(target=sample.target),
        )

        # Verify scoring matches expectations
        assert score.metadata["element_acc"] == expected_element_acc
        assert score.metadata["is_answered"] == 1
        assert score.metadata["action_f1"] == expected_action_f1, (
            f"Expected action_f1 {expected_action_f1}, got {score.metadata['action_f1']}"
        )


@pytest.mark.asyncio
async def test_mind2web_scoring_type():
    """Test scoring with TYPE action."""
    sample_10, _ = load_test_data()

    # Use sample with TYPE action from sample_10.json (zip code entry)
    sample: Sample = Sample.model_validate(
        sample_10[2]
    )  # TYPE action for zip code "60505"
    scorer_fn = mind2web_scorer()

    # Get target answer from sample metadata
    metadata = sample.metadata_as(Mind2WebMetadata)
    target_answer = metadata.answer
    letter = target_answer.letter  # "F"
    value = target_answer.value  # "60505"

    # Test exact match
    completion = f"Answer: {letter}\nAction: TYPE\nValue: {value}"

    state = TaskState(
        messages=[],
        output=ModelOutput(
            choices=[
                ChatCompletionChoice(message=ChatMessageAssistant(content=completion))
            ]
        ),
        choices=sample.choices,
        model="test-model",
        sample_id=sample.id,
        epoch=0,
        input=sample.input,
        metadata=sample.metadata,
    )
    score = await scorer_fn(
        state,
        target=Target(target=sample.target),
    )

    # For TYPE actions, verify both element_acc and action_f1
    assert score.value == "C"  # Perfect match
    assert score.metadata["element_acc"] == 1.0
    assert score.metadata["is_answered"] == 1
    assert score.metadata["action_f1"] == 1.0

    # Test partial match (correct element but different text)
    completion = f"Answer: {letter}\nAction: TYPE\nValue: Different text input"

    state = TaskState(
        messages=[],
        output=ModelOutput(
            choices=[
                ChatCompletionChoice(message=ChatMessageAssistant(content=completion))
            ]
        ),
        choices=sample.choices,
        model="test-model",
        sample_id=sample.id,
        epoch=0,
        input=sample.input,
        metadata=sample.metadata,
    )
    score = await scorer_fn(
        state,
        target=Target(target=sample.target),
    )

    # Element is correct but action text differs
    assert score.value == "C"  # Element match
    assert score.metadata["element_acc"] == 1.0
    assert score.metadata["is_answered"] == 1
    assert score.metadata["action_f1"] < 1.0  # Partial text match


@pytest.mark.asyncio
async def test_mind2web_scoring_edge_cases():
    """Test scoring with edge cases."""
    sample_10, _ = load_test_data()
    sample: Sample = Sample.model_validate(sample_10[0])  # Use first sample
    scorer_fn = mind2web_scorer()

    metadata = sample.metadata_as(Mind2WebMetadata)
    _target_answer = metadata.answer

    # Test edge cases
    test_cases = [
        ("", "I", 0, 0.0),  # Empty completion
        ("Invalid response without answer", "I", 1, 0.0),  # Malformed completion
        (
            f"Answer: {sample.target[0]}\nAction: CLICK",  # Missing value for CLICK is ok
            "C",
            1,
            1.0,  # CLICK only compares action type
        ),
    ]

    for completion, expected_value, expected_answered, expected_action_f1 in test_cases:
        state = TaskState(
            messages=[],
            output=ModelOutput(
                choices=[
                    ChatCompletionChoice(
                        message=ChatMessageAssistant(content=completion)
                    )
                ]
            ),
            choices=sample.choices,
            model="test-model",
            sample_id=sample.id,
            epoch=0,
            input=sample.input,
            metadata=sample.metadata,
        )
        score = await scorer_fn(
            state,
            target=Target(target=sample.target),
        )
        assert score.value == expected_value
        assert score.metadata["is_answered"] == expected_answered
        assert score.metadata["action_f1"] == expected_action_f1


def test_extract_answer_variations():
    """Test answer extraction with different formats."""
    # Standard format
    completion = "Answer: B\nAction: SELECT\nValue: Submit"
    result = extract_answer(completion)
    assert result.letter == "B"
    assert result.action == "SELECT"
    assert result.value == "Submit"

    # Just letter
    completion = "Answer: B"
    result = extract_answer(completion)
    assert result.letter == "B"
    assert result.action is None
    assert result.value is None

    # With quotes
    completion = 'Answer: "B"'
    result = extract_answer(completion)
    assert result.letter == "B"
    assert result.action is None
    assert result.value is None

    # With parentheses
    completion = "Answer: (C)"
    result = extract_answer(completion)
    assert result.letter == "C"
    assert result.action is None
    assert result.value is None

    # With brackets
    completion = "Answer: [D]"
    result = extract_answer(completion)
    assert result.letter == "D"
    assert result.action is None
    assert result.value is None

    # With text
    completion = "Answer: Option E"
    result = extract_answer(completion)
    assert result.letter == "E"
    assert result.action is None
    assert result.value is None

    # With mixed formatting
    completion = 'Answer: "Option [F]"'
    result = extract_answer(completion)
    assert result.letter == "F"
    assert result.action is None
    assert result.value is None

    # With explanation
    completion = """Let me analyze the options.
The best choice would be:
Answer: B
Action: TYPE
Value: Submit button
This is because..."""
    result = extract_answer(completion)
    assert result.letter == "B"
    assert result.action == "TYPE"
    assert result.value == "Submit button"


@pytest.mark.asyncio
async def test_step_and_task_success_rates():
    """Test step success rate and task success rate metrics."""
    _, sample_100 = load_test_data()
    scorer_fn = mind2web_scorer()

    # Group samples by task_id to find multi-step tasks
    task_samples = {}
    for raw_sample in sample_100:
        sample: Sample = Sample.model_validate(raw_sample)
        metadata = sample.metadata_as(Mind2WebMetadata)
        task_id = metadata.task_id
        if task_id not in task_samples:
            task_samples[task_id] = []
        task_samples[task_id].append(sample)

    # Filter for tasks with multiple steps
    multi_step_tasks = {
        task_id: samples
        for task_id, samples in task_samples.items()
        if len(samples) >= 2
    }

    # Take 3 tasks to test different scenarios
    test_tasks = list(multi_step_tasks.items())[:3]
    scores = []

    # Test each task with different success patterns
    for task_index, (task_id, samples) in enumerate(test_tasks):
        # Test first 2 steps of each task
        for step_index, sample in enumerate(samples[:2]):
            metadata = sample.metadata_as(Mind2WebMetadata)
            answer = metadata.answer
            answer_str = (
                f"{answer.letter}.\nAction: {answer.action}\nValue: {answer.value}"
            )

            # Create completion based on test scenario
            if task_index == 0 or (task_index == 1 and step_index == 0):
                # Task 1: All steps successful
                completion = f"Answer: {answer_str}"  # Prepend Answer: to the target
            elif task_index == 1 and step_index == 1:
                # Task 2: First step succeeds, second fails
                completion = "Answer: none"  # Wrong answer
            else:
                # Task 3: All steps fail
                completion = "Answer: none"  # Wrong answer

            # Create proper TaskState with all metadata
            state = TaskState(
                messages=[],
                output=ModelOutput(
                    choices=[
                        ChatCompletionChoice(
                            message=ChatMessageAssistant(content=completion)
                        )
                    ]
                ),
                choices=sample.choices,
                model="test-model",
                sample_id=sample.id,
                epoch=0,
                input=sample.input,
                metadata=sample.metadata,
            )

            # Score the step
            score = await scorer_fn(
                state,
                target=Target(target=sample.target),
            )

            # Verify expected scores
            if task_index == 0:
                # Task 1: All steps successful
                assert score.metadata["element_acc"] == 1.0
                assert score.metadata["action_f1"] == 1.0
            elif task_index == 1 and step_index == 0:
                # Task 2: First step successful
                assert score.metadata["element_acc"] == 1.0
                assert score.metadata["action_f1"] == 1.0
            else:
                # All other cases fail
                assert score.metadata["element_acc"] == 0.0
                assert score.metadata["action_f1"] == 0.0

            scores.append(SampleScore(score=score))

    # Test step success rate
    step_rate = step_success_rate()(scores)
    # 3 successful steps (2 from task 1, 1 from task 2) out of 6 total steps
    assert step_rate == 0.5

    # Test task success rate
    task_rate = task_success_rate()(scores)
    # Only task 1 is fully successful (both steps have element_acc=1.0 and action_f1=1.0)
    assert task_rate == 1.0 / 3.0  # Only task 1 is fully successful out of 3 tasks


@pytest.mark.asyncio
async def test_mind2web_scoring_none_of_above():
    """Test scoring with None of the above."""
    sample_10, _ = load_test_data()

    # Create a sample with "None of the above" target
    base_sample = Sample.model_validate(sample_10[0])  # Use first sample as base

    # Create a modified sample for "None of the above"
    none_sample = Sample(
        id=base_sample.id,
        input=base_sample.input,
        choices=base_sample.choices,
        target="A.",  # None of the above is always option A
        metadata={
            **base_sample.metadata,
            "answer": {
                "letter": "A",
                "action": None,
                "value": None,
            },
        },
    )

    scorer_fn = mind2web_scorer()

    # Test exact match
    completion = "Answer: A"
    state = TaskState(
        messages=[],
        output=ModelOutput(
            choices=[
                ChatCompletionChoice(message=ChatMessageAssistant(content=completion))
            ]
        ),
        choices=none_sample.choices,
        model="test-model",
        sample_id=none_sample.id,
        epoch=0,
        input=none_sample.input,
        metadata=none_sample.metadata,
    )

    score = await scorer_fn(
        state,
        target=Target(target=none_sample.target),
    )

    # Verify scoring - for None of the above, we only check the letter
    assert score.value == "C"  # Perfect match
    assert score.metadata["element_acc"] == 1.0
    assert score.metadata["is_answered"] == 1
    assert score.metadata["action_f1"] == 1.0  # Should be 1.0 for None of the above
