# Copyright (c) 2025, HaiyangLi <quantocean.li at gmail dot com>
# SPDX-License-Identifier: Apache-2.0

from builtins import BaseExceptionGroup, ExceptionGroup
from typing import Any
from unittest.mock import AsyncMock, patch

import anyio
import pytest
from pydantic import BaseModel

from lionpride.ln import alcall, bcall

# =============================================================================
# Test fixtures and helper functions
# =============================================================================


async def async_func(x: int, add: int = 0) -> int:
    """Simple async function for testing."""
    await anyio.sleep(0.01)
    return x + add


def sync_func(x: int, add: int = 0) -> int:
    """Simple sync function for testing."""
    return x + add


async def async_func_with_error(x: int) -> int:
    """Async function that raises error for specific input."""
    await anyio.sleep(0.01)
    if x == 3:
        raise ValueError("mock error")
    return x


def sync_func_with_error(x: int) -> int:
    """Sync function that raises error for specific input."""
    if x == 3:
        raise ValueError("mock error")
    return x


async def async_func_always_error(x: int) -> int:
    """Async function that always raises error."""
    await anyio.sleep(0.01)
    raise RuntimeError(f"Error for {x}")


class PydanticTestModel(BaseModel):
    """Pydantic model for testing MODEL_LIKE input handling."""

    value: int


# =============================================================================
# Test alcall function - Basic functionality
# =============================================================================


@pytest.mark.parametrize("anyio_backend", ["asyncio", "trio"])
class TestAlcallBasic:
    """Test alcall basic functionality."""

    @pytest.mark.anyio
    async def test_alcall_basic_async_function(self):
        """Test alcall with basic async function."""
        inputs = [1, 2, 3]
        results = await alcall(inputs, async_func, add=1)
        assert results == [2, 3, 4]

    @pytest.mark.anyio
    async def test_alcall_basic_sync_function(self):
        """Test alcall with sync function (no kwargs due to anyio limitation)."""
        inputs = [1, 2, 3]
        results = await alcall(inputs, sync_func)
        assert results == [1, 2, 3]

    @pytest.mark.anyio
    async def test_alcall_empty_input(self):
        """Test alcall with empty input."""
        results = await alcall([], async_func)
        assert results == []


# =============================================================================
# Test alcall function - func parameter validation (lines 71-82)
# =============================================================================


@pytest.mark.parametrize("anyio_backend", ["asyncio", "trio"])
class TestAlcallFuncValidation:
    """Test alcall func parameter validation."""

    @pytest.mark.anyio
    async def test_alcall_func_as_list_with_one_callable(self):
        """Test func parameter accepts list containing single callable."""
        inputs = [1, 2, 3]
        results = await alcall(inputs, [async_func])
        assert results == [1, 2, 3]

    @pytest.mark.anyio
    async def test_alcall_func_as_tuple_with_one_callable(self):
        """Test func parameter accepts tuple containing single callable."""
        inputs = [1, 2, 3]
        results = await alcall(inputs, (sync_func,))
        assert results == [1, 2, 3]

    @pytest.mark.anyio
    async def test_alcall_func_not_callable_not_iterable_raises(self):
        """Test non-callable, non-iterable func raises clear error."""
        with pytest.raises(ValueError, match="func must be callable"):
            await alcall([1, 2, 3], 123)

    @pytest.mark.anyio
    async def test_alcall_func_iterable_with_multiple_callables_raises(self):
        """Test multiple callables in iterable raises error."""
        with pytest.raises(ValueError, match="Only one callable"):
            await alcall([1, 2, 3], [async_func, sync_func])

    @pytest.mark.anyio
    async def test_alcall_func_iterable_with_non_callable_raises(self):
        """Test non-callable item in iterable raises error."""
        with pytest.raises(ValueError, match="Only one callable"):
            await alcall([1, 2, 3], ["not_callable"])

    @pytest.mark.anyio
    async def test_alcall_func_empty_iterable_raises(self):
        """Test empty iterable raises error."""
        with pytest.raises(ValueError, match="Only one callable"):
            await alcall([1, 2, 3], [])


# =============================================================================
# Test alcall function - Input processing
# =============================================================================


