"""Lazy, immutable stream with chainable map/filter/tap operations."""

from __future__ import annotations

import asyncio
import itertools
import os
import sys
from abc import ABC, abstractmethod
from collections.abc import Iterable
from concurrent.futures import (
    ProcessPoolExecutor,
    ThreadPoolExecutor,
)
from copy import deepcopy
from enum import Enum, auto
from functools import reduce

if sys.version_info >= (3, 12):
    from itertools import batched
else:
    from itertools import islice

    def batched(iterable, n):  # type: ignore[no-redef]
        """Backport of itertools.batched for <3.12."""
        it = iter(iterable)
        while batch := tuple(islice(it, n)):
            yield batch


from typing import TYPE_CHECKING, Any

import attrs

if TYPE_CHECKING:
    from collections.abc import (
        Awaitable,
        Callable,
        Generator,
    )

    MapFn = Callable[[Any], Any]
    FilterFn = Callable[[Any], bool]
    TapFn = Callable[[Any], None]

    AsyncMapFn = Callable[[Any], Awaitable[Any]]
    AsyncFilterFn = Callable[[Any], Awaitable[bool]]
    AsyncTapFn = Callable[[Any], Awaitable[None]]

    StreamFn = MapFn | FilterFn | TapFn
    AsyncStreamFn = AsyncMapFn | AsyncFilterFn | AsyncTapFn

    PlannedOps = tuple[int, StreamFn]
    AsyncPlannedOps = tuple[int, AsyncStreamFn]


@attrs.frozen
class _BaseStream(ABC):
    """
    Abstract base for lazy stream operations.
    """

    seq: tuple[Any, ...] = attrs.field(
        validator=attrs.validators.instance_of(tuple),
    )
    ops: tuple[Any, ...] = attrs.field(
        default=(),
        validator=attrs.validators.instance_of(tuple),
        repr=False,
    )

    @classmethod
    @abstractmethod
    def from_iterable(
        cls,
        it: Iterable[Any],
    ) -> _BaseStream:
        """Create a stream from an iterable."""
        ...

    @abstractmethod
    def map(self, *fns: MapFn | AsyncMapFn) -> _BaseStream:
        """Return a new stream with mapping functions applied."""
        ...

    @abstractmethod
    def filter(self, *fns: FilterFn | AsyncFilterFn) -> _BaseStream:
        """Return a new stream with filter predicates applied."""
        ...

    @abstractmethod
    def tap(self, *fns: TapFn | AsyncTapFn) -> _BaseStream:
        """Return a new stream with side-effect functions applied."""
        ...

    @abstractmethod
    def partition(self, fn: FilterFn) -> tuple[_BaseStream, _BaseStream]:
        """Split the stream into two by a predicate."""
        ...

    @abstractmethod
    def fold(
        self,
        initial: Any,
        fn: Callable[[Any, Any], Any],
        *,
        workers: int = 1,
        use_threads: bool = False,
    ) -> Any:
        """Reduce the stream with a function and initial value."""
        ...

    @abstractmethod
    def collect(self) -> tuple[Any, ...]:
        """Materialize the stream into a tuple."""
        ...

    @abstractmethod
    def par_collect(
        self,
        workers: int = 4,
        *,
        use_threads: bool = False,
    ) -> tuple[Any, ...]:
        """Materialize the stream in parallel."""
        ...

    @abstractmethod
    async def async_collect(
        self,
    ) -> tuple[Any, ...]:
        """Materialize the stream asynchronously."""
        ...

    def __bool__(self) -> bool:
        """Return True if the stream has elements."""
        return bool(self.seq)


