"""
Unit tests for the retry module.

Tests for RetryPolicy, with_retry, and RetryableAgent.
"""

from typing import Any
from unittest.mock import AsyncMock

import pytest

from persuade.agents.base import AgentPhaseResult, TwoPhaseAgent
from persuade.agents.retry import (
    RETRY_NONE,
    RETRY_PERSISTENT,
    RETRY_QUICK,
    RETRY_STANDARD,
    RetryableAgent,
    RetryPolicy,
    calculate_delay,
    with_retry,
)


class TestRetryPolicy:
    """Tests for RetryPolicy dataclass."""

    def test_default_values(self) -> None:
        """Test default policy values."""
        policy = RetryPolicy()
        assert policy.max_attempts == 3
        assert policy.initial_delay_seconds == 1.0
        assert policy.backoff_multiplier == 2.0
        assert policy.max_delay_seconds == 60.0
        assert policy.jitter is True
        assert TimeoutError in policy.retryable_errors
        assert ConnectionError in policy.retryable_errors

    def test_custom_values(self) -> None:
        """Test policy with custom values."""
        policy = RetryPolicy(
            max_attempts=5,
            initial_delay_seconds=0.5,
            backoff_multiplier=3.0,
            max_delay_seconds=30.0,
            jitter=False,
            retryable_errors=(ValueError,),
        )
        assert policy.max_attempts == 5
        assert policy.initial_delay_seconds == 0.5
        assert policy.backoff_multiplier == 3.0
        assert policy.max_delay_seconds == 30.0
        assert policy.jitter is False
        assert policy.retryable_errors == (ValueError,)


class TestPresetPolicies:
    """Tests for preset retry policies."""

    def test_retry_none(self) -> None:
        """Test RETRY_NONE disables retries."""
        assert RETRY_NONE.max_attempts == 1

    def test_retry_quick(self) -> None:
        """Test RETRY_QUICK is fast."""
        assert RETRY_QUICK.max_attempts == 2
        assert RETRY_QUICK.initial_delay_seconds == 0.5

    def test_retry_standard(self) -> None:
        """Test RETRY_STANDARD is balanced."""
        assert RETRY_STANDARD.max_attempts == 3
        assert RETRY_STANDARD.initial_delay_seconds == 1.0

    def test_retry_persistent(self) -> None:
        """Test RETRY_PERSISTENT is thorough."""
        assert RETRY_PERSISTENT.max_attempts == 5
        assert RETRY_PERSISTENT.max_delay_seconds == 30.0


class TestCalculateDelay:
    """Tests for calculate_delay function."""

    def test_no_jitter_first_attempt(self) -> None:
        """Test delay without jitter on first attempt."""
        policy = RetryPolicy(
            initial_delay_seconds=1.0,
            backoff_multiplier=2.0,
            jitter=False,
        )
        delay = calculate_delay(policy, 0)
        assert delay == 1.0

    def test_no_jitter_exponential_backoff(self) -> None:
        """Test exponential backoff without jitter."""
        policy = RetryPolicy(
            initial_delay_seconds=1.0,
            backoff_multiplier=2.0,
            jitter=False,
        )
        assert calculate_delay(policy, 0) == 1.0
        assert calculate_delay(policy, 1) == 2.0
        assert calculate_delay(policy, 2) == 4.0
        assert calculate_delay(policy, 3) == 8.0

    def test_max_delay_cap(self) -> None:
        """Test that delay is capped at max_delay_seconds."""
        policy = RetryPolicy(
            initial_delay_seconds=1.0,
            backoff_multiplier=10.0,
            max_delay_seconds=5.0,
            jitter=False,
        )
        # 1 * 10^2 = 100, but should be capped at 5
        delay = calculate_delay(policy, 2)
        assert delay == 5.0

    def test_jitter_adds_randomness(self) -> None:
        """Test that jitter adds randomness to delay."""
        policy = RetryPolicy(
            initial_delay_seconds=10.0,
            jitter=True,
        )
        # Run multiple times to check randomness
        delays = [calculate_delay(policy, 0) for _ in range(10)]
        # Should have variation (not all the same)
        assert len(set(delays)) > 1
        # Should be in range 50-150% of base (5-15)
        for d in delays:
            assert 5.0 <= d <= 15.0


