# core/loader.py

from datasets import load_dataset
import os


class DatasetLoader:

    def __init__(self):

        print("DatasetLoader geladen")

        self.hf_token = os.getenv("HF_TOKEN", None)

        self.datasets = {
            "matrix_game": "Skywork/Matrix-Game-3.0",
            "framework_storage": "sky-meilin/SkyAiFramework-storage"
        }

    def list_datasets(self):

        return list(self.datasets.keys())

    def load(self, dataset_name: str):

        if dataset_name not in self.datasets:

            raise ValueError(
                f"Dataset '{dataset_name}' nicht registriert."
            )

        dataset_path = self.datasets[dataset_name]

        dataset = load_dataset(
            dataset_path,
            token=self.hf_token
        )

        return dataset

    def load_split(
        self,
        dataset_name: str,
        split: str = None
    ):

        dataset = self.load(dataset_name)

        if split is None:

            split = list(dataset.keys())[0]

        return dataset[split]