from collections.abc import Iterable, Sequence
from typing import Protocol, TypeVar, overload

from .decorator import make_data_last

T = TypeVar('T')


class ReturnedCallable(Protocol):
    @overload
    def __call__(self, it: str, /) -> str: ...
    @overload
    def __call__(self, it: Sequence[T], /) -> Iterable[T]: ...

    def __call__(self, it: Sequence[T], /) -> Iterable[T]: ...


@overload
def swap_indices(it: str, index1: int, index2: int, /) -> str: ...


@overload
def swap_indices(it: Sequence[T], index1: int, index2: int, /) -> Iterable[T]: ...


@overload
def swap_indices(index1: int, index2: int, /) -> ReturnedCallable: ...


def swap_indices_string(it: str, index1: int, index2: int, /) -> str:
    return ''.join(swap_indices_iterable(it, index1, index2))


def swap_indices_iterable(it: Sequence[T], index1: int, index2: int, /) -> Iterable[T]:
    index1_ = index1 % len(it)
    index2_ = index2 % len(it)
    lower_i = min(index1_, index2_)
    higher_i = max(index1_, index2_)
    iterable = iter(it)
    for _ in range(lower_i):
        yield next(iterable)
    lower_val = next(iterable)
    remembered: list[T] = []
    for _ in range(higher_i - lower_i - 1):
        remembered.append(next(iterable))
    yield next(iterable)
    yield from remembered
    yield lower_val
    yield from iterable


@make_data_last
def swap_indices(iterable: Sequence[T], index1: int, index2: int, /) -> Iterable[T] | str:
    """
    Yields elements of the given iterable swapping the 2 elements at the given indices.

    Parameters
    ----------
    iterable : Iterable[T]
        Iterable to yield elements from (positional-only).
    index1 : int
        Index of the first element to swap (positional-only).
    index2 : int
        Index of the second element to swap (positional-only).

    Returns
    -------
    Iterable[T]
        Iterable with the elements at the given indices swapped.

    Examples
    --------
    Data first:
    >>> list(R.swap_indices(['a', 'b', 'c'], 0, 1))
    ['b', 'a', 'c']
    >>> list(R.swap_indices(['a', 'b', 'c'], 1, -1))
    ['a', 'c', 'b']
    >>> R.swap_indices('abc', 0, 1)
    'bac'

    Data last:
    >>> list(R.swap_indices(0, 1)(['a', 'b', 'c']))
    ['b', 'a', 'c']
    >>> R.swap_indices(0, -1)('abc')
    'cba'

    """
    # TODO: support iterable instead of Sequence
    if isinstance(iterable, str):
        return swap_indices_string(iterable, index1, index2)
    return swap_indices_iterable(iterable, index1, index2)
