"""
RetryRunner plugin
##################

RetryRunner plugin implements retry logic to improve task execution reliability.

.. warning:: For grouped tasks need to explicitly provide `connection_name` attribute
    such as `netmiko`, `napalm`, `scrapli`. Specifying `connection_name` attribute for
    standalone tasks not required. Lack of `connection_name` attribute will result in skipping
    connection retry logic and connections to all hosts initiated simultaneously up to the
    number of `num_workers`

Primary motivation for RetryRunner is to make Nornir task execution as reliable as possible
through retry mechanisms.

RetryRunner Architecture
========================

RetryRunner helps to control the rate of connections establishment by limiting the number of
connector workers.

For example, if ``num_connectors`` is 5, it means there are only 5 workers
to establish connections to devices, even if there are 100 devices, RetryRunner will connect
only with 5 of them at a time. This is very helpful when connections rate need to be limited
due to operations restrictions like AAA (TACACS, RADIUS) servers load.

When new task started and if no connection exist to device that this task makes the use of,
RetryRunner attempts to connect to device retrying up to ``connect_retry`` times.

Once connection established, task handed over to worker threads for execution, workers
will retry up to ``task_retry`` if task fails.

Connection parameters such as timeouts or usage of SSH keys handled by Nornir Connection plugins.
RetryRunner calls Nornir to start the connection, further connection establishment details
controlled by Connection plugin itself.

.. image:: ../_images/RetryRunner_v0.png

RetryRunner Sample Usage
========================

Need to instruct Nornir to use RetryRunner on instantiation::

    from nornir import InitNornir

    NornirObj = InitNornir(
        runner={
            "plugin": "RetryRunner",
            "options": {
                "num_workers": 100,
                "num_connectors": 10,
                "connect_retry": 3,
                "connect_backoff": 1000,
                "connect_splay": 100,
                "task_retry": 3,
                "task_backoff": 1000,
                "task_splay": 100
            }
        }
    )

Sample code to demonstrate usage of ``RetryRunner``, ``DictInventory`` and ``ResultSerializer`` plugins::

    import yaml
    import pprint
    from nornir import InitNornir
    from nornir.core.task import Result, Task
    from nornir_netmiko import netmiko_send_command, netmiko_send_config
    from nornir_salt.plugins.functions import ResultSerializer

    inventory_data = '''
    hosts:
      R1:
        hostname: 192.168.1.151
        platform: ios
        groups: [lab]
      R2:
        hostname: 192.168.1.153
        platform: ios
        groups: [lab]
      R3:
        hostname: 192.168.1.154
        platform: ios
        groups: [lab]

    groups:
      lab:
        username: cisco
        password: cisco
    '''

    inventory_dict = yaml.safe_load(inventory_data)

    NornirObj = InitNornir(
        runner={
            "plugin": "RetryRunner",
            "options": {
                "num_workers": 100,
                "num_connectors": 10,
                "connect_retry": 3,
                "connect_backoff": 1000,
                "connect_splay": 100,
                "task_retry": 3,
                "task_backoff": 1000,
                "task_splay": 100
            }
        },
        inventory={
            "plugin": "DictInventory",
            "options": {
                "hosts": inventory_dict["hosts"],
                "groups": inventory_dict["groups"],
                "defaults": inventory_dict.get("defaults", {})
            }
        },
    )

    def _task_group_netmiko_send_commands(task, commands):
        # run commands
        for command in commands:
            task.run(
                task=netmiko_send_command,
                command_string=command,
                name=command
            )
        return Result(host=task.host)

    # run single task
    result1 = NornirObj.run(
        task=netmiko_send_command,
        command_string="show clock"
    )

    # run grouped tasks
    result2 = NornirObj.run(
        task=_task_group_netmiko_send_commands,
        commands=["show clock", "show run | inc hostname"],
        connection_name="netmiko"
    )

    # run another single task
    result3 = NornirObj.run(
        task=netmiko_send_command,
        command_string="show run | inc hostname"
    )

    NornirObj.close_connections()

    # Print results
    formed_result1 = ResultSerializer(result1, add_details=True)
    pprint.pprint(formed_result1, width=100)

    formed_result2 = ResultSerializer(result2, add_details=True)
    pprint.pprint(formed_result2, width=100)

    formed_result3 = ResultSerializer(result3, add_details=True)
    pprint.pprint(formed_result3, width=100)

RetryRunner - Connect to hosts behind jumphost
==============================================

RetryRunner implements logic to connect with hosts behind bastion/jumphosts.

To connect to devices behind jumphost, need to define jumphost parameters in host's inventory data::

    hosts:
      R1:
        hostname: 192.168.1.151
        platform: ios
        username: test
        password: test
        data:
          jumphost:
            hostname: 10.1.1.1
            port: 22
            password: jump_host_password
            username: jump_host_user

.. note:: Only Netmiko ``connection_name="netmiko"`` and Ncclient ``connection_name="ncclient"``
    tasks, support connecting to hosts behind Jumphosts using above inventory data.

RetryRunner Reference
=====================

.. autoclass:: nornir_salt.plugins.runners.RetryRunner.RetryRunner
"""
import threading
import queue
import logging
import time
import random
from typing import List
from nornir.core.task import AggregatedResult, Task, MultiResult, Result
from nornir.core.inventory import Host

