# 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, name: str = None):
        if dataset_name not in self.datasets:
            raise ValueError(
                f"Dataset '{dataset_name}' nicht registriert."
            )

        dataset_path = self.datasets[dataset_name]
        try:
            # name erlaubt das Laden von Unterordnern/Konfigurationen, falls nötig
            dataset = load_dataset(
                dataset_path,
                name=name,
                token=self.hf_token
            )
            return dataset
        except Exception as e:
            raise RuntimeError(f"Fehler beim Laden von {dataset_name}: {str(e)}")

    def load_split(self, dataset_name: str, split: str = None, name: str = None):
        dataset = self.load(dataset_name, name=name)

        # Falls das Dataset direkt ein einzelnes Dataset ist (kein Dict mit Splits)
        if not hasattr(dataset, "keys"):
            return dataset

        if split is None:
            split = list(dataset.keys())[0]

        if split not in dataset:
            raise ValueError(
                f"Split '{split}' existiert nicht. Verfügbar: {list(dataset.keys())}"
            )

        return dataset[split]