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

from .decorator import make_data_last

T = TypeVar('T')
NestedList: TypeAlias = Iterable[T | 'NestedList[T]']


@overload
def flat(it: NestedList[T], /) -> Iterable[T]: ...


@overload
def flat() -> Callable[[NestedList[T]], Iterable[T]]: ...


@overload
def flat(it: NestedList[T], /, *, depth: int) -> NestedList[T]: ...


@overload
def flat(*, depth: int) -> Callable[[NestedList[T]], NestedList[T]]: ...


def flat_infinite(it: NestedList[T], /) -> Iterable[T]:
    for i in it:
        if isinstance(i, Iterable):
            i = cast(NestedList[T], i)
            yield from flat_infinite(i)
        else:
            yield i


def flat_with_depth(it: NestedList[T], /, *, depth: int) -> Iterable[T] | NestedList[T]:
    for i in it:
        if isinstance(i, Iterable) and depth > 0:
            i = cast(NestedList[T], i)
            yield from flat_with_depth(i, depth=depth - 1)
        else:
            yield i


@make_data_last
def flat(iterable: NestedList[T], /, *, depth: int | None = None) -> Iterable[T] | NestedList[T]:
    """
    Yields elements for the iterable and sub-iterables up to the specified depth.

    Parameters
    ----------
    iterable: NestedList[T]
        The iterable to flatten (positional-only).
    depth: int | None
        The depth to flatten to. If None, flattens infinitely (keyword-only).

    Yields
    ------
    T
        The flattened elements.

    Examples
    --------
    Data first:
    >>> list(R.flat([[1, 2], [3, 4], [5], [[6]]]))
    [1, 2, 3, 4, 5, 6]
    >>> list(R.flat([[[1]], [[[2]]]], depth=2))
    [1, [2]]

    Data last:
    >>> R.pipe([[1, 2], [3, 4], [5], [[6]]], R.flat(), list)
    [1, 2, 3, 4, 5, 6]
    >>> R.pipe([[[1]], [[[2]]]], R.flat(depth=2), list)
    [1, [2]]

    """
    if depth is None:
        yield from flat_infinite(iterable)
    else:
        yield from flat_with_depth(iterable, depth=depth)
