from typing import Callable, Any
from datasets import load_dataset
import utils.configs_loader as ProjectConfigs

# PUBLIC CLASS

class StreamingDataset:
    def __init__(
        self,
        path: str,
        split: str = "train",
        mapping: Callable[[dict], dict] | None = None,
        shuffle: bool = True,
        buffer_size: int = 10_000,
        seed: int = 42,
    ):
        self.dataset_info = ProjectConfigs.get(("DATASET", path))

        self.split = split
        self.mapping = mapping
        self.should_shuffle = shuffle
        self.buffer_size = buffer_size
        self.seed = seed

        self.dataset = None

    def take(self, amount: int):
        if self.dataset is None:
            self.load()

        return self.dataset.take(amount)

    def shuffle(self):
        if self.dataset is None:
            raise RuntimeError("Dataset must be loaded before shuffle().")

        self.dataset = self.dataset.shuffle(
            seed=self.seed,
            buffer_size=self.buffer_size,
        )

        return self.dataset

    def set_map(self, method: Callable[[dict], dict] | None):
        self.mapping = method

    def load(self):
        self.dataset = load_dataset(
            path=self.dataset_info.Path,
            name=self.dataset_info.Name,
            split=self.split,
            streaming=True,
        )

        if self.should_shuffle:
            self.shuffle()

        if self.mapping is not None:
            self.dataset = self.dataset.map(self.mapping)

        return self.dataset