src/hf_eda_mcp/error_handling.py

"""
Comprehensive error handling utilities for hf-eda-mcp.

This module provides error handling utilities including retry logic,
error suggestions, and formatted error responses for better user experience.
"""

import logging
import time
import functools
from typing import Optional, Callable, Any, List, Dict, TypeVar, cast
from requests.exceptions import RequestException, ConnectionError, Timeout, HTTPError

logger = logging.getLogger(__name__)

# Type variable for generic function return types
T = TypeVar('T')


class RetryConfig:
    """Configuration for retry logic."""
    
    def __init__(
        self,
        max_attempts: int = 3,
        initial_delay: float = 1.0,
        max_delay: float = 30.0,
        exponential_base: float = 2.0,
        jitter: bool = True
    ):
        """
        Initialize retry configuration.
        
        Args:
            max_attempts: Maximum number of retry attempts
            initial_delay: Initial delay between retries in seconds
            max_delay: Maximum delay between retries in seconds
            exponential_base: Base for exponential backoff
            jitter: Whether to add random jitter to delays
        """
        self.max_attempts = max_attempts
        self.initial_delay = initial_delay
        self.max_delay = max_delay
        self.exponential_base = exponential_base
        self.jitter = jitter


# Default retry configuration
DEFAULT_RETRY_CONFIG = RetryConfig(
    max_attempts=3,
    initial_delay=1.0,
    max_delay=30.0,
    exponential_base=2.0,
    jitter=True
)


def calculate_retry_delay(attempt: int, config: RetryConfig) -> float:
    """
    Calculate delay for retry attempt using exponential backoff.
    
    Args:
        attempt: Current attempt number (0-indexed)
        config: Retry configuration
        
    Returns:
        Delay in seconds
    """
    delay = min(
        config.initial_delay * (config.exponential_base ** attempt),
        config.max_delay
    )
    
    # Add jitter to prevent thundering herd
    if config.jitter:
        import random
        delay = delay * (0.5 + random.random())
    
    return delay


def should_retry_error(error: Exception) -> bool:
    """
    Determine if an error should trigger a retry.
    
    Args:
        error: Exception to check
        
    Returns:
        True if error is retryable, False otherwise
    """
    # Network errors are retryable
    if isinstance(error, (ConnectionError, Timeout)):
        return True
    
    # HTTP errors with specific status codes are retryable
    if isinstance(error, HTTPError):
        # Retry on 5xx server errors and 429 rate limiting
        if hasattr(error, 'response') and error.response is not None:
            status_code = error.response.status_code
            return status_code >= 500 or status_code == 429
    
    # Generic request exceptions might be retryable
    if isinstance(error, RequestException):
        # Check if it's a connection-related issue
        error_str = str(error).lower()
        retryable_keywords = ['timeout', 'connection', 'network', 'temporary']
        return any(keyword in error_str for keyword in retryable_keywords)
    
    # Don't retry other errors by default
    return False


