import asyncio
from typing import Any, AsyncGenerator

import pytest

from r2r import AsyncPipe, AsyncPipeline, PipeType


class MultiplierPipe(AsyncPipe):
    def __init__(self, multiplier=1, delay=0, name="multiplier_pipe"):
        super().__init__(
            type=PipeType.OTHER,
            config=self.PipeConfig(name=name),
        )
        self.multiplier = multiplier
        self.delay = delay

    async def _run_logic(
        self,
        input: AsyncGenerator[Any, None],
        state,
        run_id=None,
        *args,
        **kwargs,
    ) -> AsyncGenerator[Any, None]:
        async for item in input.message:
            if self.delay > 0:
                await asyncio.sleep(self.delay)  # Simulate processing delay
            if isinstance(item, list):
                processed = [x * self.multiplier for x in item]
            elif isinstance(item, int):
                processed = item * self.multiplier
            else:
                raise ValueError(f"Unsupported type: {type(item)}")
            yield processed


class FanOutPipe(AsyncPipe):
    def __init__(self, multiplier=1, delay=0, name="fan_out_pipe"):
        super().__init__(
            type=PipeType.OTHER,
            config=self.PipeConfig(name=name),
        )
        self.multiplier = multiplier
        self.delay = delay

    async def _run_logic(
        self,
        input: AsyncGenerator[Any, None],
        state,
        run_id=None,
        *args,
        **kwargs,
    ) -> AsyncGenerator[Any, None]:
        inputs = []
        async for item in input.message:
            inputs.append(item)
        for it in range(self.multiplier):
            if self.delay > 0:
                await asyncio.sleep(self.delay)
            yield [(it + 1) * ele for ele in inputs]


class FanInPipe(AsyncPipe):
    def __init__(self, delay=0, name="fan_in_pipe"):
        super().__init__(
            type=PipeType.OTHER,
            config=self.PipeConfig(name=name),
        )
        self.delay = delay

    async def _run_logic(
        self,
        input: AsyncGenerator[Any, None],
        state,
        run_id=None,
        *args,
        **kwargs,
    ) -> AsyncGenerator[Any, None]:
        total_sum = 0
        async for batch in input.message:
            if self.delay > 0:
                await asyncio.sleep(self.delay)  # Simulate processing delay
            total_sum += sum(
                batch
            )  # Assuming batch is iterable and contains numeric values
        yield total_sum


@pytest.fixture
def pipe_factory():
    def create_pipe(type, **kwargs):
        if type == "multiplier":
            return MultiplierPipe(**kwargs)
        elif type == "fan_out":
            return FanOutPipe(**kwargs)
        elif type == "fan_in":
            return FanInPipe(**kwargs)
        else:
            raise ValueError("Unsupported pipe type")

    return create_pipe


@pytest.mark.asyncio
@pytest.mark.parametrize("multiplier, delay, name", [(2, 0.1, "pipe")])
async def test_single_multiplier(pipe_factory, multiplier, delay, name):
    pipe = pipe_factory(
        "multiplier", multiplier=multiplier, delay=delay, name=name
    )

    async def input_generator():
        for i in [1, 2, 3]:
            yield i

    pipeline = AsyncPipeline()
    pipeline.add_pipe(pipe)

    result = []
    for output in await pipeline.run(input_generator()):
        result.append(output)

    expected_result = [i * multiplier for i in [1, 2, 3]]
    assert (
        result == expected_result
    ), "Pipeline output did not match expected multipliers"


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "multiplier_a, delay_a, name_a, multiplier_b, delay_b, name_b",
    [(2, 0.1, "pipe_a", 2, 0.1, "pipe_b")],
)
async def test_double_multiplier(
    pipe_factory, multiplier_a, delay_a, name_a, multiplier_b, delay_b, name_b
):
    pipe_a = pipe_factory(
        "multiplier", multiplier=multiplier_a, delay=delay_a, name=name_a
    )
    pipe_b = pipe_factory(
        "multiplier", multiplier=multiplier_b, delay=delay_b, name=name_b
    )

    async def input_generator():
        for i in [1, 2, 3]:
            yield i

    pipeline = AsyncPipeline()
    pipeline.add_pipe(pipe_a)
    pipeline.add_pipe(pipe_b)

    result = []
    for output in await pipeline.run(input_generator()):
        result.append(output)

    expected_result = [i * multiplier_a * multiplier_b for i in [1, 2, 3]]
    assert (
        result == expected_result
    ), "Pipeline output did not match expected multipliers"


