"""Retry decorators for the Trax platform.

This module provides decorators for adding retry logic to both synchronous
and asynchronous functions with exponential backoff and jitter.
"""

import asyncio
import functools
import time
from typing import Any, Callable, Optional, Type, TypeVar, Union

from .base import RetryConfig, RetryState, CircuitBreaker, DEFAULT_RETRY_CONFIG, RetryStrategy
from ..errors import wrap_error
from ..logging import get_logger

T = TypeVar('T')


def retry(
    config: Optional[RetryConfig] = None,
    retryable_exceptions: Optional[list[Type[Exception]]] = None,
    non_retryable_exceptions: Optional[list[Type[Exception]]] = None,
    circuit_breaker: Optional[CircuitBreaker] = None
):
    """Decorator for adding retry logic to synchronous functions.
    
    Args:
        config: Retry configuration. If None, uses DEFAULT_RETRY_CONFIG.
        retryable_exceptions: List of exception types that should trigger retries.
        non_retryable_exceptions: List of exception types that should not trigger retries.
        circuit_breaker: Optional circuit breaker instance.
    
    Returns:
        Decorated function with retry logic.
    """
    def decorator(func: Callable[..., T]) -> Callable[..., T]:
        @functools.wraps(func)
        def wrapper(*args, **kwargs) -> T:
            # Use provided config or default
            retry_config = config or DEFAULT_RETRY_CONFIG
            
            # Override exception lists if provided
            if retryable_exceptions is not None:
                retry_config.retryable_exceptions = retryable_exceptions
            if non_retryable_exceptions is not None:
                retry_config.non_retryable_exceptions = non_retryable_exceptions
            
            # Initialize retry state
            state = RetryState(retry_config)
            logger = get_logger(func.__module__)
            
            # Check circuit breaker if enabled
            if circuit_breaker and not circuit_breaker.can_execute():
                raise RuntimeError(f"Circuit breaker is {circuit_breaker.get_state()}")
            
            while True:
                try:
                    # Execute the function
                    result = func(*args, **kwargs)
                    
                    # Handle circuit breaker success
                    if circuit_breaker:
                        circuit_breaker.on_success()
                    
                    # Log success if retries were attempted
                    if state.attempt > 0:
                        state.log_retry_success(result)
                    
                    return result
                    
                except Exception as e:
                    # Increment attempt counter
                    state.increment_attempt()
                    state.set_last_error(e)
                    
                    # Check if we should retry
                    if not state.should_retry(e):
                        # Handle circuit breaker failure
                        if circuit_breaker:
                            circuit_breaker.on_failure()
                        
                        # Log max retries exceeded
                        if state.attempt > 1:
                            state.log_max_retries_exceeded(e)
                        
                        # Wrap and re-raise the error
                        raise wrap_error(e)
                    
                    # Calculate delay for next retry
                    delay = state.calculate_delay()
                    state.set_last_delay(delay)
                    
                    # Log retry attempt
                    state.log_retry_attempt(e, delay)
                    
                    # Wait before retrying
                    if delay > 0:
                        time.sleep(delay)
                    
                    # Continue to next iteration
                    continue
        
        return wrapper
    return decorator


def async_retry(
    config: Optional[RetryConfig] = None,
    retryable_exceptions: Optional[list[Type[Exception]]] = None,
    non_retryable_exceptions: Optional[list[Type[Exception]]] = None,
    circuit_breaker: Optional[CircuitBreaker] = None
):
    """Decorator for adding retry logic to asynchronous functions.
    
    Args:
        config: Retry configuration. If None, uses DEFAULT_RETRY_CONFIG.
        retryable_exceptions: List of exception types that should trigger retries.
        non_retryable_exceptions: List of exception types that should not trigger retries.
        circuit_breaker: Optional circuit breaker instance.
    
    Returns:
        Decorated async function with retry logic.
    """
    def decorator(func: Callable[..., T]) -> Callable[..., T]:
        @functools.wraps(func)
        async def wrapper(*args, **kwargs) -> T:
            # Use provided config or default
            retry_config = config or DEFAULT_RETRY_CONFIG
            
            # Override exception lists if provided
            if retryable_exceptions is not None:
                retry_config.retryable_exceptions = retryable_exceptions
            if non_retryable_exceptions is not None:
                retry_config.non_retryable_exceptions = non_retryable_exceptions
            
            # Initialize retry state
            state = RetryState(retry_config)
            logger = get_logger(func.__module__)
            
            # Check circuit breaker if enabled
            if circuit_breaker and not circuit_breaker.can_execute():
                raise RuntimeError(f"Circuit breaker is {circuit_breaker.get_state()}")
            
            while True:
                try:
                    # Execute the async function
                    result = await func(*args, **kwargs)
                    
                    # Handle circuit breaker success
                    if circuit_breaker:
                        circuit_breaker.on_success()
                    
                    # Log success if retries were attempted
                    if state.attempt > 0:
                        state.log_retry_success(result)
                    
                    return result
                    
                except Exception as e:
                    # Increment attempt counter
                    state.increment_attempt()
                    state.set_last_error(e)
                    
                    # Check if we should retry
                    if not state.should_retry(e):
                        # Handle circuit breaker failure
                        if circuit_breaker:
                            circuit_breaker.on_failure()
                        
                        # Log max retries exceeded
                        if state.attempt > 1:
                            state.log_max_retries_exceeded(e)
                        
                        # Wrap and re-raise the error
                        raise wrap_error(e)
                    
                    # Calculate delay for next retry
                    delay = state.calculate_delay()
                    state.set_last_delay(delay)
                    
                    # Log retry attempt
                    state.log_retry_attempt(e, delay)
                    
                    # Wait before retrying
                    if delay > 0:
                        await asyncio.sleep(delay)
                    
                    # Continue to next iteration
                    continue
        
        return wrapper
    return decorator


