"""
BGX Retry module - Retry management and backoff strategies.

This module provides the RetryManager class that handles task retry logic,
delay computation, and retry scheduling.
"""

import threading
import time
from typing import Dict, Optional

from .task import Task, TaskState, RetryPolicy


class RetryManager:
    """Manages task retries with configurable backoff strategies."""
    
    def __init__(self):
        """Initialize the retry manager."""
        self._retry_queue: Dict[str, float] = {}  # task_id -> retry_time
        self._retry_lock = threading.Lock()
        self._retry_thread = threading.Thread(
            target=self._retry_worker,
            name="bgx-retry-worker",
            daemon=True
        )
        self._shutdown_event = threading.Event()
        self._retry_thread.start()
        
        # Callback for submitting retried tasks
        self._submit_callback: Optional[callable] = None
    
    def set_submit_callback(self, callback: callable) -> None:
        """Set callback function for submitting retried tasks."""
        self._submit_callback = callback
    
    def schedule_retry(self, task: Task) -> bool:
        """Schedule a task for retry if it can be retried.
        
        Args:
            task: Task to retry
            
        Returns:
            True if task was scheduled for retry, False otherwise
        """
        if not task.can_retry():
            return False
        
        # Compute retry delay
        delay = task.compute_retry_delay()
        retry_time = time.time() + delay
        
        # Schedule retry
        with self._retry_lock:
            self._retry_queue[task.id] = retry_time
        
        # Transition to RETRYING state
        task.transition_to(TaskState.RETRYING)
        
        # Increment attempt counter
        task.attempt += 1
        
        # Call retry hook
        if task.on_retry:
            try:
                task.on_retry(task.attempt - 1, delay)
            except Exception as e:
                print(f"Retry hook error for task {task.id}: {e}")
        
        return True
    
    def _retry_worker(self) -> None:
        """Worker thread that processes retry queue."""
        while not self._shutdown_event.is_set():
            current_time = time.time()
            tasks_to_retry = []
            
            # Find tasks ready for retry
            with self._retry_lock:
                for task_id, retry_time in list(self._retry_queue.items()):
                    if current_time >= retry_time:
                        tasks_to_retry.append(task_id)
                
                # Remove scheduled tasks from queue
                for task_id in tasks_to_retry:
                    del self._retry_queue[task_id]
            
            # Submit retried tasks
            for task_id in tasks_to_retry:
                if self._submit_callback:
                    try:
                        self._submit_callback(task_id)
                    except Exception as e:
                        print(f"Error submitting retry for task {task_id}: {e}")
            
            # Sleep before next check
            time.sleep(0.1)  # 100ms resolution
    
    def get_retry_count(self) -> int:
        """Get number of tasks currently scheduled for retry."""
        with self._retry_lock:
            return len(self._retry_queue)
    
    def get_retry_info(self) -> Dict[str, float]:
        """Get information about scheduled retries."""
        with self._retry_lock:
            return self._retry_queue.copy()
    
    def cancel_retry(self, task_id: str) -> bool:
        """Cancel a scheduled retry.
        
        Args:
            task_id: ID of task to cancel retry for
            
        Returns:
            True if retry was cancelled, False if not found
        """
        with self._retry_lock:
            if task_id in self._retry_queue:
                del self._retry_queue[task_id]
                return True
            return False
    
    def shutdown(self) -> None:
        """Shutdown the retry manager."""
        self._shutdown_event.set()
        self._retry_thread.join(timeout=5.0)
        
        # Clear retry queue
        with self._retry_lock:
            self._retry_queue.clear()


class BackoffCalculator:
    """Utility class for computing backoff delays."""
    
    @staticmethod
    def compute_delay(
        strategy: str,
        attempt: int,
        base_delay: float,
        max_delay: float = 60.0,
        jitter: bool = True
    ) -> float:
        """Compute delay using specified backoff strategy.
        
        Args:
            strategy: Backoff strategy ("fixed", "linear", "exponential", "fibonacci")
            attempt: Current attempt number (1-indexed)
            base_delay: Base delay in seconds
            max_delay: Maximum delay cap
            jitter: Whether to add random jitter
            
        Returns:
            Computed delay in seconds
        """
        if strategy == "fixed":
            delay = base_delay
        elif strategy == "linear":
            delay = base_delay * attempt
        elif strategy == "exponential":
            delay = base_delay * (2 ** (attempt - 1))
        elif strategy == "fibonacci":
            delay = base_delay * BackoffCalculator._fibonacci(attempt)
        else:
            raise ValueError(f"Unknown backoff strategy: {strategy}")
        
        # Apply max delay cap
        delay = min(delay, max_delay)
        
        # Add jitter if enabled
        if jitter:
            import random
            jitter_amount = delay * 0.25
            delay += random.uniform(0, jitter_amount)
        
        return delay
    
    @staticmethod
    def _fibonacci(n: int) -> int:
        """Compute nth Fibonacci number."""
        if n <= 0:
            return 0
        elif n == 1:
            return 1
        elif n == 2:
            return 1
        else:
            a, b = 1, 1
            for _ in range(3, n + 1):
                a, b = b, a + b
            return b


# Convenience functions for common retry policies
def fixed_retry_policy(
    max_attempts: int = 3,
    delay: float = 1.0,
    jitter: bool = True,
    retry_on: list = None,
    abort_on: list = None
) -> RetryPolicy:
    """Create a fixed-delay retry policy."""
    return RetryPolicy(
        max_attempts=max_attempts,
        strategy="fixed",
        base_delay=delay,
        jitter=jitter,
        retry_on=retry_on or [Exception],
        abort_on=abort_on or []
    )


def exponential_retry_policy(
    max_attempts: int = 3,
    base_delay: float = 1.0,
    max_delay: float = 60.0,
    jitter: bool = True,
    retry_on: list = None,
    abort_on: list = None
) -> RetryPolicy:
    """Create an exponential backoff retry policy."""
    return RetryPolicy(
        max_attempts=max_attempts,
        strategy="exponential",
        base_delay=base_delay,
        max_delay=max_delay,
        jitter=jitter,
        retry_on=retry_on or [Exception],
        abort_on=abort_on or []
    )


def linear_retry_policy(
    max_attempts: int = 3,
    base_delay: float = 1.0,
    max_delay: float = 60.0,
    jitter: bool = True,
    retry_on: list = None,
    abort_on: list = None
) -> RetryPolicy:
    """Create a linear backoff retry policy."""
    return RetryPolicy(
        max_attempts=max_attempts,
        strategy="linear",
        base_delay=base_delay,
        max_delay=max_delay,
        jitter=jitter,
        retry_on=retry_on or [Exception],
        abort_on=abort_on or []
    )


def fibonacci_retry_policy(
    max_attempts: int = 3,
    base_delay: float = 1.0,
    max_delay: float = 60.0,
    jitter: bool = True,
    retry_on: list = None,
    abort_on: list = None
) -> RetryPolicy:
    """Create a Fibonacci backoff retry policy."""
    return RetryPolicy(
        max_attempts=max_attempts,
        strategy="fibonacci",
        base_delay=base_delay,
        max_delay=max_delay,
        jitter=jitter,
        retry_on=retry_on or [Exception],
        abort_on=abort_on or []
    )