@attrs.frozen
class Stream(_BaseStream):
    """
    Lazy, immutable stream with chainable operations.
    """

    @classmethod
    def from_iterable(
        cls,
        it: Iterable[Any],
    ) -> Stream:
        """
        Create a stream from *it*.
        """
        if not isinstance(it, Iterable):
            it = [it]
        return cls(seq=tuple(it))

    def map(self, *fns: MapFn | AsyncMapFn) -> Stream:
        """
        Return a new stream with *fns* mapped.
        """
        new_ops = (
            *self.ops,
            *((_MAP, fn) for fn in fns),
        )
        return attrs.evolve(self, ops=new_ops)

    def filter(self, *fns: FilterFn | AsyncFilterFn) -> Stream:
        """
        Return a new stream filtering by *fns*.
        """
        new_ops = (
            *self.ops,
            *((_FILTER, fn) for fn in fns),
        )
        return attrs.evolve(self, ops=new_ops)

    def tap(self, *fns: TapFn | AsyncTapFn) -> Stream:
        """
        Return a new stream tapping *fns*.
        """
        new_ops = (
            *self.ops,
            *((_TAP, fn) for fn in fns),
        )
        return attrs.evolve(self, ops=new_ops)

    def partition(
        self,
        fn: FilterFn,
        *,
        workers: int = 1,
        use_threads: bool = False,
    ) -> tuple[Stream, Stream]:
        """
        Split into two streams by *fn*.
        """
        seq_tuple: tuple[Any, ...]
        if workers > 1:
            seq_tuple = self.par_collect(
                workers=workers,
                use_threads=use_threads,
            )
        else:
            seq_tuple = self.collect()
        return (
            Stream(
                seq=tuple(x for x in seq_tuple if fn(x)),
            ),
            Stream(
                seq=tuple(x for x in seq_tuple if not fn(x)),
            ),
        )

    def fold(
        self,
        initial: Any,
        fn: Callable[[Any, Any], Any],
        *,
        workers: int = 1,
        use_threads: bool = False,
    ) -> Any:
        """
        Reduce the stream with *fn* and *initial*.
        """
        if workers > 1:
            return reduce(
                fn,
                self.par_collect(
                    workers=workers,
                    use_threads=use_threads,
                ),
                initial,
            )
        return reduce(fn, self.collect(), initial)

    def collect(self) -> tuple[Any, ...]:
        """
        Materialize the stream into a tuple.
        """
        return tuple(_apply_fns(self.seq, self.ops))

    def par_collect(
        self,
        workers: int = 4,
        *,
        use_threads: bool = False,
    ) -> tuple[Any, ...]:
        """
        Materialize in parallel with *workers*.
        """
        if workers == -1:
            workers = (os.cpu_count() or 5) - 1

        executor_cls = (
            ThreadPoolExecutor if use_threads else ProcessPoolExecutor
        )
        batch_size = max(4, len(self.seq) // workers)
        batches = [
            (list(chunk), self.ops)
            for chunk in batched(self.seq, n=batch_size)
        ]

        with executor_cls(max_workers=workers) as ex:
            return tuple(
                itertools.chain.from_iterable(
                    ex.map(
                        _apply_fns_worker,
                        batches,
                    )
                )
            )

    async def async_collect(
        self,
    ) -> tuple[Any, ...]:
        """
        Materialize asynchronously.
        """
        if not self.ops:
            return self.collect()

        res = await asyncio.gather(
            *(_async_apply_fns(x, self.ops) for x in self.seq)
        )
        return tuple(elem for elem in res if elem != _Nothing.NOTHING)


_MAP = 0
_FILTER = 1
_TAP = 2


class _Nothing(Enum):
    """
    Sentinel for filtered-out elements.
    """

    NOTHING = auto()


def _apply_fns_worker(
    args: tuple[tuple[Any, ...], tuple[PlannedOps, ...]],
) -> tuple[Any, ...]:
    """
    Entry point for parallel batch processing.
    """
    seq, ops = args
    return _par_apply_fns(seq, ops)


def _apply_fns(
    elements: tuple[Any, ...],
    ops: tuple[PlannedOps, ...],
) -> Generator[Any, None, None]:
    """
    Apply *ops* to *elements* lazily.
    """
    for elem in elements:
        valid = True
        res: Any = elem
        for op, op_fn in ops:
            if op == _MAP:
                res = op_fn(res)
            elif op == _FILTER and not op_fn(res):
                valid = False
                break
            elif op == _TAP:
                op_fn(deepcopy(res))
        if valid:
            yield res


def _par_apply_fns(
    elements: tuple[Any, ...],
    ops: tuple[PlannedOps, ...],
) -> tuple[Any, ...]:
    """
    Apply *ops* to *elements* eagerly.
    """
    return tuple(_apply_fns(elements, ops))


async def _async_apply_fns(
    elem: Any,
    ops: tuple[AsyncPlannedOps, ...],
) -> Any:
    """
    Apply async *ops* to a single *elem*.
    """
    res: Any = elem
    for op, op_fn in ops:
        if op == _MAP:
            res = await op_fn(res)
        elif op == _FILTER and not await op_fn(res):
            return _Nothing.NOTHING
        elif op == _TAP:
            await op_fn(deepcopy(res))
    return res
