from __future__ import annotations

import itertools
from collections.abc import Callable, Iterable, Iterator
from typing import TYPE_CHECKING, Any, overload

import cytoolz as cz

from .._core import IterWrapper

if TYPE_CHECKING:
    from .._results import Option
    from ._main import Iter


class BaseJoins[T](IterWrapper[T]):
    @overload
    def zip[T1](
        self,
        iter1: Iterable[T1],
        /,
        *,
        strict: bool = ...,
    ) -> Iter[tuple[T, T1]]: ...
    @overload
    def zip[T1, T2](
        self,
        iter1: Iterable[T1],
        iter2: Iterable[T2],
        /,
        *,
        strict: bool = ...,
    ) -> Iter[tuple[T, T1, T2]]: ...
    @overload
    def zip[T1, T2, T3](
        self,
        iter1: Iterable[T1],
        iter2: Iterable[T2],
        iter3: Iterable[T3],
        /,
        *,
        strict: bool = ...,
    ) -> Iter[tuple[T, T1, T2, T3]]: ...
    @overload
    def zip[T1, T2, T3, T4](
        self,
        iter1: Iterable[T1],
        iter2: Iterable[T2],
        iter3: Iterable[T3],
        iter4: Iterable[T4],
        /,
        *,
        strict: bool = ...,
    ) -> Iter[tuple[T, T1, T2, T3, T4]]: ...
    def zip(
        self,
        *others: Iterable[Any],
        strict: bool = False,
    ) -> Iter[tuple[Any, ...]]:
        """Yields n-length tuples, where n is the number of iterables passed as positional arguments.

        The i-th element in every tuple comes from the i-th iterable argument to `.zip()`.

        This continues until the shortest argument is exhausted.

        Args:
            *others (Iterable[Any]): Other iterables to zip with.
            strict (bool): If `True` and one of the arguments is exhausted before the others, raise a ValueError. Defaults to `False`.

        Returns:
            Iter[tuple[Any, ...]]: An `Iter` of tuples containing elements from the zipped Iter and other iterables.

        Example:
        ```python
        >>> import pyochain as pc
        >>> pc.Iter([1, 2]).zip([10, 20]).collect()
        Seq((1, 10), (2, 20))
        >>> pc.Iter(["a", "b"]).zip([1, 2, 3]).collect()
        Seq(('a', 1), ('b', 2))

        ```
        """
        return self._lazy(zip, *others, strict=strict)

    @overload
    def zip_longest[T2](
        self, iter2: Iterable[T2], /
    ) -> Iter[tuple[Option[T], Option[T2]]]: ...
    @overload
    def zip_longest[T2, T3](
        self, iter2: Iterable[T2], iter3: Iterable[T3], /
    ) -> Iter[tuple[Option[T], Option[T2], Option[T3]]]: ...
    @overload
    def zip_longest[T2, T3, T4](
        self,
        iter2: Iterable[T2],
        iter3: Iterable[T3],
        iter4: Iterable[T4],
        /,
    ) -> Iter[tuple[Option[T], Option[T2], Option[T3], Option[T4]]]: ...
    @overload
    def zip_longest[T2, T3, T4, T5](
        self,
        iter2: Iterable[T2],
        iter3: Iterable[T3],
        iter4: Iterable[T4],
        iter5: Iterable[T5],
        /,
    ) -> Iter[
        tuple[
            Option[T],
            Option[T2],
            Option[T3],
            Option[T4],
            Option[T5],
        ]
    ]: ...
    @overload
    def zip_longest(
        self,
        iter2: Iterable[T],
        iter3: Iterable[T],
        iter4: Iterable[T],
        iter5: Iterable[T],
        iter6: Iterable[T],
        /,
        *iterables: Iterable[T],
    ) -> Iter[tuple[Option[T], ...]]: ...
    def zip_longest(self, *others: Iterable[Any]) -> Iter[tuple[Option[Any], ...]]:
        """Return a zip Iterator who yield a tuple where the i-th element comes from the i-th iterable argument.

        Yield values until the longest iterable in the argument sequence is exhausted, and then it raises StopIteration.

        The longest iterable determines the length of the returned iterator, and will return `Some[T]` until exhaustion.

        When the shorter iterables are exhausted, they yield `NONE`.

        Args:
            *others (Iterable[Any]): Other iterables to zip with.

        Returns:
            Iter[tuple[Option[Any], ...]]: An iterable of tuples containing optional elements from the zipped iterables.

        Example:
        ```python
        >>> import pyochain as pc
        >>> pc.Iter([1, 2]).zip_longest([10]).collect()
        Seq((Some(1), Some(10)), (Some(2), NONE))
        >>> # Can be combined with try collect to filter out the NONE:
        >>> pc.Iter([1, 2]).zip_longest([10]).map(lambda x: pc.Iter(x).try_collect()).collect()
        Seq(Some(Vec(1, 10)), NONE)

        ```
        """
        from .._results import Option

        def _zip_longest(data: Iterable[T]) -> Iterator[tuple[Option[T], ...]]:
            return (
                tuple(Option.from_(t) for t in tup)
                for tup in itertools.zip_longest(data, *others, fillvalue=None)
            )

        return self._lazy(_zip_longest)

    @overload
    def product(self) -> Iter[tuple[T]]: ...
    @overload
    def product[T1](self, iter1: Iterable[T1], /) -> Iter[tuple[T, T1]]: ...
    @overload
    def product[T1, T2](
        self,
        iter1: Iterable[T1],
        iter2: Iterable[T2],
        /,
    ) -> Iter[tuple[T, T1, T2]]: ...
    @overload
    def product[T1, T2, T3](
        self,
        iter1: Iterable[T1],
        iter2: Iterable[T2],
        iter3: Iterable[T3],
        /,
    ) -> Iter[tuple[T, T1, T2, T3]]: ...
    @overload
    def product[T1, T2, T3, T4](
        self,
        iter1: Iterable[T1],
        iter2: Iterable[T2],
        iter3: Iterable[T3],
        iter4: Iterable[T4],
        /,
    ) -> Iter[tuple[T, T1, T2, T3, T4]]: ...

    def product(self, *others: Iterable[Any]) -> Iter[tuple[Any, ...]]:
        """Computes the Cartesian product with another iterable.

        This is the declarative equivalent of nested for-loops.

        It pairs every element from the source iterable with every element from the
        other iterable.

        Args:
            *others (Iterable[Any]): Other iterables to compute the Cartesian product with.

        Returns:
            Iter[tuple[Any, ...]]: An iterable of tuples containing elements from the Cartesian product.

        Example:
        ```python
        >>> import pyochain as pc
        >>> sizes = ["S", "M"]
        >>> pc.Iter(["blue", "red"]).product(sizes).collect()
        Seq(('blue', 'S'), ('blue', 'M'), ('red', 'S'), ('red', 'M'))

        ```
        """
        return self._lazy(itertools.product, *others)

    def diff_at(
        self,
        *others: Iterable[T],
        default: T | None = None,
        key: Callable[[T], Any] | None = None,
    ) -> Iter[tuple[T, ...]]:
        """Return those items that differ between iterables.

        Each output item is a tuple where the i-th element is from the i-th input iterable.

        If an input iterable is exhausted before others, then the corresponding output items will be filled with *default*.

        Args:
            *others (Iterable[T]): Other iterables to compare with.
            default (T | None): Value to use for missing elements. Defaults to None.
            key (Callable[[T], Any] | None): Function to apply to each item for comparison. Defaults to None.

        Returns:
            Iter[tuple[T, ...]]: An iterable of tuples containing differing elements from the input iterables.

        Example:
        ```python
        >>> import pyochain as pc
        >>> data = pc.Seq([1, 2, 3])
        >>> data.iter().diff_at([1, 2, 10, 100], default=None).collect()
        Seq((3, 10), (None, 100))
        >>> data.iter().diff_at([1, 2, 10, 100, 2, 6, 7], default=0).collect()
        Seq((3, 10), (0, 100), (0, 2), (0, 6), (0, 7))

        A key function may also be applied to each item to use during comparisons:
        ```python
        >>> import pyochain as pc
        >>> pc.Iter(["apples", "bananas"]).diff_at(["Apples", "Oranges"], key=str.lower).collect()
        Seq(('bananas', 'Oranges'),)

        ```
        """
        return self._lazy(cz.itertoolz.diff, *others, default=default, key=key)

    def join_with[R, K](
        self,
        other: Iterable[R],
        left_on: Callable[[T], K],
        right_on: Callable[[R], K],
        left_default: T | None = None,
        right_default: R | None = None,
    ) -> Iter[tuple[T, R]]:
        """Perform a relational join with another iterable.

        Args:
            other (Iterable[R]): Iterable to join with.
            left_on (Callable[[T], K]): Function to extract the join key from the left iterable.
            right_on (Callable[[R], K]): Function to extract the join key from the right iterable.
            left_default (T | None): Default value for missing elements in the left iterable. Defaults to None.
            right_default (R | None): Default value for missing elements in the right iterable. Defaults to None.

        Returns:
            Iter[tuple[T, R]]: An iterator yielding tuples of joined elements.

        Example:
        ```python
        >>> import pyochain as pc
        >>> sizes = ["S", "M"]
        >>> pc.Iter(["blue", "red"]).join_with(sizes, left_on=lambda c: c, right_on=lambda s: s).collect()
        Seq((None, 'S'), (None, 'M'), ('blue', None), ('red', None))

        ```
        """

        def _join(data: Iterable[T]) -> Iterator[tuple[T, R]]:
            return cz.itertoolz.join(
                leftkey=left_on,
                leftseq=data,
                rightkey=right_on,
                rightseq=other,
                left_default=left_default,
                right_default=right_default,
            )

        return self._lazy(_join)