@pytest.mark.parametrize("anyio_backend", ["asyncio", "trio"])
class TestAlcallInputProcessing:
    """Test alcall input processing."""

    @pytest.mark.anyio
    async def test_alcall_input_flatten(self):
        """Test input_flatten parameter flattens nested iterables."""
        inputs = [[1, 2], [3, 4]]
        results = await alcall(inputs, async_func, input_flatten=True)
        assert results == [1, 2, 3, 4]

    @pytest.mark.anyio
    async def test_alcall_input_dropna(self):
        """Test input_dropna parameter filters None values."""
        inputs = [1, None, 2, None, 3]
        results = await alcall(inputs, async_func, input_dropna=True)
        assert results == [1, 2, 3]

    @pytest.mark.anyio
    async def test_alcall_input_pydantic_model(self):
        """Test Pydantic model input normalization."""
        model = PydanticTestModel(value=5)
        results = await alcall(model, lambda x: x.value * 2)
        assert results == [10]

    @pytest.mark.anyio
    async def test_alcall_input_tuple(self):
        """Test tuple input converts to list."""
        inputs = (1, 2, 3)
        results = await alcall(inputs, async_func)
        assert results == [1, 2, 3]

    @pytest.mark.anyio
    async def test_alcall_input_generator(self):
        """Test generator input converts to list."""
        inputs = (x for x in [1, 2, 3])
        results = await alcall(inputs, async_func)
        assert results == [1, 2, 3]

    @pytest.mark.anyio
    async def test_alcall_input_range(self):
        """Test range input converts to list."""
        inputs = range(3)
        results = await alcall(inputs, async_func)
        assert results == [0, 1, 2]

    @pytest.mark.anyio
    async def test_alcall_input_non_iterable(self):
        """Test non-iterable input wraps in list."""
        result = await alcall(5, async_func)
        assert result == [5]


# =============================================================================
# Test alcall function - Retry and timeout
# =============================================================================


@pytest.mark.parametrize("anyio_backend", ["asyncio", "trio"])
class TestAlcallRetryTimeout:
    """Test alcall retry and timeout functionality."""

    @pytest.mark.anyio
    async def test_alcall_with_retries_async_func(self):
        """Test retry mechanism with async function."""
        inputs = [1, 2, 3]
        results = await alcall(
            inputs,
            async_func_with_error,
            retry_attempts=1,
            retry_default=0,
        )
        assert results == [1, 2, 0]

    @pytest.mark.anyio
    async def test_alcall_with_retries_sync_func(self):
        """Test retry mechanism with sync function."""
        inputs = [1, 2, 3]
        results = await alcall(
            inputs,
            sync_func_with_error,
            retry_attempts=1,
            retry_default=0,
        )
        assert results == [1, 2, 0]

    @pytest.mark.anyio
    async def test_alcall_timeout_async_function(self):
        """Test timeout enforcement for async functions."""

        async def slow_async_func(x: int) -> int:
            await anyio.sleep(1.0)
            return x

        inputs = [1, 2, 3]
        results = await alcall(
            inputs,
            slow_async_func,
            retry_timeout=0.05,
            retry_default="timeout",
            retry_attempts=0,
        )
        assert results == ["timeout", "timeout", "timeout"]

    @pytest.mark.anyio
    async def test_alcall_timeout_sync_function(self):
        """Test timeout enforcement for sync functions in thread pool."""

        def slow_sync_func(x: int) -> int:
            import time

            time.sleep(0.5)  # Sleep longer than timeout
            return x

        inputs = [1]  # Single input for faster test
        results = await alcall(
            inputs,
            slow_sync_func,
            retry_timeout=0.1,
            retry_default="timeout",
            retry_attempts=0,
        )
        # Note: timeout might not work reliably with sync functions in threads
        assert len(results) == 1

    @pytest.mark.anyio
    async def test_alcall_retry_backoff(self):
        """Test retry with backoff."""
        with patch("anyio.sleep", new_callable=AsyncMock) as mock_sleep:
            inputs = [3]  # Only one item that triggers error
            await alcall(
                inputs,
                async_func_with_error,
                retry_attempts=2,
                retry_initial_delay=0.1,
                retry_backoff=2,
                retry_default=0,
            )
            # Should call sleep with 0.1, then 0.2
            assert mock_sleep.call_count >= 2


