IMMREX7

# -*- coding: utf-8 -*-

import logging
from functools import partial, update_wrapper
from time import sleep
from typing import Any, Dict, Iterable, Optional, Tuple, Type, Callable


MAX_RETRIES = 3
LOGGER = logging.getLogger(__name__)


class RetriesExhaustedError(Exception):
    """A special type which signals the failure of a retry loop."""


def abbreviate_hostname_for_windows(hostname: Optional[str]) -> Optional[str]:
    """Abbreviate hostname for use on a Windows machine.

    :param hostname: the hostname
    :returns: the first non-empty domain in the hostname, excluding "www."
    """
    if hostname is None:
        return None

    if hostname.lower().startswith('www.'):
        hostname = hostname[4:]

    for domain in hostname.split('.'):
        if domain:
            return domain

    return hostname


def subdict(dict_: Dict[Any, Any], keys: Iterable[Any]) -> Dict[Any, Any]:
    """Filter a dictionary to contain only a certain set of keys.

    :param dict_: The original dictionary to be filtered.
    :param keys: A list, or other iterable, containing the desired dictionary keys.
    :returns: A dictionary containing only the desired keys.
    """
    return {k: v for k, v in dict_.items() if k in keys}


def subdict_omit(dict_: Dict[Any, Any], keys: Iterable[Any]) -> Dict[Any, Any]:
    """Filter a dictionary to omit a set of keys.

    :param dict_: The original dictionary to be filtered.
    :param keys: An iterable containing the keys to omit.
    :returns: A dictionary with the desired keys omitted.
    """
    return {k: v for k, v in dict_.items() if k not in keys}


def _should_retry(*_args, curr_attempt: int = 0,
                  max_attempts: int = MAX_RETRIES, **_kwargs) -> bool:
    return curr_attempt < max_attempts


def _retry_after(*_args, curr_attempt: int = 0, sleep_secs: int = 1, **_kwargs) -> int:
    return sleep_secs


def retry_this(on_ex_classes: Tuple[Type[Exception], ...] = (Exception,),
               max_attempts: int = MAX_RETRIES,
               sleep_secs: int = 1,
               should_retry: Optional[Callable[[Any], bool]] = None,
               retry_after: Optional[Callable[[Any], int]] = None):
    """Decorator that adds retry on error functionality to a function.
    Currently the retry strategy is 'linear' on errors. i.e. this function waits a set period
    of time before retrying the failed function again.

    :param on_ex_classes: A tuple of exceptions to retry on.
                          By default, its all exceptions that derive from the 'Exception' class.
    :param max_attempts: Limit to how many times we'll retry
    :param sleep_secs: How long to wait between retries.
    :param should_retry: A predicate which when called will return a boolean saying whether
                        the call should be retried or not. This parameter overrides the max_attempts
                        parameter and gives more control to dynamically choose on if we need to
                        continue retrying a call.
    :param retry_after: A callable that returns how long to wait between retries.
                        This parameter overrides the sleep_secs parameter and gives more control
                        to dynamically choose the wait time.
    :returns: This returns a decorator function that actually provides the retry functionality.
    """
    should_retry = should_retry or partial(_should_retry, max_attempts=max_attempts)
    retry_after = retry_after or partial(_retry_after, sleep_secs=sleep_secs)

    def wrapper(f):
        ex_classes = tuple(ex_cls for ex_cls in on_ex_classes if issubclass(ex_cls, Exception))

        def new_func(*pargs, **kwargs):
            curr_attempt = 0
            while True:
                try:
                    return f(*pargs, **kwargs)
                except ex_classes as e:  # pylint: disable=E0712
                    curr_attempt += 1
                    LOGGER.error("Exception (%s) occured while executing %s", str(e), f.__name__)

                    if not should_retry(*pargs, curr_attempt=curr_attempt, **kwargs):
                        msg = 'Max attempts exhausted for {}'.format(f.__name__)
                        # pylint: disable=E0703
                        raise RetriesExhaustedError(msg) from e

                    s_secs = retry_after(*pargs, curr_attempt=curr_attempt, **kwargs)
                    LOGGER.debug(
                        "Sleeping %s secs before retrying %s, due to exception (%s)",
                        s_secs, f.__name__, str(e))
                    sleep(s_secs)

        return update_wrapper(new_func, f)
    return wrapper