"""
Tests for the Retry module.
"""

import pytest
import time
from unittest.mock import Mock, patch

from bgx.retry import (
    RetryManager, BackoffCalculator,
    fixed_retry_policy, exponential_retry_policy,
    linear_retry_policy, fibonacci_retry_policy
)
from bgx.task import Task, TaskState, RetryPolicy


class TestBackoffCalculator:
    """Test BackoffCalculator utility."""
    
    def test_fixed_strategy(self):
        """Test fixed backoff strategy."""
        delay = BackoffCalculator.compute_delay(
            strategy="fixed",
            attempt=1,
            base_delay=2.0,
            jitter=False
        )
        assert delay == 2.0
        
        # Should be same regardless of attempt
        delay = BackoffCalculator.compute_delay(
            strategy="fixed",
            attempt=5,
            base_delay=2.0,
            jitter=False
        )
        assert delay == 2.0
    
    def test_linear_strategy(self):
        """Test linear backoff strategy."""
        delay = BackoffCalculator.compute_delay(
            strategy="linear",
            attempt=1,
            base_delay=2.0,
            jitter=False
        )
        assert delay == 2.0  # 2.0 * 1
        
        delay = BackoffCalculator.compute_delay(
            strategy="linear",
            attempt=3,
            base_delay=2.0,
            jitter=False
        )
        assert delay == 6.0  # 2.0 * 3
    
    def test_exponential_strategy(self):
        """Test exponential backoff strategy."""
        delay = BackoffCalculator.compute_delay(
            strategy="exponential",
            attempt=1,
            base_delay=2.0,
            jitter=False
        )
        assert delay == 2.0  # 2.0 * 2^(1-1)
        
        delay = BackoffCalculator.compute_delay(
            strategy="exponential",
            attempt=3,
            base_delay=2.0,
            jitter=False
        )
        assert delay == 8.0  # 2.0 * 2^(3-1)
    
    def test_fibonacci_strategy(self):
        """Test Fibonacci backoff strategy."""
        delay = BackoffCalculator.compute_delay(
            strategy="fibonacci",
            attempt=1,
            base_delay=1.0,
            jitter=False
        )
        assert delay == 1.0  # 1.0 * fib(1)
        
        delay = BackoffCalculator.compute_delay(
            strategy="fibonacci",
            attempt=2,
            base_delay=1.0,
            jitter=False
        )
        assert delay == 1.0  # 1.0 * fib(2)
        
        delay = BackoffCalculator.compute_delay(
            strategy="fibonacci",
            attempt=3,
            base_delay=1.0,
            jitter=False
        )
        assert delay == 2.0  # 1.0 * fib(3)
        
        delay = BackoffCalculator.compute_delay(
            strategy="fibonacci",
            attempt=4,
            base_delay=1.0,
            jitter=False
        )
        assert delay == 3.0  # 1.0 * fib(4)
        
        delay = BackoffCalculator.compute_delay(
            strategy="fibonacci",
            attempt=5,
            base_delay=1.0,
            jitter=False
        )
        assert delay == 5.0  # 1.0 * fib(5)
    
    def test_max_delay_cap(self):
        """Test maximum delay cap."""
        delay = BackoffCalculator.compute_delay(
            strategy="exponential",
            attempt=10,
            base_delay=10.0,
            max_delay=50.0,
            jitter=False
        )
        assert delay == 50.0  # Capped at max_delay
    
    def test_jitter(self):
        """Test jitter addition."""
        # Without jitter
        delay1 = BackoffCalculator.compute_delay(
            strategy="fixed",
            attempt=1,
            base_delay=10.0,
            jitter=False
        )
        assert delay1 == 10.0
        
        # With jitter - should be within range
        delay2 = BackoffCalculator.compute_delay(
            strategy="fixed",
            attempt=1,
            base_delay=10.0,
            jitter=True
        )
        assert 10.0 <= delay2 <= 12.5  # 10.0 + 25% jitter
    
    def test_unknown_strategy(self):
        """Test unknown strategy raises error."""
        with pytest.raises(ValueError, match="Unknown backoff strategy"):
            BackoffCalculator.compute_delay(
                strategy="unknown",
                attempt=1,
                base_delay=1.0
            )


