# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
"""This module contains the manager for handler methods used by the callback client"""
import logging
import threading
import abc
import six
from azure.iot.device.common import handle_exceptions
from azure.iot.device.common.chainable_exception import ChainableException
from azure.iot.device.iothub.client_event import (
    CONNECTION_STATE_CHANGE,
    NEW_SASTOKEN_REQUIRED,
    BACKGROUND_EXCEPTION,
)
import concurrent.futures

logger = logging.getLogger(__name__)

# Receiver Handlers
MESSAGE = "_on_message_received"
METHOD = "_on_method_request_received"
TWIN_DP_PATCH = "_on_twin_desired_properties_patch_received"

# Client Event Handler Runner
CLIENT_EVENT = "client_event"
# Client Event Names
client_events = [CONNECTION_STATE_CHANGE, NEW_SASTOKEN_REQUIRED, BACKGROUND_EXCEPTION]


class HandlerManagerException(ChainableException):
    """An exception raised by a HandlerManager"""

    pass


class HandlerRunnerKillerSentinel(object):
    """An object that functions according to the sentinel design pattern.
    Insert into an Inbox in order to indicate that the Handler Runner associated with that
    Inbox should be stopped.
    """

    pass


@six.add_metaclass(abc.ABCMeta)
class AbstractHandlerManager(object):
    """Partial class that defines handler manager functionality shared between sync/async"""

    def __init__(self, inbox_manager):
        self._inbox_manager = inbox_manager

        self._receiver_handler_runners = {
            MESSAGE: None,
            METHOD: None,
            TWIN_DP_PATCH: None,
        }
        self._client_event_runner = None

        # Receiver handlers (Each will have it's own runner)
        self._on_message_received = None
        self._on_method_request_received = None
        self._on_twin_desired_properties_patch_received = None

        # Client Event handlers (Share a single Client Event runner)
        self._on_connection_state_change = None
        self._on_new_sastoken_required = None
        self._on_background_exception = None

        # As mentioned above, Receiver handlers each get their own runner. This is because it is
        # reasonably possible that many receives (and many different types of receives) can be
        # happening in a very short time. This is because we want to be able to support processing
        # multiple receives simultaneously, but also protect different receive invocations from
        # being slowed by a poorly written handler of a different type. Thus each Receiver handler
        # gets its own unique runner that cannot end up blocked by other handlers.
        #
        # Client Events on the other hand are generated by the client itself rather than
        # unsolicited data received over the wire. As these are relatively uncommon, we don't
        # expect multiple Client Events of the same type to really start queueing up, so there
        # isn't much justification for each type of event getting it's own dedicated runner.
        # Furthermore, there is much less concern over one inefficient or slow handler blocking
        # execution of others due to this infrequency.
        #
        # However, there are other differences between these handler classes as well. Receiver
        # handlers will always be invoked with a single argument - the received data structure.
        # Client Event handlers on the other hand are more flexible - they may be invoked with
        # different numbers of arguments depending on the handler, from none to multiple. This
        # is to keep design space open for more complex Client Events in the future.
        #
        # Finally, it is possible to stop ONLY the receiver handlers via the .stop() method
        # if desired. This is useful because while receiver handlers should be stopped when the
        # client disconnects, Client Events can still occur while the client is disconnected
        # as they are propagated from the client itself rather than received data.

    def _get_inbox_for_receive_handler(self, handler_name):
        """Retrieve the inbox relevant to the handler"""
        if handler_name == METHOD:
            return self._inbox_manager.get_method_request_inbox()
        elif handler_name == TWIN_DP_PATCH:
            return self._inbox_manager.get_twin_patch_inbox()
        elif handler_name == MESSAGE:
            return self._inbox_manager.get_unified_message_inbox()
        else:
            return None

    def _get_handler_for_client_event(self, event_name):
        """Retrieve the handler relevant to the event"""
        if event_name == NEW_SASTOKEN_REQUIRED:
            return self._on_new_sastoken_required
        elif event_name == CONNECTION_STATE_CHANGE:
            return self._on_connection_state_change
        elif event_name == BACKGROUND_EXCEPTION:
            return self._on_background_exception
        else:
            return None

    @abc.abstractmethod
    def _receiver_handler_runner(self, inbox, handler_name):
        """Run infinite loop that waits for an inbox to receive an object from it, then calls
        the handler with that object
        """
        pass

    @abc.abstractmethod
    def _client_event_handler_runner(self, handler_name):
        """Run infinite loop that waits for the client event inbox to receive an event from it,
        then calls the handler that corresponds to that event
        """
        pass

    @abc.abstractmethod
    def _start_handler_runner(self, handler_name):
        """Create, and store a handler runner"""
        pass

    @abc.abstractmethod
    def _stop_receiver_handler_runner(self, handler_name):
        """Cancel and remove a handler runner"""
        pass

    @abc.abstractmethod
    def _stop_client_event_handler_runner(self):
        """Cancel the client event handler runner"""
        pass

    def _generic_receiver_handler_setter(self, handler_name, new_handler):
        """Set a handler"""
        curr_handler = getattr(self, handler_name)
        if new_handler is not None and curr_handler is None:
            # Create runner, set handler
            logger.debug("Creating new handler runner for handler: {}".format(handler_name))
            setattr(self, handler_name, new_handler)
            self._start_handler_runner(handler_name)
        elif new_handler is None and curr_handler is not None:
            # Cancel runner, remove handler
            logger.debug("Removing handler runner for handler: {}".format(handler_name))
            self._stop_receiver_handler_runner(handler_name)
            setattr(self, handler_name, new_handler)
        else:
            # Update handler, no need to change runner
            logger.debug("Updating set handler: {}".format(handler_name))
            setattr(self, handler_name, new_handler)

    @staticmethod
    def _generate_callback_for_handler(handler_name):
        """Define a callback that can handle errors during handler execution"""

        def handler_callback(future):
            try:
                e = future.exception(timeout=0)
            except Exception as raised_e:
                # This shouldn't happen because cancellation or timeout shouldn't occur...
                # But just in case...
                new_err = HandlerManagerException(
                    message="HANDLER ({}): Unable to retrieve exception data from incomplete invocation".format(
                        handler_name
                    ),
                    cause=raised_e,
                )
                handle_exceptions.handle_background_exception(new_err)
            else:
                if e:
                    new_err = HandlerManagerException(
                        message="HANDLER ({}): Error during invocation".format(handler_name),
                        cause=e,
                    )
                    handle_exceptions.handle_background_exception(new_err)
                else:
                    logger.debug(
                        "HANDLER ({}): Successfully completed invocation".format(handler_name)
                    )

        return handler_callback

    def stop(self, receiver_handlers_only=False):
        """Stop the process of invoking handlers in response to events.
        All pending items will be handled prior to stoppage.
        """
        # Stop receiver handlers
        for handler_name in self._receiver_handler_runners:
            if self._receiver_handler_runners[handler_name] is not None:
                self._stop_receiver_handler_runner(handler_name)

        # Stop the client event handler (if instructed)
        if not receiver_handlers_only and self._client_event_runner is not None:
            self._stop_client_event_handler_runner()

    def ensure_running(self):
        """Ensure the process of invoking handlers in response to events is running"""
        # Ensure any receiver handler set on the manager has a corresponding handler runner running
        for handler_name in self._receiver_handler_runners:
            if (
                self._receiver_handler_runners[handler_name] is None
                and getattr(self, handler_name) is not None
            ):
                self._start_handler_runner(handler_name)

        # Ensure client event handler runner is running if at least one client event handler is set
        # on the manager
        if self._client_event_runner is None:
            for event in client_events:
                handler = self._get_handler_for_client_event(event)
                if handler is not None:
                    self._start_handler_runner(CLIENT_EVENT)
                    break

    # ~~~Receiver Handlers~~~
    # Setting a receiver handler will start a dedicated runner for that handler
    # Removing a receiver handler will stop the dedicated runner for that handler
    @property
    def on_message_received(self):
        return self._on_message_received

    @on_message_received.setter
    def on_message_received(self, value):
        self._generic_receiver_handler_setter(MESSAGE, value)

    @property
    def on_method_request_received(self):
        return self._on_method_request_received

    @on_method_request_received.setter
    def on_method_request_received(self, value):
        self._generic_receiver_handler_setter(METHOD, value)

    @property
    def on_twin_desired_properties_patch_received(self):
        return self._on_twin_desired_properties_patch_received

    @on_twin_desired_properties_patch_received.setter
    def on_twin_desired_properties_patch_received(self, value):
        self._generic_receiver_handler_setter(TWIN_DP_PATCH, value)

    # ~~~Client Event Handlers~~~
    # Setting any client event handler will start the shared client event handler runner
    # Removing handlers will NOT stop the client event handler runner - you must use .stop()
    # Stopping when all client event handlers are removed could be added if necessary.
    @property
    def on_connection_state_change(self):
        return self._on_connection_state_change

    @on_connection_state_change.setter
    def on_connection_state_change(self, value):
        self._on_connection_state_change = value
        if self._client_event_runner is None:
            self._start_handler_runner(CLIENT_EVENT)

    @property
    def on_new_sastoken_required(self):
        return self._on_new_sastoken_required

    @on_new_sastoken_required.setter
    def on_new_sastoken_required(self, value):
        self._on_new_sastoken_required = value
        if self._client_event_runner is None:
            self._start_handler_runner(CLIENT_EVENT)

    @property
    def on_background_exception(self):
        return self._on_background_exception

    @on_background_exception.setter
    def on_background_exception(self, value):
        self._on_background_exception = value
        if self._client_event_runner is None:
            self._start_handler_runner(CLIENT_EVENT)