class TestWithRetry:
    """Tests for with_retry function."""

    @pytest.mark.asyncio
    async def test_success_first_attempt(self) -> None:
        """Test successful operation on first attempt."""
        mock_op = AsyncMock(return_value="success")

        result = await with_retry(mock_op, RETRY_STANDARD)

        assert result == "success"
        mock_op.assert_called_once()

    @pytest.mark.asyncio
    async def test_retry_on_retryable_error(self) -> None:
        """Test retry on retryable error."""
        mock_op = AsyncMock(side_effect=[TimeoutError("timeout"), "success"])
        policy = RetryPolicy(
            max_attempts=3,
            initial_delay_seconds=0.01,
            jitter=False,
        )

        result = await with_retry(mock_op, policy)

        assert result == "success"
        assert mock_op.call_count == 2

    @pytest.mark.asyncio
    async def test_no_retry_on_non_retryable_error(self) -> None:
        """Test no retry on non-retryable error."""
        mock_op = AsyncMock(side_effect=ValueError("invalid"))
        policy = RetryPolicy(
            max_attempts=3,
            retryable_errors=(TimeoutError,),  # ValueError not included
        )

        with pytest.raises(ValueError):
            await with_retry(mock_op, policy)

        mock_op.assert_called_once()

    @pytest.mark.asyncio
    async def test_exhausts_all_attempts(self) -> None:
        """Test that all attempts are exhausted before failing."""
        mock_op = AsyncMock(side_effect=TimeoutError("timeout"))
        policy = RetryPolicy(
            max_attempts=3,
            initial_delay_seconds=0.01,
            jitter=False,
        )

        with pytest.raises(TimeoutError):
            await with_retry(mock_op, policy)

        assert mock_op.call_count == 3

    @pytest.mark.asyncio
    async def test_retry_none_no_retry(self) -> None:
        """Test RETRY_NONE does not retry."""
        mock_op = AsyncMock(side_effect=TimeoutError("timeout"))

        with pytest.raises(TimeoutError):
            await with_retry(mock_op, RETRY_NONE)

        mock_op.assert_called_once()


class MockAgent(TwoPhaseAgent):
    """Mock agent for testing RetryableAgent."""

    def __init__(
        self,
        plan_results: list[AgentPhaseResult | Exception] | None = None,
        execute_results: list[AgentPhaseResult | Exception] | None = None,
    ):
        self.plan_results = plan_results or [AgentPhaseResult(success=True, data={"planned": True})]
        self.execute_results = execute_results or [
            AgentPhaseResult(success=True, data={"executed": True})
        ]
        self.plan_call_count = 0
        self.execute_call_count = 0

    async def plan(self, **kwargs: Any) -> AgentPhaseResult:
        result = self.plan_results[min(self.plan_call_count, len(self.plan_results) - 1)]
        self.plan_call_count += 1
        if isinstance(result, Exception):
            raise result
        return result

    async def execute(self, plan_result: AgentPhaseResult, **kwargs: Any) -> AgentPhaseResult:
        result = self.execute_results[min(self.execute_call_count, len(self.execute_results) - 1)]
        self.execute_call_count += 1
        if isinstance(result, Exception):
            raise result
        return result


class TestRetryableAgent:
    """Tests for RetryableAgent wrapper."""

    @pytest.mark.asyncio
    async def test_success_no_retry_needed(self) -> None:
        """Test successful execution without retry."""
        inner = MockAgent()
        agent = RetryableAgent(inner)

        plan_result = await agent.plan()
        assert plan_result.success is True
        assert inner.plan_call_count == 1

        exec_result = await agent.execute(plan_result)
        assert exec_result.success is True
        assert inner.execute_call_count == 1

    @pytest.mark.asyncio
    async def test_plan_retry_on_error(self) -> None:
        """Test plan phase retries on error."""
        inner = MockAgent(
            plan_results=[
                TimeoutError("timeout"),
                AgentPhaseResult(success=True),
            ]
        )
        agent = RetryableAgent(
            inner,
            plan_policy=RetryPolicy(
                max_attempts=3,
                initial_delay_seconds=0.01,
                jitter=False,
            ),
        )

        result = await agent.plan()
        assert result.success is True
        assert inner.plan_call_count == 2

    @pytest.mark.asyncio
    async def test_execute_retry_on_error(self) -> None:
        """Test execute phase retries on error."""
        inner = MockAgent(
            execute_results=[
                TimeoutError("timeout"),
                AgentPhaseResult(success=True),
            ]
        )
        agent = RetryableAgent(
            inner,
            execute_policy=RetryPolicy(
                max_attempts=3,
                initial_delay_seconds=0.01,
                jitter=False,
            ),
        )

        plan_result = await agent.plan()
        exec_result = await agent.execute(plan_result)

        assert exec_result.success is True
        assert inner.execute_call_count == 2

    @pytest.mark.asyncio
    async def test_plan_failure_after_exhaustion(self) -> None:
        """Test plan returns failure after exhausting retries."""
        inner = MockAgent(plan_results=[TimeoutError("always fails")])
        agent = RetryableAgent(
            inner,
            plan_policy=RetryPolicy(
                max_attempts=2,
                initial_delay_seconds=0.01,
                jitter=False,
            ),
        )

        result = await agent.plan()

        assert result.success is False
        assert result.error is not None
        assert "2 attempts" in result.error
        assert inner.plan_call_count == 2

    @pytest.mark.asyncio
    async def test_default_policies(self) -> None:
        """Test that default policies are applied."""
        inner = MockAgent()
        agent = RetryableAgent(inner)

        # Should use RETRY_STANDARD for plan (3 attempts)
        assert agent.plan_policy.max_attempts == 3
        # Should use RETRY_QUICK for execute (2 attempts)
        assert agent.execute_policy.max_attempts == 2

    @pytest.mark.asyncio
    async def test_agent_result_failure_not_retried(self) -> None:
        """Test that AgentPhaseResult failures are not retried."""
        inner = MockAgent(
            plan_results=[AgentPhaseResult(success=False, error="Business logic failure")]
        )
        agent = RetryableAgent(inner)

        result = await agent.plan()

        # Should not retry on logical failure
        assert result.success is False
        assert result.error == "Business logic failure"
        assert inner.plan_call_count == 1
