"""
Elasticsearch Retry Utilities.

Provides retry with exponential backoff and jitter for Elasticsearch operations.
"""

from __future__ import annotations

import asyncio
import random
import time
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from functools import wraps
from typing import (
    Any,
    TypeVar,
)

from smf.smf_logging import get_logger

logger = get_logger("elasticsearch.retry")

# Type variables for generic decorators
T = TypeVar("T")
F = TypeVar("F", bound=Callable[..., Any])

# Default retryable exceptions (common transient errors)
DEFAULT_RETRYABLE_EXCEPTIONS: tuple[type[Exception], ...] = (
    ConnectionError,
    TimeoutError,
    OSError,
)

# Default retryable HTTP status codes from Elasticsearch
DEFAULT_RETRYABLE_STATUS_CODES: tuple[int, ...] = (
    429,  # Too Many Requests
    502,  # Bad Gateway
    503,  # Service Unavailable
    504,  # Gateway Timeout
)


@dataclass
class RetryConfig:
    """Configuration for retry behavior.

    Attributes:
        max_retries: Maximum number of retry attempts (default: 3).
        initial_delay: Initial delay in seconds before first retry (default: 0.1).
        max_delay: Maximum delay between retries in seconds (default: 60.0).
        exponential_base: Base for exponential backoff calculation (default: 2.0).
        jitter: If True, add random jitter to delays (default: True).
        retryable_exceptions: Exception types that should trigger retry.
        retryable_status_codes: HTTP status codes that should trigger retry.
    """

    max_retries: int = 3
    initial_delay: float = 0.1
    max_delay: float = 60.0
    exponential_base: float = 2.0
    jitter: bool = True
    retryable_exceptions: tuple[type[Exception], ...] = field(
        default_factory=lambda: DEFAULT_RETRYABLE_EXCEPTIONS
    )
    retryable_status_codes: tuple[int, ...] = field(
        default_factory=lambda: DEFAULT_RETRYABLE_STATUS_CODES
    )

    def calculate_delay(self, attempt: int) -> float:
        """Calculate delay for given attempt number (0-indexed).

        Args:
            attempt: The attempt number (0 for first retry).

        Returns:
            Delay in seconds before the next retry.
        """
        # Exponential backoff: initial * base^attempt
        delay = self.initial_delay * (self.exponential_base**attempt)

        # Cap at maximum delay
        delay = min(delay, self.max_delay)

        # Add jitter if enabled (+/- 25%)
        if self.jitter:
            jitter_range = delay * 0.25
            delay = delay + random.uniform(-jitter_range, jitter_range)

        return max(0, delay)


def is_retryable_exception(
    exc: Exception,
    config: RetryConfig,
) -> bool:
    """Check if an exception should trigger a retry.

    Args:
        exc: The exception to check.
        config: Retry configuration.

    Returns:
        True if the exception is retryable.
    """
    # Check if it's a known retryable exception type
    if isinstance(exc, config.retryable_exceptions):
        return True

    # Check for Elasticsearch-specific transport errors
    try:
        from elasticsearch import TransportError

        if isinstance(exc, TransportError):
            status_code = getattr(exc, "status_code", None)
            if status_code in config.retryable_status_codes:
                return True
    except ImportError:
        pass

    # Check for connection-related errors by message
    error_msg = str(exc).lower()
    connection_keywords = ["connection", "timeout", "refused", "reset", "broken pipe"]
    return any(keyword in error_msg for keyword in connection_keywords)


def with_retry(
    config: RetryConfig | None = None,
) -> Callable[[Callable[..., T]], Callable[..., T]]:
    """Decorator for sync functions with retry logic.

    Adds exponential backoff retry behavior to a synchronous function.

    Args:
        config: Retry configuration (uses defaults if None).

    Returns:
        Decorated function with retry behavior.

    Example:
        >>> @with_retry(RetryConfig(max_retries=5))
        ... def search_documents(query: str) -> dict:
        ...     return client.search(query=query)
    """
    if config is None:
        config = RetryConfig()

    def decorator(func: Callable[..., T]) -> Callable[..., T]:
        @wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> T:
            last_exception: Exception | None = None

            for attempt in range(config.max_retries + 1):
                try:
                    return func(*args, **kwargs)
                except Exception as exc:
                    last_exception = exc

                    if not is_retryable_exception(exc, config):
                        # Not retryable, re-raise immediately
                        raise

                    if attempt == config.max_retries:
                        # Max retries exceeded
                        logger.error(
                            f"Max retries ({config.max_retries}) exceeded for {func.__name__}",
                            function=func.__name__,
                            attempt=attempt + 1,
                            error=str(exc),
                            error_type=type(exc).__name__,
                        )
                        raise

                    delay = config.calculate_delay(attempt)
                    logger.warning(
                        f"Retry {attempt + 1}/{config.max_retries} for {func.__name__}",
                        function=func.__name__,
                        attempt=attempt + 1,
                        delay_seconds=round(delay, 3),
                        error=str(exc),
                        error_type=type(exc).__name__,
                    )
                    time.sleep(delay)

            # Should not reach here, but satisfy type checker
            if last_exception:
                raise last_exception
            raise RuntimeError("Unexpected retry loop exit")

        return wrapper

    return decorator


