from urllib.parse import quote

import pytest
import requests

from inspect_evals.sec_qa.sec_qa import (
    DATASET_PATH,
)

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

HUGGING_FACE_EXPECTED_SPLITS_V1 = {
    "dev": {
        "name": "dev",
        "num_bytes": 2538,
        "num_examples": 5,
        "dataset_name": "secqa",
    },
    "val": {
        "name": "val",
        "num_bytes": 6364,
        "num_examples": 12,
        "dataset_name": "secqa",
    },
    "test": {
        "name": "test",
        "num_bytes": 53627,
        "num_examples": 110,
        "dataset_name": "secqa",
    },
}

HUGGING_FACE_EXPECTED_SPLITS_V2 = {
    "dev": {
        "name": "dev",
        "num_bytes": 3387,
        "num_examples": 5,
        "dataset_name": "secqa",
    },
    "val": {
        "name": "val",
        "num_bytes": 7494,
        "num_examples": 10,
        "dataset_name": "secqa",
    },
    "test": {
        "name": "test",
        "num_bytes": 68458,
        "num_examples": 100,
        "dataset_name": "secqa",
    },
}


def test_dataset_is_valid():
    """Test that the SecQA dataset URL is valid."""
    response = requests.get(
        "https://datasets-server.huggingface.co/is-valid?dataset=" + quote(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_splits():
    """Test that the SecQA dataset has the required splits."""
    response = requests.get(
        "https://datasets-server.huggingface.co/info?dataset=" + quote(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"]
    assert "secqa_v1" in dataset_info, (
        f"SecQA V1 subset missing from {DATASET_PATH} dataset"
    )
    assert "secqa_v2" in dataset_info, (
        f"SecQA V2 subset missing from {DATASET_PATH} dataset"
    )

    assert dataset_info["secqa_v1"]["splits"] == HUGGING_FACE_EXPECTED_SPLITS_V1, (
        "Unexpected dataset splits for SecQA V1 and especially check dev which is used in 5-shot evaluation"
    )
    assert dataset_info["secqa_v2"]["splits"] == HUGGING_FACE_EXPECTED_SPLITS_V2, (
        "Unexpected dataset splits for SecQA V2 and especially check dev which is used in 5-shot evaluation"
    )


@pytest.mark.huggingface
def test_required_columns():
    """Test that the SecQA dataset has the required columns."""
    response = requests.get(
        "https://datasets-server.huggingface.co/info?dataset=" + quote(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"]
    assert "secqa_v1" in dataset_info, (
        f"SecQA V1 subset missing from {DATASET_PATH} dataset"
    )
    assert "secqa_v2" in dataset_info, (
        f"SecQA V2 subset missing from {DATASET_PATH} dataset"
    )

    secqa_v1_features = dataset_info["secqa_v1"]["features"]
    assert "Question" in secqa_v1_features, "Question column missing from SecQA V1"
    assert "A" in secqa_v1_features, "A column missing from SecQA V1"
    assert "B" in secqa_v1_features, "B column missing from SecQA V1"
    assert "C" in secqa_v1_features, "C column missing from SecQA V1"
    assert "D" in secqa_v1_features, "D column missing from SecQA V1"
    assert "Answer" in secqa_v1_features, "Answer column missing from SecQA V1"

    secqa_v2_features = dataset_info["secqa_v2"]["features"]
    assert "Question" in secqa_v2_features, "Question column missing from SecQA V2"
    assert "A" in secqa_v2_features, "A column missing from SecQA V2"
    assert "B" in secqa_v2_features, "B column missing from SecQA V2"
    assert "C" in secqa_v2_features, "C column missing from SecQA V2"
    assert "D" in secqa_v2_features, "D column missing from SecQA V2"
    assert "Answer" in secqa_v2_features, "Answer column missing from SecQA V2"
