import importlib
import inspect

import pytest

from schemas_hf.base import DatasetSchema
from schemas_hf.naming import module_filename, schema_key
from schemas_hf.registry import registry


CURATED_DATASETS = [
    ("AI-MO/NuminaMath-LEAN", None),
    ("AI4Math/MathVista", None),
    ("HuggingFaceFW/fineweb", "sample-10BT"),
    ("ILSVRC/imagenet-1k", None),
    ("Kukedlc/spanish-train", None),
    ("MathLLMs/MathVision", None),
    ("Msun/modelnet40", None),
    ("NeelNanda/counterfact-tracing", None),
    ("Vezora/Tested-143k-Python-Alpaca", None),
    ("almogtavor/google-analogy-dataset", None),
    ("bigcode/the-stack", None),
    ("domenicrosati/TruthfulQA", None),
    ("euclaise/logician", None),
    ("facebook/xnli", "all_languages"),
    ("fancyzhx/yelp_polarity", "plain_text"),
    ("layoric/tiny-codes-alpaca", None),
    ("lmms-lab/LLaVA-Video-178K", "0_30_s_academic_v0_1"),
    ("mteb/HotelReviewSentimentClassification", None),
    ("mteb/TweetEmotionClassification", None),
    ("mteb/banking77", None),
    ("mteb/news", None),
    ("mteb/online_store_review_sentiment", None),
    ("mteb/restaurant_review_sentiment", None),
    ("mteb/sprintduplicatequestions-pairclassification", None),
    ("mteb/tweet_sarcasm", None),
    ("omni-research/DREAM-1K", "DREAM-1K"),
    ("openai/gsm8k", "main"),
    ("opus100", "en-es"),
    ("qiaojin/PubMedQA", "pqa_artificial"),
    ("qiaojin/PubMedQA", "pqa_labeled"),
    ("wikiann", "es"),
    ("wmt14", "de-en"),
    ("xnli", "all_languages"),
    ("yelp_polarity", "plain_text"),
]


def _first_schema_class(module):
    for _, obj in inspect.getmembers(module, inspect.isclass):
        if issubclass(obj, DatasetSchema) and obj is not DatasetSchema:
            return obj
    return None


@pytest.mark.parametrize("dataset_id,config_name", CURATED_DATASETS)
def test_curated_dataset_schema_exists_and_registered(dataset_id, config_name):
    module_name = module_filename(dataset_id, config_name).rsplit(".", 1)[0]
    module = importlib.import_module(f"schemas_hf.datasets.{module_name}")

    schema_cls = _first_schema_class(module)
    assert schema_cls is not None
    assert schema_cls.DATASET_NAME == dataset_id
    assert schema_cls.CONFIG_NAME == config_name

    registry.discover()
    key = schema_key(dataset_id, config_name)
    assert key in registry.list_names()

