Vec

class Vec[T](Seq[T], PyoMutableSequence[T]):
    """A `MutableSequence` wrapper with functional API.

    Implement `MutableSequence` Protocol from `collections.abc`.

    Unlike `Seq` which is immutable, `Vec` allows in-place modification of elements.

    Implement the `MutableSequence` interface, so elements can be modified in place, and passed to any function/object expecting a standard mutable sequence.

    Args:
        data (Iterable[T]): The `Iterable` to wrap.
    """

    __slots__ = ()
    _inner: list[T]  # type: ignore[override]

    def __init__(self, data: Iterable[T]) -> None:
        self._inner = list(data)  # type: ignore[override]

    @staticmethod
    def from_ref[V](data: list[V]) -> Vec[V]:
        """Create a `Vec` from a reference to an existing `list`.

        This method wraps the provided `list` without copying it, allowing for efficient creation of a `Vec`.

        This is the recommended way to create a `Vec` from foreign functions.

        Warning:
            Since the `Vec` directly references the original `list`, any modifications made to the `Vec` will also affect the original `list`, and vice versa.

        Args:
            data (list[V]): The `list` to wrap.

        Returns:
            Vec[V]: A new Vec instance wrapping the provided `list`.

        Example:
        ```python
        >>> import pyochain as pc
        >>> original_list = [1, 2, 3]
        >>> vec = pc.Vec.from_ref(original_list)
        >>> vec
        Vec(1, 2, 3)
        >>> vec[0] = 10
        >>> original_list
        [10, 2, 3]

        ```
        """
        instance: Vec[V] = Vec.__new__(Vec)  # pyright: ignore[reportUnknownVariableType]
        instance._inner = data
        return instance

    @overload
    def __setitem__(self, index: int, value: T) -> None: ...
    @overload
    def __setitem__(self, index: slice, value: Iterable[T]) -> None: ...
    def __setitem__(self, index: int | slice, value: T | Iterable[T]) -> None:
        return self._inner.__setitem__(index, value)  # type: ignore[arg-type]

    def __delitem__(self, index: int | slice) -> None:
        del self._inner[index]

    def insert(self, index: int, value: T) -> None:
        """Inserts an element at position index within the vector, shifting all elements after it to the right.

        Args:
            index (int): Position where to insert the element.
            value (T): The element to insert.

        Example:
        ```python
        >>> import pyochain as pc
        >>> vec = pc.Vec(['a', 'b', 'c'])
        >>> vec.insert(1, 'd')
        >>> vec
        Vec('a', 'd', 'b', 'c')
        >>> vec.insert(4, 'e')
        >>> vec
        Vec('a', 'd', 'b', 'c', 'e')

        ```
        """
        self._inner.insert(index, value)

    @overload
    def sort[U: SupportsRichComparison[Any]](
        self: Vec[U], *, key: None = None, reverse: bool = False
    ) -> Vec[U]: ...
    @overload
    def sort(
        self, *, key: Callable[[T], SupportsRichComparison[Any]], reverse: bool = False
    ) -> Vec[T]: ...
    @overload
    def sort(
        self,
        *,
        key: None = None,
        reverse: bool = False,
    ) -> Never: ...
    def sort(
        self,
        *,
        key: Callable[[T], SupportsRichComparison[Any]] | None = None,
        reverse: bool = False,
    ) -> Vec[Any]:
        """Sort the elements of the `Vec` in place.

        Warning:
            This method modifies the `Vec` in place and returns the same instance for chaining.

        Args:
            key (Callable[[T], SupportsRichComparison[Any]] | None): Optional function to extract a comparison key from each element.
            reverse (bool): If True, sort in descending order.

        Returns:
            Vec[Any]: The sorted `Vec` instance (self).

        Example:
        ```python
        >>> import pyochain as pc
        >>> pc.Vec((3, 1, 2)).sort()
        Vec(1, 2, 3)
        >>> pc.Iter((3, 1, 2)).map(str).collect(pc.Vec).sort(key=int)
        Vec('1', '2', '3')

        ```
        """
        self._inner.sort(key=key, reverse=reverse)  # type: ignore[arg-type]
        return self

    def extract_if(
        self, predicate: Callable[[T], bool], start: int = 0, end: int | None = None
    ) -> Iter[T]:
        """Creates an `Iter` which uses a *predicate* to determine if an element in the `Vec` should be removed.

        If the *predicate* returns `True`, the element is removed from the `Vec` and yielded.

        If the *predicate* returns `False`, the element remains in the `Vec` and will not be yielded.

        You can specify a range for the extraction.

        If the returned ExtractIf is not exhausted, e.g. because it is dropped without iterating or the iteration short-circuits, then the remaining elements will be retained.

        Use retain_mut with a negated predicate if you do not need the returned iterator.

        Using this method is equivalent to the following code:
        ```python
            data = pc.Vec([ ... ])
            for i in range(data.length()):
                if predicate(data[i]):
                    val = data.pop(i)
                    # your code here
        ```
        """

        def _extract_if_gen() -> Iterator[T]:
            effective_end = end if end is not None else len(self)
            i = start
            pop = self.pop
            while i < effective_end and i < len(self):
                if predicate(self[i]):
                    yield pop(i)
                    effective_end -= 1
                else:
                    i += 1

        return Iter(_extract_if_gen())

    def drain(self, start: int | None = None, end: int | None = None) -> Iter[T]:
        """Removes the subslice indicated by the given *start* and *end* from the `Vec`, returning an `Iterator` over the removed subslice.

        If the `Iterator` is dropped before being fully consumed, it drops the remaining removed elements.

        Note:
            In CPython, remaining elements are cleaned up when the `Iterator` is garbage collected via `__del__`.
            However, in interactive environments like doctests, garbage collection may not happen immediately.
            To guarantee cleanup, fully consume the `Iterator` or explicitly call `.collect()` on it.

        Args:
            start (int | None): Starting index of the subslice to drain. Defaults to `0` if `None`.
            end (int | None): Ending index of the subslice to drain. Defaults to `len(self)` if `None`.

        Returns:
            Iter[T]: An `Iterator` over the drained elements.

        Examples:
        ```python
        >>> import pyochain as pc
        >>> v = pc.Vec([1, 2, 3])
        >>> u = v.drain(1).collect();
        >>> v
        Vec(1)
        >>> u
        Seq(2, 3)
        >>> # A full range clears the vector, like `clear()` does
        >>> _ = v.drain().collect();
        >>> v
        Vec()

        ```
        Fully consuming the `Iterator` removes all drained elements
        ```python
        >>> import pyochain as pc
        >>> v = pc.Vec([1, 2, 3])
        >>> _ = v.drain(0, 3).collect()
        >>> v
        Vec()

        ```
        """
        return Iter(
            DrainIterator(self, start if start else 0, end if end else len(self))
        )

    def concat(self, other: list[T] | Self) -> Vec[T]:
        """Concatenate another `Vec` or `list` to **self** and return a new `Vec`.

        Note:
            This is equivalent to `list_1 + list_2` for standard lists.

        Args:
            other (list[T] | Self): The other `Vec` to concatenate.

        Returns:
            Vec[T]: The new `Vec` after concatenation.

        See Also:
            `Vec.extend()` which modifies **self** in place.

        Example:
        ```python
        >>> import pyochain as pc
        >>> v1 = pc.Vec([1, 2, 3])
        >>> v2 = [4, 5, 6] # Can also concatenate a standard list
        >>> v3 = v1.concat(v2)
        >>> v3
        Vec(1, 2, 3, 4, 5, 6)
        >>> v1.clear() # Clean up the original vec
        >>> v1
        Vec()
        >>> # New vec remains unaffected
        >>> v3
        Vec(1, 2, 3, 4, 5, 6)

        ```
        """
        match other:
            case Vec():
                data = self._inner + other._inner
            case list():
                data = self._inner + other
        return Vec.from_ref(data)