class SyncHandlerManager(AbstractHandlerManager):
    """Handler manager for use with synchronous clients"""

    def _receiver_handler_runner(self, inbox, handler_name):
        """Run infinite loop that waits for an inbox to receive an object from it, then calls
        the handler with that object
        """
        logger.debug("HANDLER RUNNER ({}): Starting runner".format(handler_name))
        _handler_callback = self._generate_callback_for_handler(handler_name)

        # Run the handler in a threadpool, so that it cannot block other handlers (from a different task),
        # or the main client thread. The number of worker threads forms an upper bound on how many instances
        # of the same handler can be running simultaneously.
        tpe = concurrent.futures.ThreadPoolExecutor(max_workers=4)
        while True:
            handler_arg = inbox.get()
            if isinstance(handler_arg, HandlerRunnerKillerSentinel):
                # Exit the runner when a HandlerRunnerKillerSentinel is found
                logger.debug(
                    "HANDLER RUNNER ({}): HandlerRunnerKillerSentinel found in inbox. Exiting.".format(
                        handler_name
                    )
                )
                tpe.shutdown()
                break
            # NOTE: we MUST use getattr here using the handler name, as opposed to directly passing
            # the handler in order for the handler to be able to be updated without cancelling
            # the running task created for this coroutine
            handler = getattr(self, handler_name)
            logger.debug("HANDLER RUNNER ({}): Invoking handler".format(handler_name))
            fut = tpe.submit(handler, handler_arg)
            fut.add_done_callback(_handler_callback)

    def _client_event_handler_runner(self):
        """Run infinite loop that waits for the client event inbox to receive an event from it,
        then calls the handler that corresponds to that event
        """
        logger.debug("HANDLER RUNNER (CLIENT EVENT): Starting runner")
        _handler_callback = self._generate_callback_for_handler("CLIENT_EVENT")

        tpe = concurrent.futures.ThreadPoolExecutor(max_workers=4)
        event_inbox = self._inbox_manager.get_client_event_inbox()
        while True:
            event = event_inbox.get()
            if isinstance(event, HandlerRunnerKillerSentinel):
                # Exit the runner when a HandlerRunnerKillerSentinel is found
                logger.debug(
                    "HANDLER RUNNER (CLIENT EVENT): HandlerRunnerKillerSentinel found in event queue. Exiting."
                )
                tpe.shutdown()
                break
            handler = self._get_handler_for_client_event(event.name)
            if handler is not None:
                logger.debug(
                    "HANDLER RUNNER (CLIENT EVENT): {} event received. Invoking {} handler".format(
                        event, handler
                    )
                )
                fut = tpe.submit(handler, *event.args_for_user)
                fut.add_done_callback(_handler_callback)
            else:
                logger.debug(
                    "No handler for event {} set. Skipping handler invocation".format(event)
                )

    def _start_handler_runner(self, handler_name):
        """Start and store a handler runner thread"""
        # Client Event handler flow
        if handler_name == CLIENT_EVENT:
            if self._client_event_runner is not None:
                # This branch of code should NOT be reachable due to checks prior to the invocation
                # of this method. The branch exists for safety.
                raise HandlerManagerException(
                    "Cannot create thread for handler runner: {}. Runner thread already exists".format(
                        handler_name
                    )
                )
            # Client events share a handler
            thread = threading.Thread(target=self._client_event_handler_runner)
            # Store the thread
            self._client_event_runner = thread

        # Receiver handler flow
        else:
            if self._receiver_handler_runners[handler_name] is not None:
                # This branch of code should NOT be reachable due to checks prior to the invocation
                # of this method. The branch exists for safety.
                raise HandlerManagerException(
                    "Cannot create thread for handler runner: {}. Runner thread already exists".format(
                        handler_name
                    )
                )
            inbox = self._get_inbox_for_receive_handler(handler_name)
            # Each receiver handler gets its own runner
            thread = threading.Thread(
                target=self._receiver_handler_runner, args=[inbox, handler_name]
            )
            # Store the thread
            self._receiver_handler_runners[handler_name] = thread

        # NOTE: It would be nice to have some kind of mechanism for making sure this thread
        # doesn't crash or raise errors, but it would require significant extra infrastructure
        # and an exception in here isn't supposed to happen anyway. Perhaps it could be added
        # later if truly necessary
        thread.daemon = True  # Don't block program exit
        thread.start()

    def _stop_receiver_handler_runner(self, handler_name):
        """Stop and remove a handler runner thread.
        All pending items in the corresponding inbox will be handled by the handler before stoppage.
        """
        logger.debug(
            "Adding HandlerRunnerKillerSentinel to inbox corresponding to {} handler runner".format(
                handler_name
            )
        )
        inbox = self._get_inbox_for_receive_handler(handler_name)
        inbox.put(HandlerRunnerKillerSentinel())

        # Wait for Handler Runner to end due to the sentinel
        logger.debug("Waiting for {} handler runner to exit...".format(handler_name))
        thread = self._receiver_handler_runners[handler_name]
        thread.join()
        self._receiver_handler_runners[handler_name] = None
        logger.debug("Handler runner for {} has been stopped".format(handler_name))

    def _stop_client_event_handler_runner(self):
        """Stop and remove a handler runner thread.
        All pending items in the client event queue will be handled by handlers (if they exist)
        before stoppage.
        """
        logger.debug("Adding HandlerRunnerKillerSentinel to client event queue")
        event_inbox = self._inbox_manager.get_client_event_inbox()
        event_inbox.put(HandlerRunnerKillerSentinel())

        # Wait for Handler Runner to end due to the stop command
        logger.debug("Waiting for client event handler runner to exit...")
        thread = self._client_event_runner
        thread.join()
        self._client_event_runner = None
        logger.debug("Handler runner for client events has been stopped")
