"""
Enhanced Retry Mechanisms for Connectors
Exponential backoff, jitter, and intelligent retry strategies
"""

import time
import random
import threading
from typing import Any, Callable, Optional, Type, Union, List, Dict
from functools import wraps
from dataclasses import dataclass
from .exceptions import RetryExhaustedError, NetworkError, TimeoutError
from .logging_config import get_logger, log_performance

logger = get_logger(__name__)


@dataclass
class RetryConfig:
    """Configuration for retry behavior"""
    
    max_attempts: int = 3
    base_delay: float = 1.0
    max_delay: float = 60.0
    exponential_base: float = 2.0
    jitter: bool = True
    jitter_factor: float = 0.1
    
    # Retry conditions
    retry_on_exceptions: List[Type[Exception]] = None
    retry_on_status_codes: List[int] = None
    retry_condition: Optional[Callable[[Exception], bool]] = None
    
    def __post_init__(self):
        """Set default values"""
        if self.retry_on_exceptions is None:
            self.retry_on_exceptions = [
                NetworkError, TimeoutError, ConnectionError, OSError
            ]
        
        if self.retry_on_status_codes is None:
            self.retry_on_status_codes = [429, 502, 503, 504]


class RetryState:
    """State tracking for retry attempts"""
    
    def __init__(self, config: RetryConfig):
        self.config = config
        self.attempt = 0
        self.total_delay = 0.0
        self.last_exception = None
        self.start_time = time.time()
    
    def should_retry(self, exception: Exception) -> bool:
        """Determine if operation should be retried"""
        self.attempt += 1
        
        if self.attempt > self.config.max_attempts:
            return False
        
        self.last_exception = exception
        
        # Check exception type
        if any(isinstance(exception, exc_type) 
               for exc_type in self.config.retry_on_exceptions):
            return True
        
        # Check custom retry condition
        if self.config.retry_condition and self.config.retry_condition(exception):
            return True
        
        # Check status code for API errors
        if hasattr(exception, 'status_code'):
            if exception.status_code in self.config.retry_on_status_codes:
                return True
        
        return False
    
    def calculate_delay(self) -> float:
        """Calculate delay before next retry"""
        delay = self.config.base_delay * (
            self.config.exponential_base ** (self.attempt - 1)
        )
        
        # Apply maximum delay limit
        delay = min(delay, self.config.max_delay)
        
        # Add jitter if enabled
        if self.config.jitter:
            jitter_amount = delay * self.config.jitter_factor
            jitter = random.uniform(-jitter_amount, jitter_amount)
            delay += jitter
        
        # Ensure non-negative delay
        delay = max(0, delay)
        
        self.total_delay += delay
        return delay
    
    def get_stats(self) -> Dict[str, Any]:
        """Get retry statistics"""
        return {
            'attempts': self.attempt,
            'total_delay': self.total_delay,
            'elapsed_time': time.time() - self.start_time,
            'last_exception': str(self.last_exception) if self.last_exception else None
        }


class RetryStrategy:
    """Base class for retry strategies"""
    
    def should_retry(self, state: RetryState, exception: Exception) -> bool:
        """Determine if operation should be retried"""
        raise NotImplementedError
    
    def calculate_delay(self, state: RetryState) -> float:
        """Calculate delay before next retry"""
        raise NotImplementedError


class ExponentialBackoffStrategy(RetryStrategy):
    """Exponential backoff retry strategy"""
    
    def __init__(self, config: RetryConfig):
        self.config = config
    
    def should_retry(self, state: RetryState, exception: Exception) -> bool:
        return state.should_retry(exception)
    
    def calculate_delay(self, state: RetryState) -> float:
        return state.calculate_delay()


class LinearBackoffStrategy(RetryStrategy):
    """Linear backoff retry strategy"""
    
    def __init__(self, config: RetryConfig):
        self.config = config
    
    def should_retry(self, state: RetryState, exception: Exception) -> bool:
        return state.should_retry(exception)
    
    def calculate_delay(self, state: RetryState) -> float:
        delay = self.config.base_delay * state.attempt
        delay = min(delay, self.config.max_delay)
        
        if self.config.jitter:
            jitter_amount = delay * self.config.jitter_factor
            jitter = random.uniform(-jitter_amount, jitter_amount)
            delay += jitter
        
        return max(0, delay)


class FixedDelayStrategy(RetryStrategy):
    """Fixed delay retry strategy"""
    
    def __init__(self, config: RetryConfig):
        self.config = config
    
    def should_retry(self, state: RetryState, exception: Exception) -> bool:
        return state.should_retry(exception)
    
    def calculate_delay(self, state: RetryState) -> float:
        delay = self.config.base_delay
        
        if self.config.jitter:
            jitter_amount = delay * self.config.jitter_factor
            jitter = random.uniform(-jitter_amount, jitter_amount)
            delay += jitter
        
        return max(0, delay)