class TestRetryManager:
    """Test RetryManager functionality."""
    
    def test_retry_manager_creation(self):
        """Test retry manager creation."""
        manager = RetryManager()
        assert manager.get_retry_count() == 0
        assert manager.get_retry_info() == {}
    
    def test_set_submit_callback(self):
        """Test setting submit callback."""
        manager = RetryManager()
        callback = Mock()
        manager.set_submit_callback(callback)
        # Callback is stored internally, we test it through schedule_retry
    
    def test_schedule_retry_success(self):
        """Test successful retry scheduling."""
        manager = RetryManager()
        callback = Mock()
        manager.set_submit_callback(callback)
        
        def sample_func():
            pass
        
        task = Task(
            func=sample_func,
            retry_policy=RetryPolicy(max_attempts=3)
        )
        task.error = ValueError("test error")
        task.transition_to(TaskState.QUEUED)
        task.transition_to(TaskState.RUNNING)
        task.transition_to(TaskState.FAILED)
        
        # Schedule retry
        result = manager.schedule_retry(task)
        
        assert result is True
        assert task.state == TaskState.RETRYING
        assert task.attempt == 2
        assert manager.get_retry_count() == 1
        
        # Check retry info
        retry_info = manager.get_retry_info()
        assert task.id in retry_info
        assert retry_info[task.id] > time.time() - 1.0  # Scheduled in the future
    
    def _create_failed_task(self, func=None, retry_policy=None):
        """Helper to create a failed task with proper state transitions."""
        if func is None:
            def sample_func():
                pass
            func = sample_func
        
        if retry_policy is None:
            retry_policy = RetryPolicy(max_attempts=3)
        
        task = Task(func=func, retry_policy=retry_policy)
        task.error = ValueError("test error")
        task.transition_to(TaskState.QUEUED)
        task.transition_to(TaskState.RUNNING)
        task.transition_to(TaskState.FAILED)
        return task
    
    def test_schedule_retry_no_policy(self):
        """Test retry scheduling without retry policy."""
        manager = RetryManager()
        
        def sample_func():
            pass
        
        task = Task(func=sample_func)
        task.error = ValueError("test error")
        task.transition_to(TaskState.QUEUED)
        task.transition_to(TaskState.RUNNING)
        task.transition_to(TaskState.FAILED)
        
        # Should not schedule retry without policy
        result = manager.schedule_retry(task)
        assert result is False
        assert manager.get_retry_count() == 0
    
    def test_schedule_retry_max_attempts(self):
        """Test retry scheduling with max attempts reached."""
        manager = RetryManager()
        
        task = self._create_failed_task(retry_policy=RetryPolicy(max_attempts=2))
        task.attempt = 2  # Already at max attempts
        
        # Should not schedule retry
        result = manager.schedule_retry(task)
        assert result is False
        assert manager.get_retry_count() == 0
    
    def test_schedule_retry_no_error(self):
        """Test retry scheduling without error."""
        manager = RetryManager()
        
        task = self._create_failed_task()
        task.error = None  # No error set
        
        # Should not schedule retry without error
        result = manager.schedule_retry(task)
        assert result is False
        assert manager.get_retry_count() == 0
    
    def test_schedule_retry_with_hook(self):
        """Test retry scheduling with on_retry hook."""
        manager = RetryManager()
        callback = Mock()
        manager.set_submit_callback(callback)
        
        hook_called = []
        
        def retry_hook(attempt, delay):
            hook_called.append((attempt, delay))
        
        task = self._create_failed_task()
        task.on_retry = retry_hook
        
        # Schedule retry
        result = manager.schedule_retry(task)
        
        assert result is True
        assert len(hook_called) == 1
        assert hook_called[0] == (1, pytest.approx(1.0, rel=0.25))  # With jitter
    
    def test_cancel_retry(self):
        """Test canceling a scheduled retry."""
        manager = RetryManager()
        
        task = self._create_failed_task()
        
        # Schedule retry
        result = manager.schedule_retry(task)
        assert result is True
        assert manager.get_retry_count() == 1
        
        # Cancel retry
        cancelled = manager.cancel_retry(task.id)
        assert cancelled is True
        assert manager.get_retry_count() == 0
        
        # Try to cancel again
        cancelled = manager.cancel_retry(task.id)
        assert cancelled is False
    
    def test_retry_worker_execution(self):
        """Test retry worker executes callbacks."""
        manager = RetryManager()
        callback = Mock()
        manager.set_submit_callback(callback)
        
        task = self._create_failed_task(retry_policy=RetryPolicy(max_attempts=3, base_delay=0.001))
        
        # Schedule retry
        manager.schedule_retry(task)
        
        # Wait for retry worker to process
        time.sleep(0.2)
        
        # Callback should have been called
        callback.assert_called_once_with(task.id)
    
    def test_shutdown(self):
        """Test retry manager shutdown."""
        manager = RetryManager()
        
        task = self._create_failed_task()
        
        # Schedule retry
        manager.schedule_retry(task)
        assert manager.get_retry_count() == 1
        
        # Shutdown
        manager.shutdown()
        
        # Wait for shutdown to complete
        time.sleep(0.1)
        
        # Should clear retry queue
        assert manager.get_retry_count() == 0


class TestRetryPolicyUtilities:
    """Test retry policy utility functions."""
    
    def test_fixed_retry_policy(self):
        """Test fixed_retry_policy utility."""
        policy = fixed_retry_policy(
            max_attempts=5,
            delay=2.0,
            jitter=False,
            retry_on=[ValueError],
            abort_on=[RuntimeError]
        )
        
        assert isinstance(policy, RetryPolicy)
        assert policy.max_attempts == 5
        assert policy.strategy == "fixed"
        assert policy.base_delay == 2.0
        assert policy.jitter is False
        assert policy.retry_on == [ValueError]
        assert policy.abort_on == [RuntimeError]
    
    def test_exponential_retry_policy(self):
        """Test exponential_retry_policy utility."""
        policy = exponential_retry_policy(
            max_attempts=4,
            base_delay=1.5,
            max_delay=30.0
        )
        
        assert isinstance(policy, RetryPolicy)
        assert policy.max_attempts == 4
        assert policy.strategy == "exponential"
        assert policy.base_delay == 1.5
        assert policy.max_delay == 30.0
        assert policy.jitter is True  # Default
    
    def test_linear_retry_policy(self):
        """Test linear_retry_policy utility."""
        policy = linear_retry_policy(
            max_attempts=3,
            base_delay=1.0,
            jitter=False
        )
        
        assert isinstance(policy, RetryPolicy)
        assert policy.max_attempts == 3
        assert policy.strategy == "linear"
        assert policy.base_delay == 1.0
        assert policy.jitter is False
    
    def test_fibonacci_retry_policy(self):
        """Test fibonacci_retry_policy utility."""
        policy = fibonacci_retry_policy(
            max_attempts=6,
            base_delay=0.5,
            max_delay=20.0
        )
        
        assert isinstance(policy, RetryPolicy)
        assert policy.max_attempts == 6
        assert policy.strategy == "fibonacci"
        assert policy.base_delay == 0.5
        assert policy.max_delay == 20.0