# Convenience decorators for common use cases
def retry_network(
    max_retries: int = 3,
    initial_delay: float = 1.0,
    max_delay: float = 30.0
):
    """Decorator for network operations with retry logic."""
    config = RetryConfig(
        max_retries=max_retries,
        initial_delay=initial_delay,
        max_delay=max_delay,
        exponential_multiplier=2.0,
        jitter_factor=0.1,
        strategy=RetryStrategy.EXPONENTIAL
    )
    return retry(config=config)


def async_retry_network(
    max_retries: int = 3,
    initial_delay: float = 1.0,
    max_delay: float = 30.0
):
    """Decorator for async network operations with retry logic."""
    config = RetryConfig(
        max_retries=max_retries,
        initial_delay=initial_delay,
        max_delay=max_delay,
        exponential_multiplier=2.0,
        jitter_factor=0.1,
        strategy=RetryStrategy.EXPONENTIAL
    )
    return async_retry(config=config)


def retry_api(
    max_retries: int = 5,
    initial_delay: float = 2.0,
    max_delay: float = 60.0
):
    """Decorator for API operations with retry logic."""
    config = RetryConfig(
        max_retries=max_retries,
        initial_delay=initial_delay,
        max_delay=max_delay,
        exponential_multiplier=2.0,
        jitter_factor=0.2,
        strategy=RetryStrategy.EXPONENTIAL
    )
    return retry(config=config)


def async_retry_api(
    max_retries: int = 5,
    initial_delay: float = 2.0,
    max_delay: float = 60.0
):
    """Decorator for async API operations with retry logic."""
    config = RetryConfig(
        max_retries=max_retries,
        initial_delay=initial_delay,
        max_delay=max_delay,
        exponential_multiplier=2.0,
        jitter_factor=0.2,
        strategy=RetryStrategy.EXPONENTIAL
    )
    return async_retry(config=config)


def retry_processing(
    max_retries: int = 2,
    initial_delay: float = 5.0,
    max_delay: float = 30.0
):
    """Decorator for processing operations with retry logic."""
    config = RetryConfig(
        max_retries=max_retries,
        initial_delay=initial_delay,
        max_delay=max_delay,
        exponential_multiplier=1.5,
        jitter_factor=0.05,
        strategy=RetryStrategy.EXPONENTIAL
    )
    return retry(config=config)


def async_retry_processing(
    max_retries: int = 2,
    initial_delay: float = 5.0,
    max_delay: float = 30.0
):
    """Decorator for async processing operations with retry logic."""
    config = RetryConfig(
        max_retries=max_retries,
        initial_delay=initial_delay,
        max_delay=max_delay,
        exponential_multiplier=1.5,
        jitter_factor=0.05,
        strategy=RetryStrategy.EXPONENTIAL
    )
    return async_retry(config=config)


# Context managers for manual retry control
class RetryContext:
    """Context manager for manual retry control."""
    
    def __init__(self, config: Optional[RetryConfig] = None):
        self.config = config or DEFAULT_RETRY_CONFIG
        self.state = RetryState(self.config)
        self.logger = get_logger(__name__)
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None:
            self.state.increment_attempt()
            self.state.set_last_error(exc_val)
            
            if self.state.should_retry(exc_val):
                delay = self.state.calculate_delay()
                self.state.set_last_delay(delay)
                self.state.log_retry_attempt(exc_val, delay)
                
                if delay > 0:
                    time.sleep(delay)
                
                return False  # Don't suppress the exception
            else:
                if self.state.attempt > 1:
                    self.state.log_max_retries_exceeded(exc_val)
        
        return False  # Don't suppress exceptions


class AsyncRetryContext:
    """Async context manager for manual retry control."""
    
    def __init__(self, config: Optional[RetryConfig] = None):
        self.config = config or DEFAULT_RETRY_CONFIG
        self.state = RetryState(self.config)
        self.logger = get_logger(__name__)
    
    async def __aenter__(self):
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None:
            self.state.increment_attempt()
            self.state.set_last_error(exc_val)
            
            if self.state.should_retry(exc_val):
                delay = self.state.calculate_delay()
                self.state.set_last_delay(delay)
                self.state.log_retry_attempt(exc_val, delay)
                
                if delay > 0:
                    await asyncio.sleep(delay)
                
                return False  # Don't suppress the exception
            else:
                if self.state.attempt > 1:
                    self.state.log_max_retries_exceeded(exc_val)
        
        return False  # Don't suppress exceptions
