import tempfile
import traceback
import warnings
from abc import abstractmethod
from copy import deepcopy
from typing import Any, Callable, Dict, Generator, Iterable, List

from datasets import Dataset, DatasetDict, IterableDataset, IterableDatasetDict

from .dataclass import Dataclass, OptionalField
from .generator_utils import CopyingReusableGenerator, ReusableGenerator
from .logging_utils import get_logger
from .settings_utils import get_settings

settings = get_settings()
logger = get_logger()


class Stream(Dataclass):
    @abstractmethod
    def __iter__(self):
        pass

    @abstractmethod
    def peek(self):
        pass

    @abstractmethod
    def take(self, n):
        pass

    @abstractmethod
    def set_copying(self, copying: bool):
        pass


class ListStream(Stream):
    instances_list: List[Dict[str, Any]]
    copying: bool = False

    def __iter__(self):
        if self.copying:
            return iter(deepcopy(self.instances_list))
        return iter(self.instances_list)

    def peek(self):
        return next(iter(self))

    def take(self, n) -> Generator:
        for i, instance in enumerate(self.instances_list):
            if i >= n:
                break
            yield instance

    def set_copying(self, copying: bool):
        self.copying = copying


class GeneratorStream(Stream):
    """A class for handling streaming data in a customizable way.

    This class provides methods for generating, caching, and manipulating streaming data.

    Attributes:
        generator (function): A generator function for streaming data. :no-index:
        gen_kwargs (dict, optional): A dictionary of keyword arguments for the generator function. :no-index:
        caching (bool): Whether the data is cached or not. :no-index:
    """

    generator: Callable
    gen_kwargs: Dict[str, Any] = OptionalField(default_factory=dict)
    caching: bool = False
    copying: bool = False

    def _get_initiator(self):
        """Private method to get the correct initiator based on the streaming and caching attributes.

        Returns:
            function: The correct initiator function.
        """
        if self.caching:
            return Dataset.from_generator

        if self.copying:
            return CopyingReusableGenerator

        return ReusableGenerator

    def _get_stream(self):
        """Private method to get the stream based on the initiator function.

        Returns:
            object: The stream object.
        """
        return self._get_initiator()(self.generator, gen_kwargs=self.gen_kwargs)

    def __iter__(self):
        return iter(self._get_stream())

    def peek(self):
        return next(iter(self))

    def take(self, n):
        for i, instance in enumerate(self):
            if i >= n:
                break
            yield instance

    def set_copying(self, copying: bool):
        self.copying = copying


class FaultyStreamError(Exception):
    """Base class for all stream-related exceptions."""

    pass


class MissingStreamError(FaultyStreamError):
    """Raised when a required stream is missing."""

    pass


class EmptyStreamError(FaultyStreamError):
    """Raised when a stream is unexpectedly empty."""

    pass


def eager_failed():
    traceback.print_exc()
    warnings.warn(
        "The eager execution has failed due to the error above.", stacklevel=2
    )


class DynamicStream(Stream):
    generator: Callable
    gen_kwargs: Dict[str, Any] = OptionalField(default_factory=dict)
    caching: bool = False
    copying: bool = False

    def __post_init__(self):
        self.stream = None
        if settings.use_eager_execution:
            try:
                instances_list = []
                for instance in self.generator(**self.gen_kwargs):
                    instances_list.append(instance)
                self.stream = ListStream(
                    instances_list=instances_list, copying=self.copying
                )
            except FaultyStreamError:
                eager_failed()
            except RuntimeError as e:
                if isinstance(e.__cause__, FaultyStreamError):
                    eager_failed()
                else:
                    raise e

        if self.stream is None:
            self.stream = GeneratorStream(
                generator=self.generator,
                gen_kwargs=self.gen_kwargs,
                caching=self.caching,
                copying=self.copying,
            )

    def __iter__(self):
        return self.stream.__iter__()

    def peek(self):
        return self.stream.peek()

    def take(self, n):
        return self.stream.take(n)

    def set_copying(self, copying: bool):
        self.stream.set_copying(copying)


class MultiStream(dict):
    """A class for handling multiple streams of data in a dictionary-like format.

    This class extends dict and its values should be instances of the Stream class.

    Attributes:
        data (dict): A dictionary of Stream objects.
    """

    def __init__(self, data=None):
        """Initializes the MultiStream with the provided data.

        Args:
            data (dict, optional): A dictionary of Stream objects. Defaults to None.

        Raises:
            AssertionError: If the values are not instances of Stream or keys are not strings.
        """
        for key, value in data.items():
            isinstance(value, Stream), "MultiStream values must be Stream"
            isinstance(key, str), "MultiStream keys must be strings"
        super().__init__(data)

    def get_generator(self, key) -> Generator:
        """Gets a generator for a specified key.

        Args:
            key (str): The key for the generator.

        Yields:
            object: The next value in the stream.
        """
        yield from self[key]

    def set_caching(self, caching: bool):
        for stream in self.values():
            stream.caching = caching

    def set_copying(self, copying: bool):
        for stream in self.values():
            stream.set_copying(copying)

    def to_dataset(self, disable_cache=True, cache_dir=None) -> DatasetDict:
        with tempfile.TemporaryDirectory() as dir_to_be_deleted:
            cache_dir = dir_to_be_deleted if disable_cache else cache_dir
            return DatasetDict(
                {
                    key: Dataset.from_generator(
                        self.get_generator,
                        keep_in_memory=disable_cache,
                        cache_dir=cache_dir,
                        gen_kwargs={"key": key},
                    )
                    for key in self.keys()
                }
            )

    def to_iterable_dataset(self) -> IterableDatasetDict:
        return IterableDatasetDict(
            {
                key: IterableDataset.from_generator(
                    self.get_generator, gen_kwargs={"key": key}
                )
                for key in self.keys()
            }
        )

    def __setitem__(self, key, value):
        assert isinstance(value, Stream), "StreamDict values must be Stream"
        assert isinstance(key, str), "StreamDict keys must be strings"
        super().__setitem__(key, value)

    @classmethod
    def from_generators(
        cls, generators: Dict[str, ReusableGenerator], caching=False, copying=False
    ):
        """Creates a MultiStream from a dictionary of ReusableGenerators.

        Args:
            generators (Dict[str, ReusableGenerator]): A dictionary of ReusableGenerators.
            caching (bool, optional): Whether the data should be cached or not. Defaults to False.
            copying (bool, optional): Whether the data should be copied or not. Defaults to False.

        Returns:
            MultiStream: A MultiStream object.
        """
        assert all(isinstance(v, ReusableGenerator) for v in generators.values())
        return cls(
            {
                key: DynamicStream(
                    generator.generator,
                    gen_kwargs=generator.gen_kwargs,
                    caching=caching,
                    copying=copying,
                )
                for key, generator in generators.items()
            }
        )

    @classmethod
    def from_iterables(
        cls, iterables: Dict[str, Iterable], caching=False, copying=False
    ):
        """Creates a MultiStream from a dictionary of iterables.

        Args:
            iterables (Dict[str, Iterable]): A dictionary of iterables.
            caching (bool, optional): Whether the data should be cached or not. Defaults to False.
            copying (bool, optional): Whether the data should be copied or not. Defaults to False.

        Returns:
            MultiStream: A MultiStream object.
        """
        return cls(
            {
                key: DynamicStream(
                    iterable.__iter__,
                    caching=caching,
                    copying=copying,
                )
                for key, iterable in iterables.items()
            }
        )