from __future__ import annotations

import asyncio

import attrs
import pytest

from codeflash_core.danom.stream import (
    _FILTER,
    _MAP,
    _TAP,
    Stream,
    _apply_fns,
    _Nothing,
    _par_apply_fns,
)


class TestStreamFromIterable:
    """Tests for Stream.from_iterable."""

    def test_from_list(self):
        """
        A list is converted into a Stream with a tuple seq.
        """
        s = Stream.from_iterable([1, 2, 3])

        assert (1, 2, 3) == s.seq

    def test_from_tuple(self):
        """
        A tuple is accepted as an iterable.
        """
        s = Stream.from_iterable((4, 5))

        assert (4, 5) == s.seq

    def test_from_generator(self):
        """
        A generator is consumed and stored as a tuple.
        """
        s = Stream.from_iterable(x * 2 for x in range(3))

        assert (0, 2, 4) == s.seq

    def test_from_empty(self):
        """
        An empty iterable produces an empty stream.
        """
        s = Stream.from_iterable([])

        assert () == s.seq

    def test_from_string(self):
        """
        A string iterable yields individual characters.
        """
        s = Stream.from_iterable("abc")

        assert ("a", "b", "c") == s.seq

    def test_from_range(self):
        """
        A range object is accepted as an iterable.
        """
        s = Stream.from_iterable(range(4))

        assert (0, 1, 2, 3) == s.seq

    def test_ops_default_empty(self):
        """
        A fresh stream has no queued operations.
        """
        s = Stream.from_iterable([1])

        assert () == s.ops


class TestStreamBool:
    """Tests for Stream.__bool__."""

    def test_truthy_when_nonempty(self):
        """
        A stream with elements is truthy.
        """
        s = Stream.from_iterable([1])

        assert bool(s) is True

    def test_falsy_when_empty(self):
        """
        An empty stream is falsy.
        """
        s = Stream.from_iterable([])

        assert bool(s) is False


class TestStreamMap:
    """Tests for Stream.map."""

    def test_single_map(self):
        """
        A single map function is applied to each element.
        """
        result = (
            Stream.from_iterable([1, 2, 3]).map(lambda x: x * 10).collect()
        )

        assert (10, 20, 30) == result

    def test_multiple_fns_in_one_call(self):
        """
        Multiple functions passed to a single map call are applied sequentially.
        """
        result = (
            Stream.from_iterable([1, 2])
            .map(lambda x: x + 1, lambda x: x * 3)
            .collect()
        )

        assert (6, 9) == result

    def test_map_empty_stream(self):
        """
        Mapping over an empty stream produces an empty tuple.
        """
        result = Stream.from_iterable([]).map(lambda x: x + 1).collect()

        assert () == result

    def test_map_preserves_immutability(self):
        """
        Calling map returns a new stream without mutating the original.
        """
        s1 = Stream.from_iterable([1, 2])
        s2 = s1.map(lambda x: x * 2)

        assert () == s1.ops
        assert 1 == len(s2.ops)
        assert s1 is not s2


class TestStreamFilter:
    """Tests for Stream.filter."""

    def test_single_filter(self):
        """
        A single filter predicate removes non-matching elements.
        """
        result = (
            Stream.from_iterable([1, 2, 3, 4])
            .filter(lambda x: x % 2 == 0)
            .collect()
        )

        assert (2, 4) == result

    def test_filter_removes_all(self):
        """
        A predicate that matches nothing produces an empty tuple.
        """
        result = (
            Stream.from_iterable([1, 3, 5]).filter(lambda x: x > 100).collect()
        )

        assert () == result

    def test_filter_keeps_all(self):
        """
        A predicate that matches everything keeps all elements.
        """
        result = (
            Stream.from_iterable([2, 4, 6])
            .filter(lambda x: x % 2 == 0)
            .collect()
        )

        assert (2, 4, 6) == result

    def test_multiple_filters(self):
        """
        Multiple filter predicates in one call are applied sequentially.
        """
        result = (
            Stream.from_iterable(range(20))
            .filter(lambda x: x % 2 == 0, lambda x: x > 5)
            .collect()
        )

        assert (6, 8, 10, 12, 14, 16, 18) == result

    def test_filter_empty_stream(self):
        """
        Filtering an empty stream produces an empty tuple.
        """
        result = Stream.from_iterable([]).filter(lambda x: True).collect()

        assert () == result


