import requests

from inspect_evals.livebench import LIVEBENCH_DATASETS

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


# https://huggingface.co/docs/dataset-viewer/valid
def test_datasets_are_valid():
    """Test that the Livebench datasets URLs are valid."""
    for ds in LIVEBENCH_DATASETS:
        response = requests.get(
            "https://datasets-server.huggingface.co/is-valid?dataset="
            + f"livebench/{ds}"
        )
        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


DATASETS_WITH_GROUND_TRUTH = [
    "math",
    "reasoning",
    "language",
    "data_analysis",
]


# https://huggingface.co/docs/dataset-viewer/info
def test_dataset_headers_with_ground_truth():
    """Test that Livebench datasets have required columns"""
    required_columns = [
        "question_id",
        "turns",
        "ground_truth",
        "category",
        "task",
        "livebench_release_date",
        "livebench_removal_date",
    ]

    for ds_name in DATASETS_WITH_GROUND_TRUTH:
        response = requests.get(
            f"https://datasets-server.huggingface.co/info?dataset=livebench/{ds_name}"
        )
        assert response.status_code == 200

        data = response.json()
        features = data["dataset_info"]["default"]["features"].keys()

        for column in required_columns:
            assert column in features, (
                f"Dataset {ds_name} is missing required column: {column}"
            )


def test_coding_dataset_headers():
    """Test that the coding dataset has the required columns"""
    required_columns = [
        "question_id",
        "turns",
        "category",
        "task",
        "question_title",
        "partial_solution",
        "public_test_cases",
        "private_test_cases",
        "original_json",
        "livebench_release_date",
        "livebench_removal_date",
    ]
    response = requests.get(
        "https://datasets-server.huggingface.co/info?dataset=livebench/coding"
    )
    assert response.status_code == 200

    data = response.json()
    features = data["dataset_info"]["default"]["features"].keys()
    assert "ground_truth" not in features
    for column in required_columns:
        assert column in features, (
            f"Dataset coding is missing required column: {column}"
        )


def test_instruction_following_dataset_headers():
    """Test that the instruction following dataset has the required columns"""
    required_columns = [
        "question_id",
        "turns",
        "category",
        "task",
        "instruction_id_list",
        "kwargs",
        "livebench_release_date",
        "livebench_removal_date",
    ]
    response = requests.get(
        "https://datasets-server.huggingface.co/info?dataset=livebench/instruction_following"
    )
    assert response.status_code == 200

    data = response.json()
    features = data["dataset_info"]["default"]["features"].keys()
    assert "ground_truth" not in features
    for column in required_columns:
        assert column in features, (
            f"Dataset instruction following is missing required column: {column}"
        )