def retry_with_backoff(
    func: Optional[Callable[..., T]] = None,
    *,
    config: Optional[RetryConfig] = None,
    retryable_exceptions: Optional[tuple] = None
) -> Callable[..., T]:
    """
    Decorator to retry a function with exponential backoff.
    
    Args:
        func: Function to decorate (when used without arguments)
        config: Retry configuration (uses default if not provided)
        retryable_exceptions: Tuple of exception types to retry on
        
    Returns:
        Decorated function with retry logic
        
    Example:
        @retry_with_backoff
        def fetch_data():
            # ... network call ...
            pass
            
        @retry_with_backoff(config=RetryConfig(max_attempts=5))
        def fetch_with_custom_config():
            # ... network call ...
            pass
    """
    if config is None:
        config = DEFAULT_RETRY_CONFIG
    
    if retryable_exceptions is None:
        retryable_exceptions = (ConnectionError, Timeout, RequestException)
    
    def decorator(f: Callable[..., T]) -> Callable[..., T]:
        @functools.wraps(f)
        def wrapper(*args: Any, **kwargs: Any) -> T:
            last_exception: Optional[Exception] = None
            
            for attempt in range(config.max_attempts):
                try:
                    return f(*args, **kwargs)
                    
                except retryable_exceptions as e:
                    last_exception = e
                    
                    # Check if we should retry this specific error
                    if not should_retry_error(e):
                        logger.warning(f"Error is not retryable: {e}")
                        raise
                    
                    # Don't sleep after the last attempt
                    if attempt < config.max_attempts - 1:
                        delay = calculate_retry_delay(attempt, config)
                        logger.warning(
                            f"Attempt {attempt + 1}/{config.max_attempts} failed: {e}. "
                            f"Retrying in {delay:.2f}s..."
                        )
                        time.sleep(delay)
                    else:
                        logger.error(
                            f"All {config.max_attempts} attempts failed. Last error: {e}"
                        )
                
                except Exception as e:
                    # Non-retryable exception, raise immediately
                    logger.error(f"Non-retryable error occurred: {e}")
                    raise
            
            # If we get here, all retries failed
            if last_exception:
                raise last_exception
            else:
                raise RuntimeError("Retry logic failed without capturing exception")
        
        return cast(Callable[..., T], wrapper)
    
    # Support both @retry_with_backoff and @retry_with_backoff()
    if func is None:
        return decorator
    else:
        return decorator(func)


def get_dataset_suggestions(dataset_id: str) -> List[str]:
    """
    Generate helpful suggestions for dataset not found errors.
    
    Args:
        dataset_id: The dataset identifier that was not found
        
    Returns:
        List of suggestion strings
    """
    suggestions = []
    
    # Check for common typos or formatting issues
    if " " in dataset_id:
        suggestions.append(
            f"Dataset ID contains spaces. Try: '{dataset_id.replace(' ', '-')}' or '{dataset_id.replace(' ', '_')}'"
        )
    
    if dataset_id.isupper():
        suggestions.append(
            f"Dataset ID is all uppercase. Try lowercase: '{dataset_id.lower()}'"
        )
    
    # Check if it looks like it might be missing organization prefix
    if "/" not in dataset_id:
        suggestions.append(
            f"Dataset might need an organization prefix. Try searching for: 'organization/{dataset_id}'"
        )
    
    # General suggestions
    suggestions.extend([
        "Verify the dataset exists on HuggingFace Hub: https://huggingface.co/datasets",
        f"Search for similar datasets: https://huggingface.co/datasets?search={dataset_id}",
        "Check if the dataset name is spelled correctly",
        "Ensure you have access if the dataset is private or gated"
    ])
    
    return suggestions


def format_authentication_error(
    dataset_id: str,
    is_gated: bool = False,
) -> Dict[str, Any]:
    """
    Format authentication error with helpful guidance.
    
    Args:
        dataset_id: The dataset identifier
        is_gated: Whether the dataset is gated (requires approval)
        
    Returns:
        Dictionary with error details and suggestions
    """
    error_details = {
        "error_type": "authentication_error",
        "dataset_id": dataset_id,
        "is_gated": is_gated,
        "message": "",
        "suggestions": []
    }
    
    if is_gated:
        error_details["message"] = (
            f"Dataset '{dataset_id}' is gated and requires approval to access."
        )
        error_details["suggestions"] = [
            f"Request access to the dataset: https://huggingface.co/datasets/{dataset_id}",
            "Wait for approval from the dataset owner",
            "Provide a valid HuggingFace token after receiving access",
            "Check your HuggingFace account for access status"
        ]
    else:
        error_details["message"] = (
            f"Authentication failed for dataset '{dataset_id}'. "
            "Your token may not have access to this dataset."
        )
        error_details["suggestions"] = [
            "Verify your token is valid and not expired",
            "Check if your token has the required permissions",
            "Ensure you have been granted access to this private dataset",
            "Try regenerating your token at: https://huggingface.co/settings/tokens"
        ]
    
    return error_details