@pytest.mark.asyncio
@pytest.mark.parametrize("multiplier, delay, name", [(3, 0.1, "pipe")])
async def test_fan_out(pipe_factory, multiplier, delay, name):
    pipe = pipe_factory(
        "fan_out", multiplier=multiplier, delay=delay, name=name
    )

    async def input_generator():
        for i in [1, 2, 3]:
            yield i

    pipeline = AsyncPipeline()
    pipeline.add_pipe(pipe)

    result = []
    for output in await pipeline.run(input_generator()):
        result.append(output)

    expected_result = [
        [i + 1, 2 * (i + 1), 3 * (i + 1)] for i in range(multiplier)
    ]
    assert (
        result == expected_result
    ), "Pipeline output did not match expected multipliers"


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "multiplier_a, delay_a, name_a, multiplier_b, delay_b, name_b",
    [
        (2, 0.1, "pipe_a", 2, 0.1, "pipe_b"),
        (4, 0.1, "pipe_a", 3, 0.1, "pipe_b"),
    ],
)
async def multiply_then_fan_out(
    pipe_factory, multiplier_a, delay_a, name_a, multiplier_b, delay_b, name_b
):
    pipe_a = pipe_factory(
        "multiplier", multiplier=multiplier_a, delay=delay_a, name=name_a
    )
    pipe_b = pipe_factory(
        "fan_out", multiplier=multiplier_b, delay=delay_b, name=name_b
    )

    async def input_generator():
        for i in [1, 2, 3]:
            yield i

    pipeline = AsyncPipeline()
    pipeline.add_pipe(pipe_a)
    pipeline.add_pipe(pipe_b)

    result = []
    async for output in await pipeline.run(input_generator()):
        result.append(output)

    expected_result = [[i * multiplier_a] async for i in input_generator()]
    assert (
        result[0] == expected_result
    ), "Pipeline output did not match expected multipliers"


@pytest.mark.asyncio
@pytest.mark.parametrize("multiplier, delay, name", [(3, 0.1, "pipe")])
async def test_fan_in_sum(pipe_factory, multiplier, delay, name):
    # Create fan-out to generate multiple streams
    fan_out_pipe = pipe_factory(
        "fan_out", multiplier=multiplier, delay=delay, name=name + "_a"
    )
    # Summing fan-in pipe
    fan_in_sum_pipe = pipe_factory("fan_in", delay=delay, name=name + "_b")

    async def input_generator():
        for i in [1, 2, 3]:
            yield i

    pipeline = AsyncPipeline()
    pipeline.add_pipe(fan_out_pipe)
    pipeline.add_pipe(fan_in_sum_pipe)

    result = await pipeline.run(input_generator())

    # Calculate expected results based on the multiplier and the sum of inputs
    expected_result = sum(
        [sum([j * i for j in [1, 2, 3]]) for i in range(1, multiplier + 1)]
    )
    assert (
        result[0] == expected_result
    ), "Pipeline output did not match expected sums"


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "multiplier_a, delay_a, name_a, multiplier_b, delay_b, name_b",
    [
        (3, 0.1, "pipe_a", 2, 0.1, "pipe_b"),
        (4, 0.1, "pipe_a", 3, 0.1, "pipe_b"),
    ],
)
async def test_fan_out_then_multiply(
    pipe_factory, multiplier_a, delay_a, name_a, multiplier_b, delay_b, name_b
):
    pipe_a = pipe_factory(
        "multiplier", multiplier=multiplier_a, delay=delay_a, name=name_a
    )
    pipe_b = pipe_factory(
        "fan_out", multiplier=multiplier_b, delay=delay_b, name=name_b
    )
    pipe_c = pipe_factory("fan_in", delay=0.1, name="pipe_c")

    async def input_generator():
        for i in [1, 2, 3]:
            yield i

    pipeline = AsyncPipeline()
    pipeline.add_pipe(pipe_a)
    pipeline.add_pipe(pipe_b)
    pipeline.add_pipe(pipe_c)

    result = await pipeline.run(input_generator())

    expected_result = sum(
        [
            sum([j * i * multiplier_a for j in [1, 2, 3]])
            for i in range(1, multiplier_b + 1)
        ]
    )
    assert (
        result[0] == expected_result
    ), "Pipeline output did not match expected multipliers"