def with_retry_async(
    config: RetryConfig | None = None,
) -> Callable[[Callable[..., Awaitable[T]]], Callable[..., Awaitable[T]]]:
    """Decorator for async functions with retry logic.

    Adds exponential backoff retry behavior to an asynchronous function.

    Args:
        config: Retry configuration (uses defaults if None).

    Returns:
        Decorated async function with retry behavior.

    Example:
        >>> @with_retry_async(RetryConfig(max_retries=5))
        ... async def search_documents(query: str) -> dict:
        ...     return await client.search(query=query)
    """
    if config is None:
        config = RetryConfig()

    def decorator(func: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
        @wraps(func)
        async def wrapper(*args: Any, **kwargs: Any) -> T:
            last_exception: Exception | None = None

            for attempt in range(config.max_retries + 1):
                try:
                    return await func(*args, **kwargs)
                except Exception as exc:
                    last_exception = exc

                    if not is_retryable_exception(exc, config):
                        raise

                    if attempt == config.max_retries:
                        logger.error(
                            f"Max retries ({config.max_retries}) exceeded for {func.__name__}",
                            function=func.__name__,
                            attempt=attempt + 1,
                            error=str(exc),
                            error_type=type(exc).__name__,
                        )
                        raise

                    delay = config.calculate_delay(attempt)
                    logger.warning(
                        f"Retry {attempt + 1}/{config.max_retries} for {func.__name__}",
                        function=func.__name__,
                        attempt=attempt + 1,
                        delay_seconds=round(delay, 3),
                        error=str(exc),
                        error_type=type(exc).__name__,
                    )
                    await asyncio.sleep(delay)

            if last_exception:
                raise last_exception
            raise RuntimeError("Unexpected retry loop exit")

        return wrapper

    return decorator


class RetryableOperation:
    """Context manager for retryable operations.

    Provides an alternative to decorators for cases where you need
    more control over the retry logic.

    Example:
        >>> config = RetryConfig(max_retries=3)
        >>> with RetryableOperation(config) as retry:
        ...     for attempt in retry:
        ...         try:
        ...             result = client.search(query=query)
        ...             break
        ...         except Exception as e:
        ...             retry.record_failure(e)
    """

    def __init__(self, config: RetryConfig | None = None) -> None:
        self.config = config or RetryConfig()
        self._attempt = 0
        self._last_exception: Exception | None = None

    def __enter__(self) -> RetryableOperation:
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: Any,
    ) -> bool:
        # Don't suppress exceptions
        return False

    def __iter__(self) -> RetryableOperation:
        self._attempt = 0
        return self

    def __next__(self) -> int:
        if self._attempt > self.config.max_retries:
            if self._last_exception:
                raise self._last_exception
            raise StopIteration

        current = self._attempt
        self._attempt += 1
        return current

    def record_failure(self, exc: Exception) -> None:
        """Record a failure and sleep if retryable.

        Args:
            exc: The exception that occurred.

        Raises:
            The exception if not retryable or max retries exceeded.
        """
        self._last_exception = exc

        if not is_retryable_exception(exc, self.config):
            raise exc

        if self._attempt > self.config.max_retries:
            raise exc

        delay = self.config.calculate_delay(self._attempt - 1)
        logger.warning(
            f"Retry {self._attempt}/{self.config.max_retries}",
            attempt=self._attempt,
            delay_seconds=round(delay, 3),
            error=str(exc),
        )
        time.sleep(delay)


__all__ = [
    "DEFAULT_RETRYABLE_EXCEPTIONS",
    "DEFAULT_RETRYABLE_STATUS_CODES",
    "RetryConfig",
    "RetryableOperation",
    "is_retryable_exception",
    "with_retry",
    "with_retry_async",
]
