"""
Retry utilities for robust error handling in Toxo.

This module provides retry mechanisms with exponential backoff,
jitter, and sophisticated error handling strategies.
"""

import asyncio
import random
import time
import logging
from typing import (
    Callable, Any, Optional, Union, Type, List, Dict,
    Awaitable, TypeVar, Tuple
)
from functools import wraps
from contextlib import asynccontextmanager, contextmanager
from dataclasses import dataclass
from enum import Enum

from .exceptions import (
    ToxoError, APIError, NetworkError, TimeoutError, RateLimitError,
    AuthenticationError, ValidationError, TrainingError, MemoryError,
    TransactionError, is_retryable_error, format_error_context
)

T = TypeVar('T')
logger = logging.getLogger(__name__)


class RetryStrategy(Enum):
    """Retry strategy types."""
    EXPONENTIAL_BACKOFF = "exponential_backoff"
    LINEAR_BACKOFF = "linear_backoff" 
    FIXED_DELAY = "fixed_delay"
    FIBONACCI = "fibonacci"


@dataclass
class RetryConfig:
    """Configuration for retry behavior."""
    max_attempts: int = 3
    base_delay: float = 1.0
    max_delay: float = 60.0
    strategy: RetryStrategy = RetryStrategy.EXPONENTIAL_BACKOFF
    backoff_multiplier: float = 2.0
    jitter: bool = True
    jitter_max: float = 0.1
    
    # Specific error handling
    retryable_exceptions: Tuple[Type[Exception], ...] = (
        APIError, NetworkError, TimeoutError, RateLimitError,
        ConnectionError, OSError
    )
    non_retryable_exceptions: Tuple[Type[Exception], ...] = (
        AuthenticationError, ValidationError
    )
    
    # Conditions for retry
    retry_on_status_codes: List[int] = None
    retry_on_error_codes: List[str] = None
    
    def __post_init__(self):
        """Validate and set defaults."""
        if self.retry_on_status_codes is None:
            self.retry_on_status_codes = [429, 500, 502, 503, 504]
        if self.retry_on_error_codes is None:
            self.retry_on_error_codes = [
                "RATE_LIMIT_ERROR", "NETWORK_ERROR", "TIMEOUT_ERROR", "API_ERROR"
            ]


class RetryState:
    """Tracks state across retry attempts."""
    
    def __init__(self, config: RetryConfig):
        self.config = config
        self.attempt = 0
        self.total_delay = 0.0
        self.errors: List[Exception] = []
        self.start_time = time.time()
    
    def should_retry(self, error: Exception) -> bool:
        """Determine if we should retry based on the error and current state."""
        # Check attempt limit
        if self.attempt >= self.config.max_attempts:
            return False
        
        # Check for non-retryable exceptions
        if isinstance(error, self.config.non_retryable_exceptions):
            return False
        
        # Check for explicitly retryable exceptions
        if isinstance(error, self.config.retryable_exceptions):
            return True
        
        # Check Toxo error retryable flag
        if isinstance(error, ToxoError):
            if not error.retryable:
                return False
            
            # Check specific error codes
            if error.error_code in self.config.retry_on_error_codes:
                return True
        
        # Check API errors with status codes
        if isinstance(error, APIError):
            status_code = error.context.get('status_code')
            if status_code in self.config.retry_on_status_codes:
                return True
        
        # Use general retryable check as fallback
        return is_retryable_error(error)
    
    def get_next_delay(self) -> float:
        """Calculate the next delay based on retry strategy."""
        if self.config.strategy == RetryStrategy.EXPONENTIAL_BACKOFF:
            delay = self.config.base_delay * (self.config.backoff_multiplier ** self.attempt)
        elif self.config.strategy == RetryStrategy.LINEAR_BACKOFF:
            delay = self.config.base_delay * (self.attempt + 1)
        elif self.config.strategy == RetryStrategy.FIXED_DELAY:
            delay = self.config.base_delay
        elif self.config.strategy == RetryStrategy.FIBONACCI:
            delay = self._fibonacci_delay()
        else:
            delay = self.config.base_delay
        
        # Apply maximum delay limit
        delay = min(delay, self.config.max_delay)
        
        # Add jitter if enabled
        if self.config.jitter:
            jitter_amount = random.uniform(0, self.config.jitter_max * delay)
            delay += jitter_amount
        
        return delay
    
    def _fibonacci_delay(self) -> float:
        """Calculate Fibonacci-based delay."""
        if self.attempt <= 1:
            return self.config.base_delay
        
        a, b = 1, 1
        for _ in range(self.attempt - 1):
            a, b = b, a + b
        
        return self.config.base_delay * b
    
    def record_attempt(self, error: Exception, delay: float):
        """Record an attempt and its error."""
        self.attempt += 1
        self.total_delay += delay
        self.errors.append(error)
        
        logger.warning(
            f"Attempt {self.attempt}/{self.config.max_attempts} failed: {error}. "
            f"Retrying in {delay:.2f}s"
        )


