import pytest
from inspect_ai.dataset import Dataset, Sample

from inspect_evals.mind2web.dataset import (
    Answer,
    Mind2WebMetadata,
    get_dataset,
)


@pytest.mark.dataset_download
def test_get_dataset():
    # Get dataset
    dataset = get_dataset(limit=100)

    # Verify it's a Dataset instance
    assert isinstance(dataset, Dataset)

    # Verify dataset has samples
    assert len(dataset.samples) > 0

    # Verify sample structure
    sample: Sample = dataset.samples[4]
    assert sample.id is not None
    assert sample.input is not None
    assert isinstance(sample.choices, list)
    assert len(sample.choices) > 3, "Dataset must have more than 3 choices"
    assert sample.target is not None

    # Verify metadata structure
    metadata = sample.metadata_as(Mind2WebMetadata)
    assert metadata.final_html is not None
    assert metadata.previous_actions is not None
    assert isinstance(metadata.target_idx, int)
    assert metadata.confirmed_task is not None
    assert metadata.task_id is not None
    assert metadata.annotation_id is not None
    assert metadata.action_uid is not None

    # If there's an answer, verify its structure
    if metadata.answer:
        answer = metadata.answer
        assert isinstance(answer, Answer)
        assert answer.letter in "ABCDEF"
        assert answer.action in (None, "CLICK", "SELECT", "TYPE")