class TestStreamTap:
    """Tests for Stream.tap."""

    def test_tap_side_effect(self):
        """
        Tap executes a side effect without altering elements.
        """
        seen = []
        result = Stream.from_iterable([1, 2, 3]).tap(seen.append).collect()

        assert (1, 2, 3) == result
        assert [1, 2, 3] == seen

    def test_tap_receives_deepcopy(self):
        """
        Tap receives a deep copy so mutations do not affect the stream.
        """
        originals = [{"a": 1}, {"a": 2}]
        captured = []

        def mutating_tap(x):
            x["a"] = 999
            captured.append(x)

        result = Stream.from_iterable(originals).tap(mutating_tap).collect()

        assert ({"a": 1}, {"a": 2}) == result
        assert [{"a": 999}, {"a": 999}] == captured


class TestStreamChaining:
    """Tests for chaining multiple operations."""

    def test_map_then_filter(self):
        """
        Map followed by filter applies both in order.
        """
        result = (
            Stream.from_iterable([1, 2, 3, 4, 5])
            .map(lambda x: x * 2)
            .filter(lambda x: x > 4)
            .collect()
        )

        assert (6, 8, 10) == result

    def test_filter_then_map(self):
        """
        Filter followed by map applies both in order.
        """
        result = (
            Stream.from_iterable([1, 2, 3, 4, 5])
            .filter(lambda x: x > 2)
            .map(lambda x: x**2)
            .collect()
        )

        assert (9, 16, 25) == result

    def test_map_filter_tap_chain(self):
        """
        All three operation types can be chained together.
        """
        tapped = []
        result = (
            Stream.from_iterable([1, 2, 3, 4])
            .map(lambda x: x + 10)
            .filter(lambda x: x % 2 == 0)
            .tap(tapped.append)
            .collect()
        )

        assert (12, 14) == result
        assert [12, 14] == tapped

    def test_multiple_map_calls(self):
        """
        Multiple chained map calls compose sequentially.
        """
        result = (
            Stream.from_iterable([1, 2])
            .map(lambda x: x + 1)
            .map(lambda x: x * 10)
            .collect()
        )

        assert (20, 30) == result


class TestStreamFold:
    """Tests for Stream.fold."""

    def test_sum(self):
        """
        Fold with addition accumulates a sum.
        """
        result = Stream.from_iterable([1, 2, 3, 4]).fold(
            0, lambda acc, x: acc + x
        )

        assert 10 == result

    def test_fold_with_operations(self):
        """
        Fold applies queued operations before reducing.
        """
        result = (
            Stream.from_iterable([1, 2, 3, 4])
            .filter(lambda x: x % 2 == 0)
            .map(lambda x: x * 10)
            .fold(0, lambda acc, x: acc + x)
        )

        assert 60 == result

    def test_fold_empty_stream(self):
        """
        Folding an empty stream returns the initial value.
        """
        result = Stream.from_iterable([]).fold(42, lambda acc, x: acc + x)

        assert 42 == result

    def test_fold_string_concat(self):
        """
        Fold can concatenate strings.
        """
        result = Stream.from_iterable(["a", "b", "c"]).fold(
            "", lambda acc, x: acc + x
        )

        assert "abc" == result


class TestStreamPartition:
    """Tests for Stream.partition."""

    def test_basic_partition(self):
        """
        Partition splits elements by the predicate into two streams.
        """
        truthy, falsy = Stream.from_iterable([1, 2, 3, 4, 5]).partition(
            lambda x: x > 3
        )

        assert (4, 5) == truthy.collect()
        assert (1, 2, 3) == falsy.collect()

    def test_partition_all_true(self):
        """
        When all elements match, the false stream is empty.
        """
        truthy, falsy = Stream.from_iterable([2, 4, 6]).partition(
            lambda x: x % 2 == 0
        )

        assert (2, 4, 6) == truthy.collect()
        assert () == falsy.collect()

    def test_partition_all_false(self):
        """
        When no elements match, the true stream is empty.
        """
        truthy, falsy = Stream.from_iterable([1, 3, 5]).partition(
            lambda x: x % 2 == 0
        )

        assert () == truthy.collect()
        assert (1, 3, 5) == falsy.collect()

    def test_partition_empty(self):
        """
        Partitioning an empty stream yields two empty streams.
        """
        truthy, falsy = Stream.from_iterable([]).partition(lambda x: x > 0)

        assert () == truthy.collect()
        assert () == falsy.collect()

    def test_partition_applies_pending_ops(self):
        """
        Partition materializes pending operations before splitting.
        """
        truthy, falsy = (
            Stream.from_iterable([1, 2, 3, 4])
            .map(lambda x: x * 10)
            .partition(lambda x: x >= 30)
        )

        assert (30, 40) == truthy.collect()
        assert (10, 20) == falsy.collect()


