# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.


from .huggingface import HuggingFaceDataset
from .interface import DatasetSource, IDataset


class DatasetRegistry:
    _dataset_registry: dict[str, DatasetSource] = {}

    @staticmethod
    def register_dataset(dataset_id: str, dataset_source: DatasetSource):
        DatasetRegistry._dataset_registry[dataset_id] = dataset_source

    @staticmethod
    def get_dataset(dataset_id: str) -> IDataset:
        config = DatasetRegistry._dataset_registry[dataset_id]
        if config.type == "huggingface":
            hf_dataset = HuggingFaceDataset(dataset_id, config)
            hf_dataset.load()
            return hf_dataset
        else:
            raise ValueError(f"Unsupported dataset type: {config.type}")


# # extracted from https://huggingface.co/datasets/MMMU/MMMU_Pro
# # packaged/converted into prompts by https://huggingface.co/datasets/llamastack/MMMU_Pro/blob/main/hf2llamastack.py
# DATASET_REGISTRY["MMMU_Pro_standard"] = HuggingfaceDataSource(
#     args={
#         "path": "llamastack/MMMU_Pro",
#         "split": "validation",
#         "name": "standard (10 options)",
#     }
# )

# # extracted from https://huggingface.co/datasets/MMMU/MMMU_Pro
# # packaged/converted into prompts by https://huggingface.co/datasets/llamastack/MMMU_Pro/blob/main/hf2llamastack.py
# DATASET_REGISTRY["MMMU_Pro_vision"] = HuggingfaceDataSource(
#     args={"path": "llamastack/MMMU_Pro", "split": "validation", "name": "vision"}
# )

# # extracted from https://huggingface.co/datasets/lmms-lab/ai2d
# # packaged/converted into prompts by https://huggingface.co/datasets/llamastack/ai2d/blob/main/hf2llamastack.py
# DATASET_REGISTRY["ai2d"] = HuggingfaceDataSource(
#     args={
#         "path": "llamastack/ai2d",
#         "split": "test",
#     }
# )