try:
    import paramiko

    HAS_PARAMIKO = True
except ImportError:
    HAS_PARAMIKO = False

log = logging.getLogger(__name__)
LOCK = threading.Lock()


def worker(
    stop_event,
    connectors_q,
    work_q,
    task_backoff: int,
    task_splay: int,
    task_retry: int,
    connect_retry: int,
    reconnect_on_fail: bool,
):
    while not stop_event.is_set():
        try:
            work = work_q.get(block=True, timeout=0.1)
        except queue.Empty:
            continue
        task, host, params, result = work
        # check if has exception recorded in params by connector thread
        if params.get("connect_exception"):
            work_result = MultiResult(task.name)
            work_result.append(
                Result(
                    host,
                    result=params["connect_exception"],
                    failed=True,
                    exception=params["connect_exception"],
                    name=task.name,
                )
            )
        else:
            # check if need backoff task retry for this host
            if params["task_retry"] > 0:
                elapsed = time.time() - params["timestamp"]
                should_wait = (params["task_retry"] * task_backoff) / 1000
                if elapsed < should_wait:
                    work_q.put(work)
                    work_q.task_done()
                    continue
            log.info(
                "nornir_salt:RetryRunner {} - running task '{}'".format(
                    host.name, task.name
                )
            )
            time.sleep(random.randrange(0, task_splay) / 1000)  # nosec
            work_result = task.start(host)
        # check if need to run task retry logic
        if task.results.failed:
            log.error(
                "nornir_salt:RetryRunner {} - task execution retry attempt {} failed: '{}'".format(
                    host.name, params["task_retry"], work_result[0].exception
                )
            )
            if (
                params["task_retry"] < task_retry
                and params["connection_retry"] < connect_retry
            ):
                params["task_retry"] += 1
                params["timestamp"] = time.time()
                # recover task results for not to count task as failed
                for r in task.results:
                    r.failed = False
                    r.exception = None
                    r.skip_results = (
                        True  # for ResultSerializer to skip failed attempts
                    )
                if reconnect_on_fail:
                    # close host connections to retry them
                    close_host_connection(host, params["connection_name"])
                    connectors_q.put(work)
                    params["connection_retry"] += 1
                else:
                    work_q.put(work)
                work_q.task_done()
                continue
        # enreach result objects with runner statistics
        for result_item in work_result:
            result_item.connection_retry = params["connection_retry"]
            result_item.task_retry = params["task_retry"]
        with LOCK:
            result[host.name] = work_result
            del work_result
        work_q.task_done()
        log.info(
            "nornir_salt:RetryRunner {} - task '{}' completed".format(
                host.name, task.name
            )
        )


def connector(
    stop_event,
    connectors_q,
    work_q,
    connect_backoff: int,
    connect_splay: int,
    connect_retry: int,
    jumphosts_connections,
):
    while not stop_event.is_set():
        try:
            connection = connectors_q.get(block=True, timeout=0.1)
        except queue.Empty:
            continue
        task, host, params, result = connection
        # check if need back off connection retry for this host
        if params["connection_retry"] > 0:
            elapsed = time.time() - params["timestamp"]
            should_wait = (params["connection_retry"] * connect_backoff) / 1000
            if elapsed < should_wait:
                connectors_q.put(connection)
                connectors_q.task_done()
                continue
        # initiate connection to host
        connection_name = task.task.__globals__.get("CONNECTION_NAME", None)
        connection_name = task.params.pop("connection_name", connection_name)
        # on connect retry get connection name from params
        connection_name = params.get("connection_name", connection_name)
        params.setdefault("connection_name", connection_name)
        if connection_name and connection_name not in host.connections:
            try:
                time.sleep(random.randrange(0, connect_splay) / 1000)  # nosec
                if host.get("jumphost") and connection_name in ["netmiko", "ncclient"]:
                    extras = host.get_connection_parameters(connection_name).extras
                    extras["sock"] = connect_to_device_behind_jumphost(
                        host, jumphosts_connections
                    )
                    host.open_connection(
                        connection_name, configuration=task.nornir.config, extras=extras
                    )
                else:
                    host.open_connection(
                        connection_name, configuration=task.nornir.config
                    )
                log.info(
                    "nornir_salt:RetryRunner {} - started connection: '{}'".format(
                        host.name, connection_name
                    )
                )
            except Exception as e:
                # close host connections to retry them
                close_host_connection(host, connection_name)
                err_msg = "nornir_salt:RetryRunner {} - connection retry attempt {}, error: '{}'".format(
                    host.name, params["connection_retry"], e
                )
                log.error(err_msg)
                if params["connection_retry"] < connect_retry:
                    params["connection_retry"] += 1
                    params["timestamp"] = time.time()
                    connectors_q.put(connection)
                    connectors_q.task_done()
                    continue
                else:
                    # record exception in params for worker thread to react on it
                    params["connect_exception"] = err_msg
        connectors_q.task_done()
        work_q.put((task, host, params, result))


