import time
from .logger import get_logger
from .exceptions import ChatException, InvalidAPIKeyException, ValidationException

logger = get_logger(f"dracula.{__name__}")

NON_RETRYABLE_EXCEPTIONS = (InvalidAPIKeyException, ValidationException)

RETRYABLE_STATUS_CODES = [
    "429",
    "500",
    "502",
    "503",
    "504",
]


def is_retryable(exception: Exception) -> bool:
    """
    Check if an exception is retryable.

    Args:
        exception (Exception): The exception to check.

    Returns:
        bool: True if the exception is retryable, False otherwise.
    """
    # Never retry these
    if isinstance(exception, NON_RETRYABLE_EXCEPTIONS):
        return False

    # Check if the error message contains a retryable status code
    error_message = str(exception)
    for code in RETRYABLE_STATUS_CODES:
        if code in error_message:
            return True

    # Retry on generic network errors
    if isinstance(exception, (ConnectionError, TimeoutError)):
        return True

    return True


def retry(func, max_retries: int = 3, retry_delay: float = 1.0, *args, **kwargs):
    """
    Execute a function with automatic retry on failure.

    Uses exponential backoff — each retry waits twice as long as the previous.

    Args:
        func: The function to execute.
        max_retries (int): Maximum number of retry attempts. Defaults to 3.
        retry_delay (float): Initial delay in seconds between retries. Defaults to 1.0.
        *args: Positional arguments to pass to the function.
        **kwargs: Keyword arguments to pass to the function.

    Returns:
        Any: The function's return value.

    Raises:
        Exception: The last exception if all retries are exhausted.

    Example:
        >>> result = retry(ai.chat, max_retries=3, retry_delay=1.0, "Hello!")
    """
    last_exception = None

    for attempt in range(max_retries + 1):
        try:
            return func(*args, **kwargs)

        except NON_RETRYABLE_EXCEPTIONS:
            # Don't retry — re-raise immediately
            raise

        except Exception as e:
            last_exception = e

            if attempt == max_retries:
                # All retries exhausted
                logger.error(
                    f"All {max_retries} retries exhausted. " f"Last error: {str(e)}"
                )
                raise ChatException(
                    f"Failed after {max_retries} retries. " f"Last error: {str(e)}"
                )

            if not is_retryable(e):
                raise

            # Calculate exponential backoff delay
            delay = retry_delay * (2**attempt)

            logger.warning(
                f"Attempt {attempt + 1}/{max_retries} failed: {str(e)}. "
                f"Retrying in {delay:.1f} seconds..."
            )

            time.sleep(delay)

    raise last_exception


async def async_retry(
    func, max_retries: int = 3, retry_delay: float = 1.0, *args, **kwargs
):
    """
    Execute an async function with automatic retry on failure.

    Uses exponential backoff — each retry waits twice as long as the previous.

    Args:
        func: The async function to execute.
        max_retries (int): Maximum number of retry attempts. Defaults to 3.
        retry_delay (float): Initial delay in seconds between retries. Defaults to 1.0.
        *args: Positional arguments to pass to the function.
        **kwargs: Keyword arguments to pass to the function.

    Returns:
        Any: The function's return value.

    Raises:
        Exception: The last exception if all retries are exhausted.
    """
    import asyncio

    last_exception = None

    for attempt in range(max_retries + 1):
        try:
            return await func(*args, **kwargs)

        except NON_RETRYABLE_EXCEPTIONS:
            raise

        except Exception as e:
            last_exception = e

            if attempt == max_retries:
                logger.error(
                    f"All {max_retries} retries exhausted. " f"Last error: {str(e)}"
                )
                raise ChatException(
                    f"Failed after {max_retries} retries. " f"Last error: {str(e)}"
                )

            if not is_retryable(e):
                raise

            delay = retry_delay * (2**attempt)

            logger.warning(
                f"Attempt {attempt + 1}/{max_retries} failed: {str(e)}. "
                f"Retrying in {delay:.1f} seconds..."
            )

            await asyncio.sleep(delay)

    raise last_exception