def retry(
    config: Optional[RetryConfig] = None,
    max_attempts: int = 3,
    base_delay: float = 1.0,
    strategy: RetryStrategy = RetryStrategy.EXPONENTIAL_BACKOFF,
    **kwargs
):
    """
    Decorator for adding retry functionality to functions.
    
    Args:
        config: RetryConfig instance (overrides individual parameters)
        max_attempts: Maximum number of retry attempts
        base_delay: Base delay between retries
        strategy: Retry strategy to use
        **kwargs: Additional RetryConfig parameters
    """
    if config is None:
        config = RetryConfig(
            max_attempts=max_attempts,
            base_delay=base_delay,
            strategy=strategy,
            **kwargs
        )
    
    def decorator(func: Callable[..., T]) -> Callable[..., T]:
        if asyncio.iscoroutinefunction(func):
            return _async_retry_wrapper(func, config)
        else:
            return _sync_retry_wrapper(func, config)
    
    return decorator


def _sync_retry_wrapper(func: Callable[..., T], config: RetryConfig) -> Callable[..., T]:
    """Synchronous retry wrapper."""
    @wraps(func)
    def wrapper(*args, **kwargs) -> T:
        state = RetryState(config)
        
        while True:
            try:
                result = func(*args, **kwargs)
                if state.attempt > 0:
                    logger.info(f"Function succeeded after {state.attempt} retries")
                return result
            
            except Exception as error:
                if not state.should_retry(error):
                    # Log final failure
                    logger.error(
                        f"Function failed permanently after {state.attempt} attempts: {error}",
                        extra={"error_context": format_error_context(error)}
                    )
                    raise
                
                delay = state.get_next_delay()
                state.record_attempt(error, delay)
                
                # Handle rate limiting with respect for retry-after header
                if isinstance(error, RateLimitError):
                    retry_after = error.context.get('retry_after')
                    if retry_after:
                        delay = max(delay, retry_after)
                
                time.sleep(delay)
    
    return wrapper


def _async_retry_wrapper(func: Callable[..., Awaitable[T]], config: RetryConfig) -> Callable[..., Awaitable[T]]:
    """Asynchronous retry wrapper."""
    @wraps(func)
    async def wrapper(*args, **kwargs) -> T:
        state = RetryState(config)
        
        while True:
            try:
                result = await func(*args, **kwargs)
                if state.attempt > 0:
                    logger.info(f"Function succeeded after {state.attempt} retries")
                return result
            
            except Exception as error:
                if not state.should_retry(error):
                    # Log final failure
                    logger.error(
                        f"Function failed permanently after {state.attempt} attempts: {error}",
                        extra={"error_context": format_error_context(error)}
                    )
                    raise
                
                delay = state.get_next_delay()
                state.record_attempt(error, delay)
                
                # Handle rate limiting with respect for retry-after header
                if isinstance(error, RateLimitError):
                    retry_after = error.context.get('retry_after')
                    if retry_after:
                        delay = max(delay, retry_after)
                
                await asyncio.sleep(delay)
    
    return wrapper


class Circuit:
    """Circuit breaker for preventing cascading failures."""
    
    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
    
    def __call__(self, func):
        """Decorator for circuit breaker."""
        @wraps(func)
        def wrapper(*args, **kwargs):
            return self._call_with_circuit_breaker(func, *args, **kwargs)
        return wrapper
    
    def _call_with_circuit_breaker(self, func, *args, **kwargs):
        """Execute function with circuit breaker logic."""
        if self._state == "open":
            if self._should_attempt_reset():
                self._state = "half-open"
            else:
                raise APIError("Circuit breaker is open", retryable=False)
        
        try:
            result = func(*args, **kwargs)
            self._on_success()
            return result
        except self.expected_exception as e:
            self._on_failure()
            raise
    
    def _should_attempt_reset(self) -> bool:
        """Check if we should attempt to reset the circuit."""
        return (
            self._last_failure_time and
            time.time() - self._last_failure_time >= self.recovery_timeout
        )
    
    def _on_success(self):
        """Handle successful execution."""
        self._failure_count = 0
        self._state = "closed"
    
    def _on_failure(self):
        """Handle failed execution."""
        self._failure_count += 1
        self._last_failure_time = time.time()
        
        if self._failure_count >= self.failure_threshold:
            self._state = "open"