def close_host_connection(host, connection_name):
    try:
        host.close_connection(connection_name)
    except:
        _ = host.connections.pop(connection_name, None)
    if host.get("jumphost"):
        channel_name = "jumphost_{}_channel".format(host["jumphost"]["hostname"])
        try:
            host.close_connection(channel_name)
        except:
            _ = host.connections.pop(channel_name, None)


def connect_to_device_behind_jumphost(host, jumphosts_connections):
    """
    Establish connection to devices behind jumphost/bastion
    """

    jumphost = {"timeout": 3, "look_for_keys": False, "allow_agent": False}
    jumphost.update(host["jumphost"])
    # Initiate connection to jumphost if not initiated already or
    # if it failed before
    #
    # jumphosts_connections dictionary shared between threads,
    # first thread to start will detect that no connection to
    # jumphost exists and will try to establish it, while doing
    # so, jumphost key will be added to jumphosts_connections
    # dictionary with '__connecting__' value, next threads
    # will see '__connecting__' as value for jumphost and will
    # continue sleeping until either connection succeeded or
    # failed. On connection failure, first thread will set value
    # to '__failed__' to signal other threads to exit.
    if (
        jumphost["hostname"] not in jumphosts_connections
        or jumphosts_connections.get(jumphost["hostname"]) == "__failed__"
    ):
        try:
            with LOCK:
                jumphosts_connections[jumphost["hostname"]] = "__connecting__"
            jumphost_ssh_client = paramiko.client.SSHClient()
            jumphost_ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            jumphost_ssh_client.connect(**jumphost)
            with LOCK:
                jumphosts_connections[jumphost["hostname"]] = {
                    "jumphost_ssh_client": jumphost_ssh_client,
                    "jumphost_ssh_transport": jumphost_ssh_client.get_transport(),
                }
            # add jumphost to host connections to close it
            # on nornir.close_connections() call
            host.connections[
                "jumphost_{}".format(jumphost["hostname"])
            ] = jumphost_ssh_client
            log.info(
                "nornir_salt:RetryRunner Started connection to jumphost '{}' - '{}'".format(
                    jumphost["hostname"], jumphost_ssh_client
                )
            )
        except Exception as e:
            with LOCK:
                jumphosts_connections[jumphost["hostname"]] = "__failed__"
            # add exception info to host data to include in results
            error_msg = "nornir_salt:RetryRunner failed connection to jumphost '{}', error - {}".format(
                host["jumphost"]["hostname"], e
            )
            # host["exception"] = error_msg
            log.error(error_msg)
            raise RuntimeError(error_msg)
    else:
        # sleep random time waiting for connection to jumphost to establish
        while jumphosts_connections[jumphost["hostname"]] == "__connecting__":
            time.sleep(random.randrange(0, 500) / 1000)  # nosec
        if jumphosts_connections[jumphost["hostname"]] == "__failed__":
            # add exception info to host data to include in results
            error_msg = "nornir_salt:RetryRunner failed connection to jumphost '{}' in another thread".format(
                host["jumphost"]["hostname"]
            )
            # host["exception"] = error_msg
            log.error(error_msg)
            raise RuntimeError(error_msg)
    # connect to host
    channel_name = "jumphost_{}_channel".format(jumphost["hostname"])
    jconn = jumphosts_connections[jumphost["hostname"]]
    # check if connection to jumphost was closed, reconnect if so
    if not jconn["jumphost_ssh_transport"].is_active():
        log.debug(
            "nornir_salt:RetryRunner jumphost '{}' disconnected, reconnecting".format(
                jumphost["hostname"]
            )
        )
        jconn["jumphost_ssh_client"].close()
        jumphosts_connections.pop(jumphost["hostname"])
        return connect_to_device_behind_jumphost(host, jumphosts_connections)
    # proceed with openning new channel via jumphost
    if not host.connections.get(channel_name):
        host.connections[channel_name] = jconn["jumphost_ssh_transport"].open_channel(
            kind="direct-tcpip",
            dest_addr=(host.hostname, host.port or 22),
            src_addr=("localhost", 7777),
            timeout=3,
        )
        log.info(
            "nornir_salt:RetryRunner {} - started new channel via jumphost '{}' - '{}'".format(
                host.name, jumphost["hostname"], jconn["jumphost_ssh_client"]
            )
        )
    return host.connections[channel_name]


