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

from .decorator import make_data_last2

T = TypeVar('T')
U = TypeVar('U')
V = TypeVar('V')


@overload
def zip_with(first: Iterable[T], second: Iterable[U], fn: Callable[[T, U], V], /) -> Iterable[V]: ...


@overload
def zip_with(second: Iterable[U], fn: Callable[[T, U], V], /) -> Callable[[Iterable[T]], Iterable[tuple[T, U]]]: ...
@overload
def zip_with(fn: Callable[[T, U], V], /) -> Callable[[Iterable[T], Iterable[U]], Iterable[tuple[T, U]]]: ...


@make_data_last2
def zip_with(
    first: Iterable[T],
    second: Iterable[U],
    function: Callable[[T, U], V],
    /,
    *,
    strict: bool = False,
) -> Iterable[V]:
    """
    Yields the result of applying the function to pairs of elements from two iterables.

    Parameters
    ----------
    first: Iterable[T]
        First iterable (positional-only).
    second: Iterable[U]
        Second iterable (positional-only).
    function: Callable[[T, U], V]
        Function to apply to pairs of elements (positional-only).
    strict: bool
        Whether to raise StopIteration if the iterables are of different lengths (keyword-only, optional).

    Yields
    ------
    V
        Result of applying the function to pairs of elements from two iterables.

    Examples
    --------
    Data first:
    >>> list(R.zip_with(['1', '2', '3'], ['a', 'b', 'c'], R.add))
    ['1a', '2b', '3c']

    Data last:
    >>> R.pipe(['1', '2', '3'], R.zip_with(['a', 'b', 'c'], R.add), list)
    ['1a', '2b', '3c']

    """
    for i, j in zip(first, second, strict=strict):
        yield function(i, j)