def format_network_error(
    error: Exception,
    operation: str = "operation"
) -> Dict[str, Any]:
    """
    Format network error with helpful guidance.
    
    Args:
        error: The network exception
        operation: Description of the operation that failed
        
    Returns:
        Dictionary with error details and suggestions
    """
    error_details = {
        "error_type": "network_error",
        "operation": operation,
        "message": f"Network error during {operation}: {str(error)}",
        "suggestions": []
    }
    
    # Determine specific error type and provide targeted suggestions
    if isinstance(error, Timeout):
        error_details["error_subtype"] = "timeout"
        error_details["suggestions"] = [
            "The request timed out. Try again in a moment",
            "Check your internet connection",
            "The HuggingFace Hub might be experiencing high load",
            "Try with a smaller sample size or different dataset"
        ]
    elif isinstance(error, ConnectionError):
        error_details["error_subtype"] = "connection"
        error_details["suggestions"] = [
            "Unable to connect to HuggingFace Hub",
            "Check your internet connection",
            "Verify you can access https://huggingface.co",
            "Check if you're behind a firewall or proxy",
            "Try again in a few moments"
        ]
    else:
        error_details["error_subtype"] = "general"
        error_details["suggestions"] = [
            "A network error occurred. Please try again",
            "Check your internet connection",
            "The HuggingFace Hub might be temporarily unavailable",
            "Try again in a few moments"
        ]
    
    return error_details


def format_error_response(
    error: Exception,
    context: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
    """
    Format any error into a structured response with helpful information.
    
    Args:
        error: The exception to format
        context: Optional context information (dataset_id, operation, etc.)
        
    Returns:
        Dictionary with formatted error information
    """
    from hf_eda_mcp.integrations.hf_client import (
        DatasetNotFoundError,
        AuthenticationError,
        NetworkError
    )
    
    context = context or {}
    
    # Handle specific error types
    if isinstance(error, DatasetNotFoundError):
        dataset_id = context.get("dataset_id", "unknown")
        return {
            "error_type": "dataset_not_found",
            "message": str(error),
            "dataset_id": dataset_id,
            "suggestions": get_dataset_suggestions(dataset_id)
        }
    
    elif isinstance(error, AuthenticationError):
        dataset_id = context.get("dataset_id", "unknown")
        is_gated = "gated" in str(error).lower()
        return format_authentication_error(dataset_id, is_gated)
    
    elif isinstance(error, NetworkError):
        operation = context.get("operation", "operation")
        # Extract the original exception if available
        original_error = error.__cause__ or error
        return format_network_error(original_error, operation)
    
    elif isinstance(error, (ConnectionError, Timeout, RequestException)):
        operation = context.get("operation", "operation")
        return format_network_error(error, operation)
    
    elif isinstance(error, ValueError):
        return {
            "error_type": "validation_error",
            "message": str(error),
            "suggestions": [
                "Check that all input parameters are valid",
                "Refer to the tool documentation for parameter requirements"
            ]
        }
    
    else:
        # Generic error
        return {
            "error_type": "unknown_error",
            "message": f"An unexpected error occurred: {str(error)}",
            "error_class": type(error).__name__,
            "suggestions": [
                "Try the operation again",
                "Check the logs for more details",
                "If the problem persists, report it as an issue"
            ]
        }


def log_error_with_context(
    error: Exception,
    context: Optional[Dict[str, Any]] = None,
    level: int = logging.ERROR
) -> None:
    """
    Log an error with contextual information.
    
    Args:
        error: The exception to log
        context: Optional context information
        level: Logging level (default: ERROR)
    """
    context = context or {}
    
    # Build context string
    context_parts = [f"{k}={v}" for k, v in context.items()]
    context_str = ", ".join(context_parts) if context_parts else "no context"
    
    # Log with full details
    logger.log(
        level,
        f"Error occurred: {type(error).__name__}: {str(error)} | Context: {context_str}",
        exc_info=True
    )