from collections.abc import Callable, Iterable
from typing import Any, TypeVar, overload

import remedapy as R

from .decorator import make_data_last

T = TypeVar('T')


@overload
def for_each(count: Iterable[T], fn: Callable[[], Any], /) -> list[T]: ...


@overload
def for_each(count: Iterable[T], fn: Callable[[T], Any], /) -> list[T]: ...


@overload
def for_each(fn: Callable[[], Any], /) -> Callable[[Iterable[T]], list[T]]: ...


@overload
def for_each(fn: Callable[[T], Any], /) -> Callable[[Iterable[T]], list[T]]: ...


# TODO: index and whole array?
@make_data_last
def for_each(iterable: Iterable[T], callbackfn: Callable[[], Any] | Callable[[T], Any], /) -> Iterable[T]:
    """
    Applies the given function to each element of the iterable and then yields the element.

    Parameters
    ----------
    iterable: Iterable[T]
        The iterable to flatten (positional-only).
    callbackfn: Callable[[], Any] |
                Callable[[T], Any]
        The function to apply to each element of the iterable (positional-only).

    Yields
    ------
    T
        The elements of the original iterable.

    Examples
    --------
    Data first:
    >>> x = []
    >>> result = list(R.for_each([1, 2, 3], lambda i: x.append(i*10)))
    >>> x
    [10, 20, 30]
    >>> result
    [1, 2, 3]

    Data last:
    >>> x = []
    >>> result = list(R.for_each(lambda i: x.append(i*10))([1, 2, 3]))
    >>> x
    [10, 20, 30]
    >>> result
    [1, 2, 3]

    """
    apply_fn = R.apply(callbackfn)
    for i in iterable:
        apply_fn(i)
        yield i
