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

import remedapy as R

from .decorator import make_data_last

T = TypeVar('T')


@overload
def times(n: int, function: Callable[[], T], /) -> Iterable[T]: ...


@overload
def times(n: int, function: Callable[[int], T], /) -> Iterable[T]: ...


@overload
def times(function: Callable[[], T], /) -> Callable[[int], Iterable[T]]: ...


@overload
def times(function: Callable[[int], T], /) -> Callable[[int], Iterable[T]]: ...


@make_data_last
def times(n: int, function: Callable[[], T] | Callable[[int], T], /) -> Iterable[T]:
    """
    Yields the result of applying function n times, passing the index as an argument.

    If the function accepts 0 arguments, the index is not passed.

    Parameters
    ----------
    n : int
        Number of times to apply the function.
    function : Callable[[], T] | Callable[[int], T]
        Function to apply (positional-only).

    Returns
    -------
    Iterable[T]
        Result of applying function to data.

    Examples
    --------
    Data first:
    >>> list(R.times(5, R.identity()))
    [0, 1, 2, 3, 4]
    >>> list(R.times(5, R.constant('a')))
    ['a', 'a', 'a', 'a', 'a']

    Data last:
    >>> list(R.times(R.identity())(5))
    [0, 1, 2, 3, 4]
    >>> list(R.times(R.constant('a'))(5))
    ['a', 'a', 'a', 'a', 'a']

    """
    apply = R.apply(function)
    for i in range(n):
        yield apply(i)
