# Copyright (c) 2025, NVIDIA CORPORATION.
# SPDX-License-Identifier: BSD-3-Clause

from contextlib import contextmanager
from dataclasses import dataclass
from typing import (
    Any,
    Generator,
    Generic,
    Iterable,
    List,
    Optional,
    Sequence,
    Tuple,
    TypeVar,
    Union,
)

from megatron.energon.flavors.base_dataset import MergedState, Sample, SavableDataset, State
from megatron.energon.worker import WorkerConfig

T = TypeVar("T")
T_sample = TypeVar("T_sample", covariant=True)
T_sample_out = TypeVar("T_sample_out", covariant=True)
T_sample_in = TypeVar("T_sample_in", covariant=True)


@dataclass
class BaseSingleWrapperState(State):
    """Base class for dataset states."""

    #: The class name of the dataset saving this state for assertion
    dataset_type: str

    #: State of the inner dataset
    dataset_state: State


@dataclass
class BaseSingleWrapperMergedState(MergedState):
    """Base class for dataset states."""

    #: The class name of the dataset saving/merging this state for assertion
    dataset_type: str

    #: State of the inner dataset
    dataset_state: MergedState


class BaseWrapperDataset(SavableDataset[T_sample], Generic[T_sample]):
    """Base class for dataset wrappers. All dataset wrappers should derive from this. A dataset
    wrapper takes one dataset and modifies its samples to make a new dataset. This can be for
    shuffling samples or applying custom functions to the data. Some wrappers only modify the
    length of the dataset or how it's repeated."""

    def __init__(
        self,
        dataset: Union[SavableDataset, Iterable[SavableDataset]],
        *,
        worker_config: WorkerConfig,
    ):
        super().__init__(worker_config=worker_config)

        if isinstance(dataset, SavableDataset):
            dataset = [dataset]

        for d in dataset:
            # Check that the dataset worker configs are the same as the wrapper worker config
            assert (
                d.worker_config == self.worker_config
            ), "Dataset and wrapper worker configs must match."


class BaseSingleWrapperDataset(
    BaseWrapperDataset[T_sample_out], Generic[T_sample_in, T_sample_out]
):
    """Base class for dataset wrappers that wrap a single dataset. Provides default implementations
    for saving and restoring the dataset state."""

    dataset: SavableDataset[T_sample_in]

    def __init__(self, dataset: SavableDataset[T_sample_in], *, worker_config: WorkerConfig):
        super().__init__(dataset, worker_config=worker_config)
        assert isinstance(dataset, SavableDataset)
        self.dataset = dataset

    def save_state(self) -> BaseSingleWrapperState:
        return BaseSingleWrapperState(
            dataset_type=type(self).__name__,
            dataset_state=self.dataset.save_state(),
        )

    def merge_states(
        self, states: Sequence[Optional[BaseSingleWrapperState]]
    ) -> BaseSingleWrapperMergedState:
        assert all(s is None or isinstance(s, BaseSingleWrapperState) for s in states)
        assert all(s is None or s.dataset_type == type(self).__name__ for s in states)
        return BaseSingleWrapperMergedState(
            dataset_type=type(self).__name__,
            dataset_state=self.dataset.merge_states(
                [None if s is None else s.dataset_state for s in states]
            ),
        )

    def restore_state(self, state: Optional[BaseSingleWrapperMergedState]) -> None:
        if state is None:
            self.dataset.restore_state(None)
        else:
            assert isinstance(state, BaseSingleWrapperMergedState)
            assert state.dataset_type == type(self).__name__
            self.dataset.restore_state(state.dataset_state)

    def can_restore_sample(self) -> bool:
        return self.dataset.can_restore_sample()

    def assert_can_restore(self) -> None:
        self.dataset.assert_can_restore()

    def restore_sample(self, index: Tuple[Union[str, int, tuple], ...]) -> T_sample_out:
        return self.dataset.restore_sample(index)

    def worker_has_samples(self) -> bool:
        return self.dataset.worker_has_samples()


class SampleIndex:
    """A simple class to hold the sample index for each worker."""

    worker_config: WorkerConfig
    _sample_index: List[int]
    __rank: Optional[int] = None

    actives = 0

    def __init__(self, worker_config: WorkerConfig, *, src: Any) -> None:
        self.worker_config = worker_config
        self._sample_index = [0] * max(self.worker_config.num_workers, 1)
        self.src = src

    @property
    def _rank(self) -> int:
        if self.__rank is None:
            self.__rank = self.worker_config.rank_worker_id()
        return self.__rank

    def get_next(self) -> int:
        res = self._sample_index[self._rank]
        self._sample_index[self._rank] += 1
        return res

    @property
    def current_idx(self) -> int:
        return self._sample_index[self._rank]

    @contextmanager
    def ctx(self, sample_idx: Optional[int] = None):
        if sample_idx is None:
            sample_idx = self.get_next()
        assert WorkerConfig.active_worker_config is not None
        WorkerConfig.active_worker_config.worker_push_sample_index(sample_idx)
        # print("  " * SampleIndex.actives + f"Activated from {type(self.src).__name__}({id(self.src)}) {sample_idx} -> {WorkerConfig.active_worker_config._sample_index_stack}")
        SampleIndex.actives += 1
        try:
            yield sample_idx
        finally:
            assert WorkerConfig.active_worker_config is not None
            popped = WorkerConfig.active_worker_config.worker_pop_sample_index()
            SampleIndex.actives -= 1
            # print("  " * SampleIndex.actives + f"Deactivate from {type(self.src).__name__}({id(self.src)}) {sample_idx} -> {WorkerConfig.active_worker_config._sample_index_stack}")
            assert popped == sample_idx, f"Expected {sample_idx}, got {popped}"

    def iter_ctx(
        self,
        it: Iterable[T_sample],
        sample_idx: Optional[int] = None,
    ) -> Generator[Tuple[int, T_sample], None, None]:
        it = iter(it)
        try:
            while True:
                try:
                    with self.ctx(sample_idx) as res_sample_idx:
                        x = next(it)
                    yield res_sample_idx, x
                except StopIteration:
                    break
        finally:
            if hasattr(it, "close"):
                it.close()

    def save_state(self) -> int:
        return self._sample_index[self._rank]

    def merge_states(self, states: List[int]) -> List[int]:
        return states

    def restore_state(self, state: Optional[List[int]]) -> None:
        if state is None:
            self._sample_index = [0] * max(self.worker_config.num_workers, 1)
        else:
            assert len(state) == len(self._sample_index)
            self._sample_index = state


def get_sample_restore_key(sample: Any) -> Optional[Union[str, int]]:
    """Gets the restore key from an arbitrary sample."""
    if isinstance(sample, Sample) or hasattr(sample, "__restore_key__"):
        return sample.__restore_key__
    elif isinstance(sample, dict) and "__restore_key__" in sample:
        return sample["__restore_key__"]
    else:
        return None