# =============================================================================
# Test alcall function - Exception handling
# =============================================================================


@pytest.mark.parametrize("anyio_backend", ["asyncio", "trio"])
class TestAlcallExceptionHandling:
    """Test alcall exception handling."""

    @pytest.mark.anyio
    async def test_alcall_exception_reraises_after_retry_exhaustion(self):
        """Test exception re-raises when retries exhausted and no default."""
        inputs = [1, 2, 3]
        # Exceptions in task groups are wrapped in ExceptionGroup
        try:
            await alcall(
                inputs,
                async_func_always_error,
                retry_attempts=2,
                # No retry_default, should re-raise
            )
            raise AssertionError("Should have raised exception")
        except BaseExceptionGroup as eg:
            # Verify all sub-exceptions are RuntimeError
            for exc in eg.exceptions:
                assert isinstance(exc, RuntimeError)

    @pytest.mark.anyio
    async def test_alcall_exception_with_retry_default_no_reraise(self):
        """Test exception with retry_default does not re-raise."""
        inputs = [1, 2, 3]
        results = await alcall(
            inputs,
            async_func_always_error,
            retry_attempts=2,
            retry_default="failed",
        )
        assert results == ["failed", "failed", "failed"]


# =============================================================================
# Test alcall function - Concurrency and throttling
# =============================================================================


@pytest.mark.parametrize("anyio_backend", ["asyncio", "trio"])
class TestAlcallConcurrency:
    """Test alcall concurrency and throttling."""

    @pytest.mark.anyio
    async def test_alcall_max_concurrent(self):
        """Test max_concurrent parameter."""
        inputs = [1, 2, 3, 4, 5]
        results = await alcall(inputs, async_func, max_concurrent=2)
        assert results == [1, 2, 3, 4, 5]

    @pytest.mark.anyio
    async def test_alcall_throttle_period(self):
        """Test throttle_period parameter."""
        inputs = [1, 2, 3]
        results = await alcall(inputs, async_func, throttle_period=0.01)
        assert results == [1, 2, 3]

    @pytest.mark.anyio
    async def test_alcall_delay_before_start(self):
        """Test delay_before_start parameter."""
        with patch("anyio.sleep", new_callable=AsyncMock) as mock_sleep:
            inputs = [1, 2, 3]
            await alcall(inputs, async_func, delay_before_start=0.5)
            mock_sleep.assert_any_call(0.5)


# =============================================================================
# Test alcall function - Output processing
# =============================================================================


@pytest.mark.parametrize("anyio_backend", ["asyncio", "trio"])
class TestAlcallOutputProcessing:
    """Test alcall output processing."""

    @pytest.mark.anyio
    async def test_alcall_output_flatten(self):
        """Test output_flatten parameter."""

        async def func_returning_list(x: int) -> list:
            return [x, x * 2]

        inputs = [1, 2, 3]
        results = await alcall(inputs, func_returning_list, output_flatten=True)
        assert results == [1, 2, 2, 4, 3, 6]

    @pytest.mark.anyio
    async def test_alcall_output_dropna(self):
        """Test output_dropna parameter."""

        async def func_with_none(x: int) -> Any:
            return None if x == 2 else x

        inputs = [1, 2, 3]
        results = await alcall(inputs, func_with_none, output_dropna=True)
        assert results == [1, 3]

    @pytest.mark.anyio
    async def test_alcall_output_unique(self):
        """Test output_unique parameter."""

        async def func_with_duplicates(x: int) -> list:
            return [x, x]

        inputs = [1, 2, 3]
        results = await alcall(
            inputs,
            func_with_duplicates,
            output_flatten=True,
            output_unique=True,
        )
        assert sorted(results) == [1, 2, 3]


# =============================================================================
# Test bcall function
# =============================================================================


