from typing import Callable, Iterator, Optional, TypeVar

from ray.data._internal.block_batching.interfaces import ResolvedBlock
from ray.data._internal.block_batching.util import (
    _MappingIterator,
    blocks_to_batches,
    collate,
    format_batches,
)
from ray.data._internal.stats import DatasetStats
from ray.data.block import Block, DataBatch

T = TypeVar("T")


def batch_blocks(
    blocks: Iterator[Block],
    *,
    stats: Optional[DatasetStats] = None,
    batch_size: Optional[int] = None,
    batch_format: str = "default",
    drop_last: bool = False,
    collate_fn: Optional[Callable[[DataBatch], DataBatch]] = None,
    shuffle_buffer_min_size: Optional[int] = None,
    shuffle_seed: Optional[int] = None,
    ensure_copy: bool = False,
) -> Iterator[DataBatch]:
    """Create formatted batches of data from 1 or more blocks.

    This function takes in an iterator of already fetched blocks. Consequently, this
    function doesn't support block prefetching.
    """
    # TODO: make stage timings optional at _BatchingIterator so this
    # shim can be removed. map() avoids holding block references.
    wrapped_blocks = map(lambda b: ResolvedBlock(block=b), blocks)

    # Build the processing pipeline
    batch_iter = format_batches(
        blocks_to_batches(
            block_iter=wrapped_blocks,
            stats=stats,
            batch_size=batch_size,
            drop_last=drop_last,
            shuffle_buffer_min_size=shuffle_buffer_min_size,
            shuffle_seed=shuffle_seed,
            ensure_copy=ensure_copy,
        ),
        batch_format=batch_format,
        stats=stats,
        ensure_copy=ensure_copy,
    )

    if collate_fn is not None:
        batch_iter = collate(batch_iter, collate_fn=collate_fn, stats=stats)

    return _UserTimingIterator(
        _MappingIterator(batch_iter, lambda batch: batch.data), stats
    )


class _UserTimingIterator(Iterator[DataBatch]):
    def __init__(self, iter: Iterator[DataBatch], stats: Optional[DatasetStats]):
        self._iter = iter
        self._stats = stats
        self._active_timer = None

    def __iter__(self) -> Iterator[DataBatch]:
        return self

    def __next__(self) -> DataBatch:
        # Since we're tracking time spent in user-code, we stop
        # the timer immediately when `__next__` is called
        self._stop_timer()

        try:
            res = next(self._iter)
            # Reset timer and return
            #
            # NOTE: It's crucial that we reset the timer only after we
            #       retrieved the result to avoid starting the timer before
            #       we retrieve the next value
            self._reset_timer()
            return res
        except StopIteration:
            self._stop_timer()
            raise

    def _stop_timer(self):
        if not self._stats:
            return

        if self._active_timer:
            self._active_timer.__exit__(None, None, None)
            self._active_timer = None

    def _reset_timer(self):
        if not self._stats:
            return

        self._active_timer = self._stats.iter_user_s.timer()
        self._active_timer.__enter__()