"""Tests for retry mechanism."""

import pytest

from swarm.retry import ERROR_TYPE_TO_RETRY_ON, ToolRetryExecutor


class MockRetryPolicy:
    """Mock retry policy that satisfies RetryPolicyProtocol."""

    def __init__(
        self,
        max_retries: int = 3,
        backoff: str = "exponential",
        base_seconds: int = 1,
        jitter: bool = False,
        retry_on: list = None,
    ):
        self._max_retries = max_retries
        self._backoff = backoff
        self._base_seconds = base_seconds
        self._jitter = jitter
        self._retry_on = retry_on or ["timeout", "transient_error"]

    @property
    def max_retries(self) -> int:
        return self._max_retries

    @property
    def backoff(self) -> str:
        return self._backoff

    @property
    def base_seconds(self) -> int:
        return self._base_seconds

    @property
    def jitter(self) -> bool:
        return self._jitter

    @property
    def retry_on(self) -> list:
        return self._retry_on


class MockResult:
    """Mock result object for testing."""

    def __init__(self, value: str, is_error: bool = False, error_type=None):
        self.value = value
        self.is_error = is_error
        self.error_type = error_type


class MockErrorType:
    """Mock error type enum."""

    def __init__(self, value: str):
        self.value = value


class TestToolRetryExecutor:
    """Test ToolRetryExecutor."""

    def test_init_with_policy(self):
        """Test initialization with policy."""
        policy = MockRetryPolicy(max_retries=5)
        executor = ToolRetryExecutor(policy)
        assert executor.policy.max_retries == 5

    def test_is_retryable_transient_error(self):
        """Test that transient errors are retryable."""
        policy = MockRetryPolicy(retry_on=["transient_error"])
        executor = ToolRetryExecutor(policy)

        # "transient" maps to "transient_error" in ERROR_TYPE_TO_RETRY_ON
        assert executor.is_retryable("transient") is True

    def test_is_retryable_non_retryable_error(self):
        """Test that non-retryable errors are not retried."""
        policy = MockRetryPolicy(retry_on=["transient_error"])
        executor = ToolRetryExecutor(policy)

        # "validation" maps to "validation_error" which is not in retry_on
        assert executor.is_retryable("validation") is False

    def test_is_retryable_none_error_type(self):
        """Test that None error type is not retryable."""
        policy = MockRetryPolicy()
        executor = ToolRetryExecutor(policy)
        assert executor.is_retryable(None) is False

    def test_calculate_backoff_constant(self):
        """Test constant backoff calculation."""
        policy = MockRetryPolicy(backoff="constant", base_seconds=5, jitter=False)
        executor = ToolRetryExecutor(policy)

        assert executor.calculate_backoff(1) == 5.0
        assert executor.calculate_backoff(2) == 5.0
        assert executor.calculate_backoff(3) == 5.0

    def test_calculate_backoff_linear(self):
        """Test linear backoff calculation."""
        policy = MockRetryPolicy(backoff="linear", base_seconds=2, jitter=False)
        executor = ToolRetryExecutor(policy)

        assert executor.calculate_backoff(1) == 2.0
        assert executor.calculate_backoff(2) == 4.0
        assert executor.calculate_backoff(3) == 6.0

    def test_calculate_backoff_exponential(self):
        """Test exponential backoff calculation."""
        policy = MockRetryPolicy(backoff="exponential", base_seconds=1, jitter=False)
        executor = ToolRetryExecutor(policy)

        assert executor.calculate_backoff(1) == 1.0  # 1 * 2^0
        assert executor.calculate_backoff(2) == 2.0  # 1 * 2^1
        assert executor.calculate_backoff(3) == 4.0  # 1 * 2^2

    def test_calculate_backoff_max_cap(self):
        """Test that backoff is capped at 300 seconds."""
        policy = MockRetryPolicy(backoff="exponential", base_seconds=100, jitter=False)
        executor = ToolRetryExecutor(policy)

        # 100 * 2^3 = 800, should be capped at 300
        assert executor.calculate_backoff(4) == 300.0

    @pytest.mark.asyncio
    async def test_execute_with_retry_success_first_try(self):
        """Test successful execution on first try."""
        policy = MockRetryPolicy()
        executor = ToolRetryExecutor(policy)

        async def success_func():
            return MockResult(value="success", is_error=False)

        result = await executor.execute_with_retry(success_func, tool_name="test_tool")
        assert result.value == "success"
        assert result.is_error is False

    @pytest.mark.asyncio
    async def test_execute_with_retry_retries_on_transient_error(self):
        """Test that transient errors trigger retries."""
        policy = MockRetryPolicy(
            max_retries=2,
            backoff="constant",
            base_seconds=0,  # No delay for testing
            retry_on=["transient_error"],
        )
        executor = ToolRetryExecutor(policy)

        call_count = 0

        async def failing_then_success_func():
            nonlocal call_count
            call_count += 1
            if call_count < 2:
                return MockResult(
                    value="error",
                    is_error=True,
                    error_type=MockErrorType("transient"),
                )
            return MockResult(value="success", is_error=False)

        result = await executor.execute_with_retry(
            failing_then_success_func, tool_name="test_tool"
        )
        assert result.value == "success"
        assert call_count == 2

    @pytest.mark.asyncio
    async def test_execute_with_retry_stops_on_non_retryable_error(self):
        """Test that non-retryable errors don't trigger retries."""
        policy = MockRetryPolicy(
            max_retries=3,
            retry_on=["transient_error"],
        )
        executor = ToolRetryExecutor(policy)

        call_count = 0

        async def validation_error_func():
            nonlocal call_count
            call_count += 1
            return MockResult(
                value="validation failed",
                is_error=True,
                error_type=MockErrorType("validation"),  # Not in retry_on
            )

        result = await executor.execute_with_retry(
            validation_error_func, tool_name="test_tool"
        )
        assert result.is_error is True
        assert call_count == 1  # No retries

    @pytest.mark.asyncio
    async def test_execute_with_retry_max_retries_exhausted(self):
        """Test that retries stop after max_retries."""
        policy = MockRetryPolicy(
            max_retries=2,
            backoff="constant",
            base_seconds=0,
            retry_on=["transient_error"],
        )
        executor = ToolRetryExecutor(policy)

        call_count = 0

        async def always_failing_func():
            nonlocal call_count
            call_count += 1
            return MockResult(
                value="error",
                is_error=True,
                error_type=MockErrorType("transient"),
            )

        result = await executor.execute_with_retry(
            always_failing_func, tool_name="test_tool"
        )
        assert result.is_error is True
        # Initial call + max_retries = 1 + 2 = 3
        assert call_count == 3


class TestErrorTypeMapping:
    """Test ERROR_TYPE_TO_RETRY_ON mapping."""

    def test_mapping_completeness(self):
        """Test that common error types are mapped."""
        assert "internal" in ERROR_TYPE_TO_RETRY_ON
        assert "transient" in ERROR_TYPE_TO_RETRY_ON
        assert "validation" in ERROR_TYPE_TO_RETRY_ON

    def test_mapping_values(self):
        """Test mapping values are correct."""
        assert ERROR_TYPE_TO_RETRY_ON["internal"] == "internal_error"
        assert ERROR_TYPE_TO_RETRY_ON["transient"] == "transient_error"
        assert ERROR_TYPE_TO_RETRY_ON["validation"] == "validation_error"