class TestStreamCollect:
    """Tests for Stream.collect."""

    def test_collect_no_ops(self):
        """
        Collecting with no operations returns the original elements.
        """
        result = Stream.from_iterable([1, 2, 3]).collect()

        assert (1, 2, 3) == result

    def test_collect_returns_tuple(self):
        """
        Collect always returns a tuple.
        """
        result = Stream.from_iterable([1]).collect()

        assert isinstance(result, tuple)

    def test_collect_single_element(self):
        """
        A single-element stream collects to a one-element tuple.
        """
        result = Stream.from_iterable([42]).map(lambda x: x + 1).collect()

        assert (43,) == result


class TestStreamParCollect:
    """Tests for Stream.par_collect."""

    def test_par_collect_threads(self):
        """
        Parallel collection with threads produces the same result as sequential.
        """
        result = (
            Stream.from_iterable(range(20))
            .map(lambda x: x * 2)
            .par_collect(workers=2, use_threads=True)
        )

        assert tuple(x * 2 for x in range(20)) == result

    def test_par_collect_empty(self):
        """
        Parallel collection of an empty stream returns an empty tuple.
        """
        result = Stream.from_iterable([]).par_collect(
            workers=2, use_threads=True
        )

        assert () == result

    def test_par_collect_with_filter(self):
        """
        Parallel collection respects filter operations.
        """
        result = (
            Stream.from_iterable(range(10))
            .filter(lambda x: x % 2 == 0)
            .par_collect(workers=2, use_threads=True)
        )

        assert (0, 2, 4, 6, 8) == result


class TestStreamAsyncCollect:
    """Tests for Stream.async_collect."""

    @pytest.mark.asyncio
    async def test_async_map(self):
        """
        Async map functions are awaited during collection.
        """

        async def times_ten(x):
            return x * 10

        result = await (
            Stream.from_iterable([1, 2, 3]).map(times_ten).async_collect()
        )

        assert (10, 20, 30) == result

    @pytest.mark.asyncio
    async def test_async_collect_no_ops(self):
        """
        Async collect with no ops falls back to synchronous collect.
        """
        result = await Stream.from_iterable([1, 2]).async_collect()

        assert (1, 2) == result

    @pytest.mark.asyncio
    async def test_async_filter(self):
        """
        Async filter predicates remove non-matching elements.
        """

        async def gt_two(x):
            return x > 2

        result = await (
            Stream.from_iterable([1, 2, 3, 4]).filter(gt_two).async_collect()
        )

        assert (3, 4) == result

    @pytest.mark.asyncio
    async def test_async_tap(self):
        """
        Async tap executes side effects without altering elements.
        """
        seen = []

        async def record(x):
            seen.append(x)

        result = await (
            Stream.from_iterable([10, 20]).tap(record).async_collect()
        )

        assert (10, 20) == result
        assert [10, 20] == seen


class TestApplyFns:
    """Tests for the _apply_fns helper."""

    def test_map_op(self):
        """
        MAP operations transform elements.
        """
        ops = ((_MAP, lambda x: x + 1),)
        result = tuple(_apply_fns((1, 2, 3), ops))

        assert (2, 3, 4) == result

    def test_filter_op(self):
        """
        FILTER operations remove non-matching elements.
        """
        ops = ((_FILTER, lambda x: x > 2),)
        result = tuple(_apply_fns((1, 2, 3, 4), ops))

        assert (3, 4) == result

    def test_tap_op(self):
        """
        TAP operations run side effects and pass elements through.
        """
        seen = []
        ops = ((_TAP, seen.append),)
        result = tuple(_apply_fns((1, 2), ops))

        assert (1, 2) == result
        assert 2 == len(seen)

    def test_combined_ops(self):
        """
        Mixed MAP, FILTER, and TAP ops are applied in sequence.
        """
        tapped = []
        ops = (
            (_MAP, lambda x: x * 2),
            (_FILTER, lambda x: x > 4),
            (_TAP, tapped.append),
        )
        result = tuple(_apply_fns((1, 2, 3, 4), ops))

        assert (6, 8) == result
        assert [6, 8] == tapped

    def test_empty_ops(self):
        """
        No operations yields elements unchanged.
        """
        result = tuple(_apply_fns((1, 2), ()))

        assert (1, 2) == result

    def test_empty_elements(self):
        """
        Empty elements yields nothing regardless of operations.
        """
        ops = ((_MAP, lambda x: x + 1),)
        result = tuple(_apply_fns((), ops))

        assert () == result

    def test_filter_short_circuits(self):
        """
        Once a filter rejects an element, subsequent ops are skipped.
        """
        map_calls = []

        def tracking_map(x):
            map_calls.append(x)
            return x

        ops = (
            (_FILTER, lambda x: x > 5),
            (_MAP, tracking_map),
        )
        tuple(_apply_fns((1, 10), ops))

        assert [10] == map_calls