class RetryRunner:
    """
    RetryRunner is a modification of QueueRunner that strives to
    make task execution as reliable as possible.

    Arguments:
        num_workers: number of threads for tasks execution
        num_connectors: number of threads for device connections
        connect_retry: number of connection attempts
        connect_backoff: exponential backoff timer in milliseconds
        connect_splay: random interval between 0 and splay for each connection in milliseconds
        task_retry: number of attempts to run task
        task_backoff: exponential backoff timer in milliseconds
        task_splay: random interval between 0 and splay before task start in milliseconds
        reconnect_on_fail: boolean, default True, perform reconnect to host on task failure
        task_timeout: int, seconds to wait for task to complete before closing all queues
            and stopping connectors and workers threads, default 600
    """

    def __init__(
        self,
        num_workers: int = 100,
        num_connectors: int = 20,
        connect_retry: int = 3,
        connect_backoff: int = 5000,
        connect_splay: int = 100,
        task_retry: int = 1,
        task_backoff: int = 5000,
        task_splay: int = 100,
        reconnect_on_fail: bool = True,
        task_timeout: int = 600,
    ) -> None:
        self.num_workers = num_workers
        self.num_connectors = num_connectors
        self.connect_retry = connect_retry
        self.connect_backoff = connect_backoff
        self.connect_splay = connect_splay
        self.task_retry = task_retry
        self.task_backoff = task_backoff
        self.task_splay = task_splay
        self.jumphosts_connections = {}
        self.reconnect_on_fail = reconnect_on_fail
        self.task_timeout = task_timeout

    def run(self, task: Task, hosts: List[Host]) -> AggregatedResult:
        connectors_q = queue.Queue()
        work_q = queue.Queue()
        stop_event = threading.Event()
        task_timeout_event = threading.Event()
        task_timeout_timer = threading.Timer(
            self.task_timeout, lambda: task_timeout_event.set()
        )
        result = AggregatedResult(task.name)
        connector_threads = []
        worker_threads = []
        # enqueue hosts in connectors queue
        for host in hosts:
            connectors_q.put(
                (task.copy(), host, {"connection_retry": 0, "task_retry": 0}, result)
            )
        # start connectors threads
        for _ in range(self.num_connectors):
            t = threading.Thread(
                target=connector,
                args=(
                    stop_event,
                    connectors_q,
                    work_q,
                    self.connect_backoff,
                    self.connect_splay,
                    self.connect_retry,
                    self.jumphosts_connections,
                ),
            )
            t.start()
            connector_threads.append(t)
        # start worker threads
        for _ in range(self.num_workers):
            t = threading.Thread(
                target=worker,
                args=(
                    stop_event,
                    connectors_q,
                    work_q,
                    self.task_backoff,
                    self.task_splay,
                    self.task_retry,
                    self.connect_retry,
                    self.reconnect_on_fail,
                ),
            )
            t.start()
            worker_threads.append(t)
        # wait until all hosts completed task or timeout reached
        task_timeout_timer.start()
        while not task_timeout_event.is_set():
            with LOCK:
                hosts_no_result = [h.name for h in hosts if h.name not in result]
            if hosts_no_result == []:
                task_timeout_timer.cancel()
                break
            time.sleep(0.1)
        else:
            log.warning(
                "nornir_salt:RetryRunner task '{}', '{}'s task_timeout reached, hosts no results '{}'".format(
                    task.name, self.task_timeout, hosts_no_result
                )
            )
        # block until all queues empty
        connectors_q.join()
        work_q.join()
        # stop and delete connectors and workers threads
        stop_event.set()
        while connector_threads:
            _ = connector_threads.pop().join()
        while worker_threads:
            _ = worker_threads.pop().join()
        return result