class Transaction:
    """Transaction-like context manager for rollback operations."""
    
    def __init__(self, name: str = None):
        self.name = name or f"transaction_{int(time.time())}"
        self.operations: List[Callable] = []
        self.rollback_operations: List[Callable] = []
        self.committed = False
        self.rolled_back = False
    
    def add_operation(self, operation: Callable, rollback: Callable = None):
        """Add an operation with optional rollback."""
        self.operations.append(operation)
        if rollback:
            self.rollback_operations.append(rollback)
    
    def __enter__(self):
        """Enter transaction context."""
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Exit transaction context."""
        if exc_type is not None:
            # Exception occurred, rollback
            self.rollback()
            return False  # Re-raise the exception
        else:
            # Normal completion, commit
            self.commit()
            return True
    
    async def __aenter__(self):
        """Async enter transaction context."""
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Async exit transaction context."""
        if exc_type is not None:
            await self.arollback()
            return False
        else:
            await self.acommit()
            return True
    
    def commit(self):
        """Commit the transaction."""
        if self.rolled_back:
            raise TransactionError("Cannot commit a rolled-back transaction")
        
        try:
            for operation in self.operations:
                operation()
            self.committed = True
            logger.info(f"Transaction {self.name} committed successfully")
        except Exception as e:
            logger.error(f"Failed to commit transaction {self.name}: {e}")
            self.rollback()
            raise
    
    async def acommit(self):
        """Async commit the transaction."""
        if self.rolled_back:
            raise TransactionError("Cannot commit a rolled-back transaction")
        
        try:
            for operation in self.operations:
                if asyncio.iscoroutinefunction(operation):
                    await operation()
                else:
                    operation()
            self.committed = True
            logger.info(f"Transaction {self.name} committed successfully")
        except Exception as e:
            logger.error(f"Failed to commit transaction {self.name}: {e}")
            await self.arollback()
            # Re-raise the original exception instead of wrapping it
            raise
    
    def rollback(self):
        """Rollback the transaction."""
        if self.rolled_back:
            logger.debug(f"Transaction {self.name} already rolled back, skipping")
            return
            
        if self.committed:
            logger.warning(f"Rolling back already committed transaction {self.name}")
        
        # Execute rollback operations in reverse order
        for rollback_op in reversed(self.rollback_operations):
            try:
                rollback_op()
            except Exception as e:
                logger.error(f"Rollback operation failed in transaction {self.name}: {e}")
        
        self.rolled_back = True
        logger.info(f"Transaction {self.name} rolled back")
    
    async def arollback(self):
        """Async rollback the transaction."""
        if self.rolled_back:
            logger.debug(f"Transaction {self.name} already rolled back, skipping")
            return
            
        if self.committed:
            logger.warning(f"Rolling back already committed transaction {self.name}")
        
        # Execute rollback operations in reverse order
        for rollback_op in reversed(self.rollback_operations):
            try:
                if asyncio.iscoroutinefunction(rollback_op):
                    await rollback_op()
                else:
                    rollback_op()
            except Exception as e:
                logger.error(f"Rollback operation failed in transaction {self.name}: {e}")
        
        self.rolled_back = True
        logger.info(f"Transaction {self.name} rolled back")


@contextmanager
def transaction(name: str = None):
    """Context manager for transaction operations."""
    txn = Transaction(name)
    try:
        yield txn
        # If we get here, no exception occurred, so commit
        txn.commit()
    except Exception:
        txn.rollback()
        raise


@asynccontextmanager
async def atransaction(name: str = None):
    """Async context manager for transaction operations."""
    txn = Transaction(name)
    try:
        yield txn
        # If we get here, no exception occurred, so commit
        await txn.acommit()
    except Exception:
        await txn.arollback()
        raise


# Utility functions
def create_retry_config(
    error_type: str = "api",
    **kwargs
) -> RetryConfig:
    """Create retry config for common scenarios."""
    presets = {
        "api": RetryConfig(
            max_attempts=3,
            base_delay=1.0,
            strategy=RetryStrategy.EXPONENTIAL_BACKOFF,
            retryable_exceptions=(APIError, NetworkError, TimeoutError)
        ),
        "training": RetryConfig(
            max_attempts=2,
            base_delay=5.0,
            strategy=RetryStrategy.LINEAR_BACKOFF,
            retryable_exceptions=(TrainingError, MemoryError)
        ),
        "file_io": RetryConfig(
            max_attempts=3,
            base_delay=0.5,
            strategy=RetryStrategy.FIXED_DELAY,
            retryable_exceptions=(OSError, IOError)
        )
    }
    
    config = presets.get(error_type, RetryConfig())
    
    # Override with provided kwargs
    for key, value in kwargs.items():
        if hasattr(config, key):
            setattr(config, key, value)
    
    return config 