@pytest.mark.parametrize("anyio_backend", ["asyncio", "trio"])
class TestBcall:
    """Test bcall function."""

    @pytest.mark.anyio
    async def test_bcall_basic(self):
        """Test bcall basic functionality."""
        inputs = [1, 2, 3, 4, 5]
        batches = []
        async for batch in bcall(inputs, async_func, batch_size=2):
            batches.append(batch)
        assert batches == [[1, 2], [3, 4], [5]]

    @pytest.mark.anyio
    async def test_bcall_with_retries(self):
        """Test bcall with retries."""
        inputs = [1, 2, 3, 4, 5]
        batches = []
        async for batch in bcall(
            inputs,
            async_func_with_error,
            batch_size=2,
            retry_attempts=1,
            retry_default=0,
        ):
            batches.append(batch)
        assert batches == [[1, 2], [0, 4], [5]]

    @pytest.mark.anyio
    async def test_bcall_with_kwargs(self):
        """Test bcall with kwargs."""
        inputs = [1, 2, 3, 4, 5]
        batches = []
        async for batch in bcall(inputs, async_func, batch_size=2, add=10):
            batches.append(batch)
        assert batches == [[11, 12], [13, 14], [15]]

    @pytest.mark.anyio
    async def test_bcall_with_all_options(self):
        """Test bcall with all options."""
        inputs = [1, 2, 3, 4, 5]
        batches = []
        async for batch in bcall(
            inputs,
            async_func,
            batch_size=2,
            input_flatten=False,
            output_flatten=False,
            max_concurrent=2,
            throttle_period=0.01,
        ):
            batches.append(batch)
        assert batches == [[1, 2], [3, 4], [5]]


# =============================================================================
# Test edge cases and combinations
# =============================================================================


@pytest.mark.parametrize("anyio_backend", ["asyncio", "trio"])
class TestEdgeCases:
    """Test edge cases and combinations."""

    @pytest.mark.anyio
    async def test_alcall_combined_input_output_processing(self):
        """Test combined input and output processing."""

        async def func_returning_list(x: int) -> list:
            return [x, x * 2]

        inputs = [[1, 2], [3, 4]]
        results = await alcall(
            inputs,
            func_returning_list,
            input_flatten=True,
            output_flatten=True,
        )
        assert results == [1, 2, 2, 4, 3, 6, 4, 8]

    @pytest.mark.anyio
    async def test_alcall_with_both_flatten_and_unique(self):
        """Test combined flatten and unique."""

        async def func_with_duplicates(x: int) -> list:
            return [x, x, x + 1]

        inputs = [1, 2, 3]
        results = await alcall(
            inputs,
            func_with_duplicates,
            output_flatten=True,
            output_unique=True,
        )
        assert sorted(results) == [1, 2, 3, 4]

    @pytest.mark.anyio
    async def test_alcall_max_concurrent_with_throttle(self):
        """Test max_concurrent with throttle_period."""
        inputs = [1, 2, 3, 4, 5]
        results = await alcall(
            inputs,
            async_func,
            max_concurrent=2,
            throttle_period=0.01,
        )
        assert results == [1, 2, 3, 4, 5]


# =============================================================================
# Return Exceptions Tests
# =============================================================================