class AdaptiveRetryStrategy(RetryStrategy):
    """Adaptive retry strategy that adjusts based on failure patterns"""
    
    def __init__(self, config: RetryConfig):
        self.config = config
        self.failure_history = []
        self.max_history = 10
    
    def should_retry(self, state: RetryState, exception: Exception) -> bool:
        # Record failure
        self.failure_history.append(time.time())
        
        # Keep only recent failures
        if len(self.failure_history) > self.max_history:
            self.failure_history.pop(0)
        
        return state.should_retry(exception)
    
    def calculate_delay(self, state: RetryState) -> float:
        base_delay = state.calculate_delay()
        
        # Adaptive factor based on recent failure frequency
        if len(self.failure_history) > 1:
            recent_failures = [t for t in self.failure_history 
                             if time.time() - t < 300]  # Last 5 minutes
            failure_rate = len(recent_failures) / 300.0
            
            # Increase delay if failures are frequent
            if failure_rate > 0.1:  # More than 1 failure per 10 seconds
                base_delay *= 2.0
        
        return base_delay


def retry(
    max_attempts: int = 3,
    base_delay: float = 1.0,
    max_delay: float = 60.0,
    exponential_base: float = 2.0,
    jitter: bool = True,
    jitter_factor: float = 0.1,
    retry_on_exceptions: Optional[List[Type[Exception]]] = None,
    retry_on_status_codes: Optional[List[int]] = None,
    retry_condition: Optional[Callable[[Exception], bool]] = None,
    strategy: Optional[RetryStrategy] = None,
    on_retry: Optional[Callable[[RetryState], None]] = None,
    log_retries: bool = True
):
    """
    Decorator for retrying functions with configurable backoff strategy
    
    Args:
        max_attempts: Maximum number of retry attempts
        base_delay: Base delay between retries in seconds
        max_delay: Maximum delay between retries in seconds
        exponential_base: Base for exponential backoff
        jitter: Whether to add jitter to delays
        jitter_factor: Factor for jitter calculation (0-1)
        retry_on_exceptions: List of exception types to retry on
        retry_on_status_codes: List of HTTP status codes to retry on
        retry_condition: Custom function to determine if should retry
        strategy: Custom retry strategy
        on_retry: Callback function called before each retry
        log_retries: Whether to log retry attempts
    """
    
    def decorator(func: Callable) -> Callable:
        @wraps(func)
        @log_performance
        def wrapper(*args, **kwargs):
            # Create retry configuration
            config = RetryConfig(
                max_attempts=max_attempts,
                base_delay=base_delay,
                max_delay=max_delay,
                exponential_base=exponential_base,
                jitter=jitter,
                jitter_factor=jitter_factor,
                retry_on_exceptions=retry_on_exceptions,
                retry_on_status_codes=retry_on_status_codes,
                retry_condition=retry_condition
            )
            
            # Create retry state
            state = RetryState(config)
            
            # Use custom strategy if provided
            retry_strategy = strategy or ExponentialBackoffStrategy(config)
            
            while True:
                try:
                    result = func(*args, **kwargs)
                    
                    # Log successful retry if applicable
                    if state.attempt > 1 and log_retries:
                        stats = state.get_stats()
                        logger.info(
                            f"Function {func.__name__} succeeded after {state.attempt} attempts",
                            extra={'retry_stats': stats}
                        )
                    
                    return result
                
                except Exception as e:
                    if not retry_strategy.should_retry(state, e):
                        # Create retry exhausted error
                        stats = state.get_stats()
                        raise RetryExhaustedError(
                            f"Function {func.__name__} failed after {state.attempt} attempts: {e}",
                            attempts=state.attempt,
                            last_error=e
                        ) from e
                    
                    # Calculate delay
                    delay = retry_strategy.calculate_delay(state)
                    
                    # Log retry attempt
                    if log_retries:
                        logger.warning(
                            f"Function {func.__name__} failed (attempt {state.attempt}/{config.max_attempts}), "
                            f"retrying in {delay:.2f}s: {e}",
                            extra={
                                'function': func.__name__,
                                'attempt': state.attempt,
                                'max_attempts': config.max_attempts,
                                'delay': delay,
                                'exception': str(e)
                            }
                        )
                    
                    # Call retry callback if provided
                    if on_retry:
                        try:
                            on_retry(state)
                        except Exception as callback_error:
                            logger.error(f"Retry callback failed: {callback_error}")
                    
                    # Wait before retry
                    if delay > 0:
                        time.sleep(delay)
        
        return wrapper
    return decorator


