import requests

from inspect_evals.lab_bench.lab_bench import DATASET_PATH, DatasetSubsets

HUGGING_FACE_IS_VALID_RESPONSE = {
    "preview": True,
    "viewer": True,
    "search": True,
    "filter": True,
    "statistics": True,
}


def test_dataset_is_valid():
    """Test that the LAB-Bench dataset URL is valid."""
    response = requests.get(
        "https://datasets-server.huggingface.co/is-valid?dataset=" + DATASET_PATH
    )
    assert response.status_code == 200, (
        f"Hugging Face dataset `/is-valid` returned status code {response.status_code} check documentation at `https://huggingface.co/docs/dataset-viewer/en/valid`"
    )
    assert response.json() == HUGGING_FACE_IS_VALID_RESPONSE


def test_required_subsets_and_columns():
    """Test that the LAB-Bench dataset has the required subsets."""
    response = requests.get(
        "https://datasets-server.huggingface.co/info?dataset=" + DATASET_PATH
    )
    assert response.status_code == 200, (
        f"Hugging Face dataset `/info` returned status code {response.status_code} check documentation at `https://huggingface.co/docs/dataset-viewer/en/info`"
    )

    data = response.json()
    assert "dataset_info" in data, (
        f"Hugging Face dataset `/info` returned unexpected json {data} check documentation at `https://huggingface.co/docs/dataset-viewer/en/info`"
    )

    dataset_info = data["dataset_info"]
    for dataset_subset in DatasetSubsets:
        subset = dataset_subset.value
        assert subset in dataset_info, (
            f"LAB-Bench {subset} subset missing from {DATASET_PATH} dataset"
        )

        subset_features = dataset_info[subset]["features"]

        # Check common columns
        assert "question" in subset_features, (
            f"question column missing from LAB-Bench {subset}"
        )
        assert "ideal" in subset_features, (
            f"ideal column missing from LAB-Bench {subset}"
        )
        assert "distractors" in subset_features, (
            f"distractors column missing from LAB-Bench {subset}"
        )

        # Check additional subset-specific columns
        if subset == DatasetSubsets.LitQA2.value:
            assert "sources" in subset_features, (
                f"sources column missing from LAB-Bench {subset}"
            )

        if subset == DatasetSubsets.SuppQA.value:
            assert "source" in subset_features, (
                f"source column missing from LAB-Bench {subset}"
            )
            assert "paper-title" in subset_features, (
                f"paper-title column missing from LAB-Bench {subset}"
            )

        if subset == DatasetSubsets.FigQA.value:
            assert "figure" in subset_features, (
                f"figure column missing from LAB-Bench {subset}"
            )

        if subset == DatasetSubsets.TableQA.value:
            assert "tables" in subset_features, (
                f"tables column missing from LAB-Bench {subset}"
            )

        if subset == DatasetSubsets.ProtocolQA.value:
            assert "protocol" in subset_features, (
                f"protocol column missing from LAB-Bench {subset}"
            )
