"""
Retry utilities for robust operation handling.

Implements:
- Exponential backoff with jitter
- Configurable retry strategies
- Error classification for intelligent retries
"""

import functools
import random
import time
from typing import Any, Callable, Optional, Tuple, Type

from ..infra.logging import get_logger

logger = get_logger(__name__)


# Errors that are typically transient and worth retrying
TRANSIENT_ERRORS = (
    TimeoutError,
    ConnectionError,
    RuntimeError,  # Includes tokenization errors
    OSError,
)


def is_retryable_error(error: Exception) -> bool:
    """
    Check if an error is retryable.
    
    Args:
        error: The exception to check
        
    Returns:
        True if the error should be retried
    """
    error_msg = str(error).lower()
    
    # Tokenization/embedding errors that can be retried with sanitization
    retryable_patterns = [
        "negative int to unsigned",
        "overflow",
        "memory",
        "timeout",
        "connection",
        "temporary",
        "resource",
    ]
    
    if any(pattern in error_msg for pattern in retryable_patterns):
        return True
    
    # Check error type
    if isinstance(error, TRANSIENT_ERRORS):
        return True
    
    return False


def retry_with_backoff(
    max_retries: int = 3,
    initial_delay: float = 0.5,
    max_delay: float = 30.0,
    exponential_base: float = 2.0,
    jitter: bool = True,
    retryable_exceptions: Optional[Tuple[Type[Exception], ...]] = None,
    on_retry: Optional[Callable[[Exception, int], None]] = None
):
    """
    Decorator for retrying functions with exponential backoff.
    
    Args:
        max_retries: Maximum number of retry attempts
        initial_delay: Initial delay in seconds
        max_delay: Maximum delay in seconds
        exponential_base: Base for exponential backoff
        jitter: Add randomness to delay
        retryable_exceptions: Tuple of exception types to retry
        on_retry: Callback function called on each retry (error, attempt)
        
    Returns:
        Decorated function
        
    Example:
        @retry_with_backoff(max_retries=3)
        def call_api():
            ...
    """
    retryable_exceptions = retryable_exceptions or TRANSIENT_ERRORS
    
    def decorator(func: Callable) -> Callable:
        @functools.wraps(func)
        def wrapper(*args, **kwargs) -> Any:
            last_exception = None
            
            for attempt in range(max_retries + 1):
                try:
                    return func(*args, **kwargs)
                except retryable_exceptions as e:
                    last_exception = e
                    
                    if attempt == max_retries:
                        logger.error(
                            f"Function {func.__name__} failed after {max_retries + 1} attempts: {e}"
                        )
                        raise
                    
                    # Calculate delay with exponential backoff
                    delay = min(
                        initial_delay * (exponential_base ** attempt),
                        max_delay
                    )
                    
                    if jitter:
                        delay = delay * (0.5 + random.random())
                    
                    logger.warning(
                        f"Attempt {attempt + 1}/{max_retries + 1} for {func.__name__} failed: {e}. "
                        f"Retrying in {delay:.2f}s..."
                    )
                    
                    if on_retry:
                        on_retry(e, attempt)
                    
                    time.sleep(delay)
                except Exception as e:
                    # Check if this error is retryable based on message
                    if is_retryable_error(e) and attempt < max_retries:
                        last_exception = e
                        delay = min(
                            initial_delay * (exponential_base ** attempt),
                            max_delay
                        )
                        if jitter:
                            delay = delay * (0.5 + random.random())
                        
                        logger.warning(
                            f"Retryable error in {func.__name__}: {e}. "
                            f"Retrying in {delay:.2f}s..."
                        )
                        time.sleep(delay)
                    else:
                        raise
            
            # Should not reach here, but just in case
            if last_exception:
                raise last_exception
                
        return wrapper
    return decorator


class RetryContext:
    """
    Context manager for retry operations with state.
    
    Allows for more complex retry logic where state needs to be
    maintained between attempts.
    
    Example:
        with RetryContext(max_retries=3) as ctx:
            while ctx.should_retry():
                try:
                    result = do_something()
                    ctx.success()
                except Exception as e:
                    ctx.record_error(e)
    """
    
    def __init__(
        self,
        max_retries: int = 3,
        initial_delay: float = 0.5,
        max_delay: float = 30.0,
        exponential_base: float = 2.0,
        jitter: bool = True
    ):
        self.max_retries = max_retries
        self.initial_delay = initial_delay
        self.max_delay = max_delay
        self.exponential_base = exponential_base
        self.jitter = jitter
        
        self.attempt = 0
        self.last_error: Optional[Exception] = None
        self._succeeded = False
    
    def __enter__(self) -> "RetryContext":
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        pass
    
    def should_retry(self) -> bool:
        """Check if we should attempt (another) retry."""
        if self._succeeded:
            return False
        if self.attempt > self.max_retries:
            if self.last_error:
                raise self.last_error
            return False
        return True
    
    def record_error(self, error: Exception) -> None:
        """Record an error and prepare for retry."""
        self.last_error = error
        self.attempt += 1
        
        if self.attempt <= self.max_retries:
            delay = min(
                self.initial_delay * (self.exponential_base ** (self.attempt - 1)),
                self.max_delay
            )
            if self.jitter:
                delay = delay * (0.5 + random.random())
            
            logger.warning(
                f"Attempt {self.attempt}/{self.max_retries + 1} failed: {error}. "
                f"Retrying in {delay:.2f}s..."
            )
            time.sleep(delay)
    
    def success(self) -> None:
        """Mark the operation as successful."""
        self._succeeded = True


def with_fallback(
    primary_fn: Callable,
    fallback_fn: Callable,
    error_types: Tuple[Type[Exception], ...] = (Exception,)
) -> Callable:
    """
    Create a function that falls back to an alternative on error.
    
    Args:
        primary_fn: Primary function to try
        fallback_fn: Fallback function if primary fails
        error_types: Exception types that trigger fallback
        
    Returns:
        Wrapped function that tries primary then fallback
    """
    @functools.wraps(primary_fn)
    def wrapper(*args, **kwargs):
        try:
            return primary_fn(*args, **kwargs)
        except error_types as e:
            logger.warning(f"Primary function failed: {e}. Using fallback.")
            return fallback_fn(*args, **kwargs)
    
    return wrapper