class TestParApplyFns:
    """Tests for the _par_apply_fns helper."""

    def test_returns_tuple(self):
        """
        _par_apply_fns returns a tuple, not a generator.
        """
        result = _par_apply_fns((1, 2, 3), ((_MAP, lambda x: x + 1),))

        assert isinstance(result, tuple)
        assert (2, 3, 4) == result

    def test_empty_elements(self):
        """
        Empty elements produce an empty tuple.
        """
        result = _par_apply_fns((), ((_MAP, lambda x: x),))

        assert () == result

    def test_filter_and_map(self):
        """
        Combined filter and map work in the eager variant.
        """
        ops = (
            (_FILTER, lambda x: x % 2 == 0),
            (_MAP, lambda x: x * 10),
        )
        result = _par_apply_fns((1, 2, 3, 4), ops)

        assert (20, 40) == result


class TestNothing:
    """Tests for the _Nothing sentinel."""

    def test_nothing_is_singleton(self):
        """
        NOTHING is the sole member of the _Nothing enum.
        """
        assert 1 == len(_Nothing)
        assert _Nothing.NOTHING is _Nothing.NOTHING

    def test_nothing_is_not_equal_to_none(self):
        """
        NOTHING is distinct from None.
        """
        assert _Nothing.NOTHING != None  # noqa: E711


class TestStreamImmutability:
    """Tests for Stream's frozen/immutable behavior."""

    def test_cannot_set_seq(self):
        """
        Assigning to seq on a frozen stream raises an error.
        """
        s = Stream.from_iterable([1, 2])

        with pytest.raises(attrs.exceptions.FrozenInstanceError):
            s.seq = (3, 4)

    def test_cannot_set_ops(self):
        """
        Assigning to ops on a frozen stream raises an error.
        """
        s = Stream.from_iterable([1, 2])

        with pytest.raises(attrs.exceptions.FrozenInstanceError):
            s.ops = ()

    def test_seq_must_be_tuple(self):
        """
        Constructing a stream with a non-tuple seq raises TypeError.
        """
        with pytest.raises(TypeError):
            Stream(seq=[1, 2, 3])


class TestEdgeCases:
    """Tests for edge cases and boundary conditions."""

    def test_single_element_map(self):
        """
        A stream with one element processes map correctly.
        """
        result = Stream.from_iterable([5]).map(lambda x: x * 3).collect()

        assert (15,) == result

    def test_single_element_filter_keeps(self):
        """
        A single element matching the filter is kept.
        """
        result = Stream.from_iterable([5]).filter(lambda x: x > 0).collect()

        assert (5,) == result

    def test_single_element_filter_removes(self):
        """
        A single element not matching the filter is removed.
        """
        result = Stream.from_iterable([5]).filter(lambda x: x > 10).collect()

        assert () == result

    def test_none_values_in_stream(self):
        """
        None values are valid stream elements.
        """
        result = Stream.from_iterable([None, 1, None]).collect()

        assert (None, 1, None) == result

    def test_nested_streams(self):
        """
        Streams can contain other streams as elements.
        """
        inner = Stream.from_iterable([1, 2])
        outer = Stream.from_iterable([inner, inner])
        result = outer.collect()

        assert 2 == len(result)
        assert (1, 2) == result[0].collect()

    def test_large_stream(self):
        """
        A large stream processes without error.
        """
        result = (
            Stream.from_iterable(range(10000))
            .map(lambda x: x + 1)
            .filter(lambda x: x % 1000 == 0)
            .collect()
        )

        assert (
            1000,
            2000,
            3000,
            4000,
            5000,
            6000,
            7000,
            8000,
            9000,
            10000,
        ) == result

    def test_fold_with_workers_threads(self):
        """
        Fold with workers>1 and use_threads produces the same result.
        """
        result = Stream.from_iterable(range(100)).fold(
            0, lambda acc, x: acc + x, workers=2, use_threads=True
        )

        assert 4950 == result