def async_retry(
    max_attempts: int = 3,
    base_delay: float = 1.0,
    max_delay: float = 60.0,
    exponential_base: float = 2.0,
    jitter: bool = True,
    jitter_factor: float = 0.1,
    retry_on_exceptions: Optional[List[Type[Exception]]] = None,
    retry_condition: Optional[Callable[[Exception], bool]] = None,
    strategy: Optional[RetryStrategy] = None,
    on_retry: Optional[Callable[[RetryState], None]] = None,
    log_retries: bool = True
):
    """
    Async decorator for retrying functions with configurable backoff strategy
    """
    import asyncio
    
    def decorator(func: Callable) -> Callable:
        @wraps(func)
        async def wrapper(*args, **kwargs):
            # Create retry configuration
            config = RetryConfig(
                max_attempts=max_attempts,
                base_delay=base_delay,
                max_delay=max_delay,
                exponential_base=exponential_base,
                jitter=jitter,
                jitter_factor=jitter_factor,
                retry_on_exceptions=retry_on_exceptions,
                retry_condition=retry_condition
            )
            
            # Create retry state
            state = RetryState(config)
            
            # Use custom strategy if provided
            retry_strategy = strategy or ExponentialBackoffStrategy(config)
            
            while True:
                try:
                    result = await func(*args, **kwargs)
                    
                    # Log successful retry if applicable
                    if state.attempt > 1 and log_retries:
                        stats = state.get_stats()
                        logger.info(
                            f"Async function {func.__name__} succeeded after {state.attempt} attempts",
                            extra={'retry_stats': stats}
                        )
                    
                    return result
                
                except Exception as e:
                    if not retry_strategy.should_retry(state, e):
                        # Create retry exhausted error
                        stats = state.get_stats()
                        raise RetryExhaustedError(
                            f"Async function {func.__name__} failed after {state.attempt} attempts: {e}",
                            attempts=state.attempt,
                            last_error=e
                        ) from e
                    
                    # Calculate delay
                    delay = retry_strategy.calculate_delay(state)
                    
                    # Log retry attempt
                    if log_retries:
                        logger.warning(
                            f"Async function {func.__name__} failed (attempt {state.attempt}/{config.max_attempts}), "
                            f"retrying in {delay:.2f}s: {e}",
                            extra={
                                'function': func.__name__,
                                'attempt': state.attempt,
                                'max_attempts': config.max_attempts,
                                'delay': delay,
                                'exception': str(e)
                            }
                        )
                    
                    # Call retry callback if provided
                    if on_retry:
                        try:
                            on_retry(state)
                        except Exception as callback_error:
                            logger.error(f"Retry callback failed: {callback_error}")
                    
                    # Wait before retry
                    if delay > 0:
                        await asyncio.sleep(delay)
        
        return wrapper
    return decorator


class RetryCircuitBreaker:
    """Circuit breaker pattern for retry operations"""
    
    def __init__(
        self,
        failure_threshold: int = 5,
        recovery_timeout: float = 60.0,
        expected_exception: Type[Exception] = Exception
    ):
        self.failure_threshold = failure_threshold
        self.recovery_timeout = recovery_timeout
        self.expected_exception = expected_exception
        
        self.failure_count = 0
        self.last_failure_time = None
        self.state = 'CLOSED'  # CLOSED, OPEN, HALF_OPEN
        self._lock = threading.Lock()
    
    def __call__(self, func: Callable) -> Callable:
        @wraps(func)
        def wrapper(*args, **kwargs):
            with self._lock:
                if self.state == 'OPEN':
                    if time.time() - self.last_failure_time > self.recovery_timeout:
                        self.state = 'HALF_OPEN'
                        logger.info("Circuit breaker transitioning to HALF_OPEN")
                    else:
                        raise RetryExhaustedError(
                            f"Circuit breaker is OPEN. Call to {func.__name__} blocked."
                        )
            
            try:
                result = func(*args, **kwargs)
                
                with self._lock:
                    if self.state == 'HALF_OPEN':
                        self.state = 'CLOSED'
                        self.failure_count = 0
                        logger.info("Circuit breaker transitioning to CLOSED")
                
                return result
            
            except self.expected_exception as e:
                with self._lock:
                    self.failure_count += 1
                    self.last_failure_time = time.time()
                    
                    if self.failure_count >= self.failure_threshold:
                        self.state = 'OPEN'
                        logger.warning(
                            f"Circuit breaker transitioning to OPEN after {self.failure_count} failures"
                        )
                
                raise
        
        return wrapper
    
    def reset(self):
        """Reset the circuit breaker to CLOSED state"""
        with self._lock:
            self.failure_count = 0
            self.last_failure_time = None
            self.state = 'CLOSED'
            logger.info("Circuit breaker reset to CLOSED")
    
    def get_state(self) -> Dict[str, Any]:
        """Get circuit breaker state"""
        with self._lock:
            return {
                'state': self.state,
                'failure_count': self.failure_count,
                'last_failure_time': self.last_failure_time,
                'failure_threshold': self.failure_threshold,
                'recovery_timeout': self.recovery_timeout
            }