class TestReturnExceptions:
    """Tests for return_exceptions parameter in alcall and bcall."""

    @pytest.mark.anyio
    async def test_alcall_return_exceptions_true_returns_exceptions(self):
        """Test that return_exceptions=True returns exceptions in results list."""
        inputs = [1, 2, 3, 4, 5]
        results = await alcall(inputs, async_func_with_error, return_exceptions=True)

        # Result at index 2 should be ValueError, others should be successful
        assert results[0] == 1
        assert results[1] == 2
        assert isinstance(results[2], ValueError)
        assert str(results[2]) == "mock error"
        assert results[3] == 4
        assert results[4] == 5

    @pytest.mark.anyio
    async def test_alcall_return_exceptions_false_raises_exceptiongroup(self):
        """Test that return_exceptions=False (default) raises ExceptionGroup."""
        inputs = [1, 2, 3, 4, 5]

        with pytest.raises(ExceptionGroup) as exc_info:
            await alcall(inputs, async_func_with_error, return_exceptions=False)

        # Verify ExceptionGroup contains ValueError
        assert any(isinstance(e, ValueError) for e in exc_info.value.exceptions)

    @pytest.mark.anyio
    async def test_alcall_return_exceptions_default_raises(self):
        """Test that default behavior (no parameter) raises ExceptionGroup."""
        inputs = [1, 2, 3, 4, 5]

        with pytest.raises(ExceptionGroup) as exc_info:
            await alcall(inputs, async_func_with_error)

        assert any(isinstance(e, ValueError) for e in exc_info.value.exceptions)

    @pytest.mark.anyio
    async def test_alcall_return_exceptions_true_all_failures(self):
        """Test return_exceptions=True when all tasks fail."""

        async def always_fail(x: int) -> int:
            raise RuntimeError(f"fail_{x}")

        inputs = [1, 2, 3]
        results = await alcall(inputs, always_fail, return_exceptions=True)

        # All results should be exceptions
        assert all(isinstance(r, RuntimeError) for r in results)
        assert str(results[0]) == "fail_1"
        assert str(results[1]) == "fail_2"
        assert str(results[2]) == "fail_3"

    @pytest.mark.anyio
    async def test_alcall_return_exceptions_true_all_success(self):
        """Test return_exceptions=True when all tasks succeed."""
        inputs = [1, 2, 3, 4, 5]
        results = await alcall(inputs, async_func, return_exceptions=True)

        # No exceptions, all successful
        assert results == [1, 2, 3, 4, 5]
        assert not any(isinstance(r, BaseException) for r in results)

    @pytest.mark.anyio
    async def test_alcall_return_exceptions_preserves_order(self):
        """Test that return_exceptions preserves original input order."""
        inputs = [5, 4, 3, 2, 1]
        results = await alcall(inputs, async_func_with_error, return_exceptions=True)

        # Order should match input order
        assert results[0] == 5
        assert results[1] == 4
        assert isinstance(results[2], ValueError)  # x=3 fails
        assert results[3] == 2
        assert results[4] == 1

    @pytest.mark.anyio
    async def test_alcall_return_exceptions_with_sync_function(self):
        """Test return_exceptions=True with synchronous function."""
        inputs = [1, 2, 3, 4, 5]
        results = await alcall(inputs, sync_func_with_error, return_exceptions=True)

        assert results[0] == 1
        assert results[1] == 2
        assert isinstance(results[2], ValueError)
        assert results[3] == 4
        assert results[4] == 5

    @pytest.mark.anyio
    async def test_bcall_return_exceptions_true(self):
        """Test bcall with return_exceptions=True."""
        inputs = [1, 2, 3, 4, 5, 6]
        batches = []

        async for batch in bcall(
            inputs,
            async_func_with_error,
            batch_size=3,
            return_exceptions=True,
        ):
            batches.append(batch)

        # First batch: [1, 2, 3] -> [1, 2, ValueError]
        assert len(batches) == 2
        assert batches[0][0] == 1
        assert batches[0][1] == 2
        assert isinstance(batches[0][2], ValueError)

        # Second batch: [4, 5, 6] -> [4, 5, 6]
        assert batches[1] == [4, 5, 6]

    @pytest.mark.anyio
    async def test_bcall_return_exceptions_false_raises(self):
        """Test bcall with return_exceptions=False raises ExceptionGroup."""
        inputs = [1, 2, 3, 4, 5]

        with pytest.raises(ExceptionGroup):
            async for _batch in bcall(
                inputs,
                async_func_with_error,
                batch_size=3,
                return_exceptions=False,
            ):
                pass  # Should raise before completing iteration

    @pytest.mark.anyio
    async def test_alcall_return_exceptions_with_retries(self):
        """Test return_exceptions=True with retry logic."""
        call_counts = {}

        async def flaky_func(x: int) -> int:
            call_counts[x] = call_counts.get(x, 0) + 1
            if x == 3 and call_counts[x] < 3:
                raise ValueError("temporary error")
            return x

        inputs = [1, 2, 3]
        results = await alcall(
            inputs,
            flaky_func,
            retry_attempts=2,
            return_exceptions=True,
        )

        # x=3 should succeed after retries
        assert results == [1, 2, 3]
        assert call_counts[3] == 3  # Called 3 times (initial + 2 retries)

    @pytest.mark.anyio
    async def test_alcall_return_exceptions_with_retry_exhaustion(self):
        """Test return_exceptions=True when retries are exhausted."""
        call_counts = {}

        async def always_fail(x: int) -> int:
            call_counts[x] = call_counts.get(x, 0) + 1
            raise ValueError(f"persistent error {x}")

        inputs = [1, 2]
        results = await alcall(
            inputs,
            always_fail,
            retry_attempts=2,
            return_exceptions=True,
        )

        # Should return exceptions after exhausting retries
        assert all(isinstance(r, ValueError) for r in results)
        assert call_counts[1] == 3  # initial + 2 retries
        assert call_counts[2] == 3

    @pytest.mark.anyio
    async def test_alcall_return_exceptions_with_max_concurrent(self):
        """Test return_exceptions=True with concurrency limit."""
        inputs = [1, 2, 3, 4, 5]
        results = await alcall(
            inputs,
            async_func_with_error,
            max_concurrent=2,
            return_exceptions=True,
        )

        assert results[0] == 1
        assert results[1] == 2
        assert isinstance(results[2], ValueError)
        assert results[3] == 4
        assert results[4] == 5


# =============================================================================
# Test timeout path without triggering timeout
# =============================================================================


class TestTimeoutWithoutTriggering:
    """Test timeout set but function completes before timeout."""

    @pytest.mark.anyio
    async def test_alcall_async_with_timeout_but_completes_fast(self):
        """Test async function completes within timeout."""

        async def fast_async_func(x: int) -> int:
            # Very fast function - completes well before timeout
            await anyio.sleep(0.001)
            return x * 2

        inputs = [1, 2, 3]
        # Set timeout but function completes before it
        results = await alcall(inputs, fast_async_func, retry_timeout=1.0)
        assert results == [2, 4, 6]

    @pytest.mark.anyio
    async def test_alcall_sync_with_timeout_but_completes_fast(self):
        """Test sync function completes within timeout."""
        import time

        def fast_sync_func(x: int) -> int:
            # Small delay to ensure we enter the timeout scope but don't trigger it
            time.sleep(0.001)
            return x * 3

        inputs = [1, 2, 3]
        # Set timeout but function completes before it
        results = await alcall(inputs, fast_sync_func, retry_timeout=0.5)
        assert results == [3, 6, 9]

    @pytest.mark.anyio
    async def test_alcall_sync_timeout_raises_error(self):
        """Test sync function timeout behavior."""
        import time

        def slow_sync_func(x: int) -> int:
            # Sleep with periodic checks to allow cancellation
            for _ in range(100):
                time.sleep(0.01)
            return x

        # Test with return_exceptions to catch the TimeoutError
        results = await alcall(
            [1], slow_sync_func, retry_timeout=0.001, return_exceptions=True, retry_attempts=0
        )
        # With a very short timeout, should get TimeoutError
        assert len(results) == 1
        # Note: Sync function cancellation may not work reliably in all backends
        assert isinstance(results[0], (TimeoutError, int))


# =============================================================================
# Test cancellation exception handling
# =============================================================================


class TestCancellationHandling:
    """Test cancellation exception re-raise."""

    @pytest.mark.anyio
    async def test_alcall_cancellation_exception_reraises(self):
        """Test cancellation propagates correctly."""

        async def cancellable_func(x: int) -> int:
            # Long-running function that can be cancelled
            await anyio.sleep(1.0)
            return x

        # Test cancellation by using a cancel scope
        with anyio.CancelScope() as scope:
            async with anyio.create_task_group() as tg:

                async def run_alcall():
                    await alcall([1, 2, 3, 4, 5], cancellable_func)

                tg.start_soon(run_alcall)
                # Cancel immediately
                await anyio.sleep(0.001)
                scope.cancel()

        # If we reach here, cancellation was handled correctly
        assert scope.cancel_called

    @pytest.mark.anyio
    async def test_alcall_pure_cancellation_exceptiongroup(self):
        """Test ExceptionGroup containing only cancellation exceptions."""

        async def cancellable_func(x: int) -> int:
            # Function that will be cancelled
            await anyio.sleep(10.0)  # Long sleep ensures cancellation
            return x

        # Create a cancel scope that cancels during execution
        try:
            with anyio.fail_after(0.01):  # Very short timeout to force cancellation
                await alcall([1, 2, 3], cancellable_func)
        except anyio.get_cancelled_exc_class():
            # Expected - pure cancellation scenario
            pass
        except TimeoutError:
            # Alternative - timeout is also acceptable
            pass
