"""Contains executor implementation for aplex."""

import asyncio
import atexit
import concurrent.futures
import enum
import inspect
import itertools
import os
import queue
import signal
import threading
import time
import warnings
import weakref
from typing import (Any, AsyncIterator, Callable, Dict, Generator,
                    Iterable, List, Optional, Set, Tuple, Union)

from . import compat
from .futures import ConcurrentFuture, AsyncioFuture
from .load_balancers import LoadBalancer, RoundRobin
if compat.PY36:
    from typing import AsyncGenerator


DEFAULT_POOL_SIZE = os.cpu_count() or 2  # os.cpu_count() may return None.
DEFAULT_MAX_WORKS_PER_WORKER = 300

# XXX(Lun): Let users choose whether to shut down if broken or exception
# happens in worker? Below are both true for now.
# SHUTDOWN_IF_BROKEN = False
# SHUTDOWN_IF_EXCEPTION = False

_executor_track_set_lock = threading.Lock()
_executor_track_set = weakref.WeakSet()  # type: Set[BaseAsyncPoolExecutor]


class AplexWorkerError(OSError):
    """The executor are broken or an BaseException raised."""


class HandlerCommands(enum.Enum):
    """Commands sent from handler to workers."""

    CANCEL = 'Cancel'
    CLOSE = 'Close'


class WorkStates(enum.Enum):
    """States of work in workers."""

    SUCCESS = 'Success'
    CANCELLED = 'Cancelled'
    EXCEPTION = 'Exception'


class WorkerStates(enum.Enum):
    """States of workers."""

    BROKEN = 'Broken'  # For worker closed due to unknown reason.
    ERROR = 'Error'  # For BaseException raised in worker.
    CLOSING = 'Closing'


class SubmitForms(enum.Enum):
    """Forms of submit items."""

    WORK = 'Work'
    COMMAND = 'Command'


class ResultForms(enum.Enum):
    """Forms of result items."""

    WORK_STATE = 'Work state'
    WORKER_STATE = 'Worker state'


class SubmitItem(object):
    """Object sent from executor to workers"""

    __slots__ = (['form'] +
                 ['work', 'args', 'kwargs',
                  'work_id', 'load_balancing_meta'] +
                 ['command', 'command_meta'])

    @classmethod
    def work_form(cls,
                  work: Callable,
                  args: Tuple[Any],
                  kwargs: Dict[str, Any],
                  work_id: int,
                  load_balancing_meta: Optional[Any] = None
                  ) -> 'SubmitItem':
        """Generates a submit item in work form."""

        item = cls.__new__(cls)
        item.form = SubmitForms.WORK

        item.work = work  # Named "work" to differentiate from asyncio task.
        item.args = args
        item.kwargs = kwargs
        item.work_id = work_id
        item.load_balancing_meta = load_balancing_meta
        return item

    @classmethod
    def command_form(cls,
                     command: HandlerCommands,
                     command_meta: Optional[Any] = None
                     ) -> 'SubmitItem':
        """Generates a submit item in command form."""

        item = cls.__new__(cls)
        item.form = SubmitForms.COMMAND

        item.command = command
        item.command_meta = command_meta
        return item


class ResultItem(object):
    """Object received from workers."""

    __slots__ = (['form'] +
                 ['work_id', 'work_state', 'work_state_meta'] +
                 ['worker_state', 'worker_state_meta'])

    @classmethod
    def work_state_form(cls,
                        work_id: int,
                        work_state: WorkStates,
                        work_state_meta: Optional[Any] = None,
                        ) -> 'ResultItem':
        """Generates a result item in work state form."""

        item = cls.__new__(cls)
        item.form = ResultForms.WORK_STATE

        item.work_id = work_id
        item.work_state = work_state
        item.work_state_meta = work_state_meta
        return item

    @classmethod
    def worker_state_form(cls,
                          worker_state: WorkerStates,
                          worker_state_meta: Optional[Any] = None
                          ) -> 'ResultItem':
        """Generates a result item in worker state form."""

        item = cls.__new__(cls)
        item.form = ResultForms.WORKER_STATE

        item.worker_state = worker_state
        item.worker_state_meta = worker_state_meta
        return item


class BaseWorkerMixin(object):
    """Implementation of worker.
    
    This is a mixin for threading.Thread and multiprocessing.Process.
    """
    _submit_queue = None

    def __init__(self, result_queue, loop_factory=None):
        super().__init__()

        self._result_queue = result_queue
        self._loop_factory = loop_factory

    def submit(self, item: SubmitItem):
        self._submit_queue.put(item)

    def cancel_work_by_id(self, work_id):
        self._submit_queue.put(
            SubmitItem.command_form(HandlerCommands.CANCEL,
                                    command_meta=work_id))

    def close(self):
        """"""
        self._submit_queue.put(
            SubmitItem.command_form(HandlerCommands.CLOSE, None))

    def clear(self):
        self._submit_queue = None
        self._result_queue = None

    def run(self):
        """Target for threading.Thread and multiprocessing.Process."""

        # Subtasks generated in corouines are not included to prevent
        # from cancel a task twice. Users have to handle CancelledError
        # properly to cancel their subtasks.
        running_tasks = {}

        def cancel(work_id):
            # Although not creating task for callables that are not
            # a coroutine function, The work id here must map to
            # a task, since only this type of work is allowed
            # to call cancel interface.
            try:
                running_tasks[work_id].cancel()
            except KeyError:
                pass  # Task has been done and popped.

        def run_uncancellable(uncancellable, args, kwargs, work_id):
            """Runs uncancellable, i.e., works that are not coroutine
            functions.
            """
            try:
                result = uncancellable(*args, **kwargs)
            except Exception as e:
                item = ResultItem.work_state_form(work_id,
                                                  WorkStates.EXCEPTION,
                                                  self._make_picklable(e))
            else:
                item = ResultItem.work_state_form(work_id,
                                                  WorkStates.SUCCESS,
                                                  result)
            self._send_result_back(work_id, item)

        def task_done_callback(task):
            work_id = task.work_id
            running_tasks.pop(work_id)

            if task.cancelled():
                item = ResultItem.work_state_form(work_id,
                                                  WorkStates.CANCELLED,
                                                  None)
            else:
                exception = task.exception()
                if exception is not None:
                    exception = self._make_picklable(exception)
                    item = ResultItem.work_state_form(work_id,
                                                      WorkStates.EXCEPTION,
                                                      exception)
                else:
                    item = ResultItem.work_state_form(work_id,
                                                      WorkStates.SUCCESS,
                                                      task.result())
            self._send_result_back(work_id, item)

        def run_coroutine(coro_function, args, kwargs, work_id):
            """Creates a task of the coroutine generated from 
            coroutine function work.
            """
            task = loop.create_task(coro_function(*args, **kwargs))
            task.work_id = work_id
            running_tasks[work_id] = task
            task.add_done_callback(task_done_callback)

        # TODO(Lun): Wrap this with try-except to catch
        # unexpected exception and than close worker.
        def handle_submit_queue(loop):
            """Target for the thread handling submit queue."""
            while True:
                item = self._submit_queue.get()

                # Prevents from doing anything after a BaseException raises.
                if item is None:
                    return

                if item.form == SubmitForms.COMMAND:
                    # XXX(Lun): Use dict as switch in C if too many cases.
                    if item.command == HandlerCommands.CLOSE:
                        loop.call_soon_threadsafe(loop.stop)
                        return
                    elif item.command == HandlerCommands.CANCEL:
                        work_id = item.command_meta
                        loop.call_soon_threadsafe(cancel, work_id)
                        continue
                    assert False, 'Should be unreachable here.'

                if not inspect.iscoroutinefunction(item.work):
                    # Not create future for a callable that is not a
                    # coroutine function since it can not be cancelled.
                    loop.call_soon_threadsafe(
                        run_uncancellable,
                        item.work, item.args, item.kwargs, item.work_id)
                else:
                    loop.call_soon_threadsafe(
                        run_coroutine,
                        item.work, item.args, item.kwargs, item.work_id)

        def cancel_tasks(*tasks):
            if not tasks:
                return
            for task in tasks:
                task.cancel()
            loop.run_until_complete(
                asyncio.gather(*tasks, return_exceptions=True))

        try:
            if self._loop_factory is None:
                loop = asyncio.new_event_loop()
            else:
                loop = self._loop_factory()
            asyncio.set_event_loop(loop)

            handle_submit_thread = threading.Thread(
                target=handle_submit_queue, args=(loop,), daemon=True)
            handle_submit_thread.start()

            loop.run_forever()
        except BaseException as e:
            # Close handle submit thread so that any command from
            # handler is not allowed.
            self._submit_queue.put(None)  # Close handle submit thread.

            self._result_queue.put(
                ResultItem.worker_state_form(WorkerStates.ERROR,
                                             self._make_picklable(e)))
        finally:
            try:
                cancel_tasks(*running_tasks.values())

                # Cancels tasks the user forgot to handle when
                # asyncio.CancelledError is raised.
                if compat.PY37:
                    # Must pass loop for get_running_loop().
                    orphan_tasks = asyncio.tasks.all_tasks(loop)
                else:
                    orphan_tasks = asyncio.Task.all_tasks()
                cancel_tasks(*orphan_tasks)

                if compat.PY36:
                    loop.run_until_complete(loop.shutdown_asyncgens())

                    # Run scheduled callbacks.
                    loop.call_soon(loop.stop)
                    loop.run_forever()

                # Should report only when close command is received and
                # cancellation succeeds. See annotation in below except
                # BaseException clause.
                self._result_queue.put(
                    ResultItem.worker_state_form(WorkerStates.CLOSING,
                                                 self.ident))
            except BaseException as e:
                # Cancellation fails with new BaseException when closing
                # or fails due to previous BaseException. See details in
                # asyncio.Task._wakeup. The previous BaseException would
                # raised again due to `future.result()``.

                self._result_queue.put(
                    ResultItem.worker_state_form(WorkerStates.ERROR,
                                                 self._make_picklable(e)))
            finally:
                loop.close()

    def _send_result_back(self, work_id, item):
        try:
            self._result_queue.put(item)
        except TypeError as e:
            # Item unpicklable.
            item = ResultItem.work_state_form(work_id,
                                              WorkStates.EXCEPTION,
                                              self._make_picklable(e))
            self._result_queue.put(item)

    def _make_picklable(self, stuff):
        """Makes some stuff picklable or not loss information when pickling.

        This should be overwritten by ProcessWorker subclass.
        """
        return stuff


class BaseWorkerManager(object):
    """A manager of worker, or in other words, a broker."""

    # All workers of a executor share a common result queue.
    _result_queue = None
    _Worker = None

    def __init__(self,
                 pool_size,
                 max_works_per_worker,
                 load_balancer,
                 worker_loop_factory):
        self._workers = []
        for __ in range(pool_size):
            worker = self._Worker(self._result_queue, worker_loop_factory)
            worker.start()
            self._workers.append(worker)

        self._work_owner = {}  # To find worker with work_id.
        # Dict for load balancing, shared with self._load_balancer.
        self._workloads = dict.fromkeys(self._workers, 0)

        self._load_balancer = load_balancer(self._workers,
                                            self._workloads,
                                            max_works_per_worker)

    def submit(self, item: SubmitItem, future, load_balancing_meta):
        work_id = item.work_id
        worker = self._load_balancer.get_proper_worker(load_balancing_meta)

        # Caculates workload here.
        self._work_owner[work_id] = worker
        self._workloads[worker] += 1

        def work_done_callback(done_future):
            worker = self._work_owner.pop(work_id)
            self._workloads[worker] -= 1

        future.add_done_callback(work_done_callback)
        worker.submit(item)

    def cancel_work(self, future: ConcurrentFuture):
        work_id = id(future)

        worker = self._work_owner[work_id]
        worker.cancel_work_by_id(work_id)

    def wait_return_item(self) -> ResultItem:
        """Waits for workers until result item is returned
        
        Returns:
            Result item from some worker.
        
        Raises:
            ConnectionAbortedError: If some workers are broken.
        """
        raise NotImplementedError()

    def remove_closed_worker(self, worker_id):
        for worker in self._workers:
            if worker.ident == worker_id:
                target = worker
                break
        assert self._workloads[target] == 0

        target.clear()
        self._workers.remove(target)
        del self._workloads[target]

        return target  # For subclass' super(). No need to find it again.

    def are_all_workers_closed(self):
        return False if self._workers else True

    def wakeup(self):
        """To wake up form wait_return_item method."""
        self._result_queue.put(None)

    def close(self):
        for worker in self._workers:
            worker.close()

    def join(self):
        # XXX(Lun): Join removed worker?
        for worker in self._workers:
            worker.join()

    def clear(self):
        self._result_queue = None
        for worker in self._workers:
            worker.clear()
        self._workers.clear()
        self._workloads.clear()


class _HandlerThread(threading.Thread):
    """Handler of executor."""

    def __init__(self, executor, pool_size, max_works_per_worker,
                 load_balancer, worker_loop_factory,
                 WorkerManager):
        super().__init__()

        self._queue = queue.Queue()
        self._executor = weakref.ref(executor)

        self._running_futures = {}
        self._max_runnings = pool_size * max_works_per_worker

        self._worker_manager = WorkerManager(pool_size,
                                             max_works_per_worker,
                                             load_balancer,
                                             worker_loop_factory)

        self._executor_shutdown_called = False
        self._close_lock = threading.Lock()

    def submit(self, item: Tuple[SubmitItem, ConcurrentFuture, Any]):
        self._queue.put(item)

        # To wake up from wait_return_item in self.run().
        self._worker_manager.wakeup()

    def cancel_work(self, future: ConcurrentFuture):
        """An interface set in future to cancel work."""
        self._worker_manager.cancel_work(future)

    def close(self):
        with self._close_lock:
            self._executor_shutdown_called = True
            self._worker_manager.close()

    def run(self):
        try:
            self._run()
        except AplexWorkerError:
            pass  # TODO(Lun): logger.error
        except BaseException:
            # Close workers if some bug happens.
            self._worker_manager.close()
            raise
        else:
            self._worker_manager.join()
        finally:
            self._clear()

    def _run(self):
        while True:
            with self._close_lock:
                if not self._executor_shutdown_called:
                    self._transfer_to_worker_manager()

            try:
                # Blocks here until some result is returned or
                # some worker is terminated.
                item = self._worker_manager.wait_return_item()
            except ConnectionAbortedError as e:
                item = ResultItem.worker_state_form(WorkerStates.BROKEN, e)

            if item is None:
                # Just a wakeup sentinel from wait_return_item to take a
                # new work submitted by the user.
                continue

            # Handles worker state information.

            if item.form == ResultForms.WORKER_STATE:
                self._handle_worker_state(item.worker_state,
                                          item.worker_state_meta)
                if self._worker_manager.are_all_workers_closed():
                    break
                else:
                    continue

            # Handles received result from finished or cancelled work.

            future = self._running_futures.pop(item.work_id)
            state = item.work_state

            # A condition to avoid cancelling by the user.
            with future._condition:
                if state == WorkStates.CANCELLED:
                    future._cancel_running()
                elif future._cancel_interface_called:
                    exc = concurrent.futures.CancelledError(
                        'Work was done before cancellation function '
                        'is scheduled.')
                    future.set_exception(exc)
                elif state == WorkStates.SUCCESS:
                    future.set_result(item.work_state_meta)
                elif state == WorkStates.EXCEPTION:
                    # concurrent.futures.asyncio.CancelledError and
                    # asyncio.CancelledError are included here.
                    future.set_exception(item.work_state_meta)
                else:
                    assert False, 'Should be unreachable here.'

    def _transfer_to_worker_manager(self):
        """Sends received work to worker manager, i.e., broker."""

        while True:
            # All workers reach the limit of max works per worker.
            if len(self._running_futures) >= self._max_runnings:
                return

            try:
                item = self._queue.get(block=False)
            except queue.Empty:
                return
            else:
                submit_item, future, load_balancing_meta = item
                if future.set_running_or_notify_cancel():
                    self._running_futures[submit_item.work_id] = future
                    self._worker_manager.submit(
                        submit_item, future, load_balancing_meta)

    def _handle_worker_state(self,
                             state: WorkerStates,
                             state_meta: Optional[BaseException]):
        if state == WorkerStates.CLOSING:
            # This will happen only if executor shutdown is called by users.
            assert self._executor_shutdown_called

            # Removes this worker from listening list and clear pending
            # works. All the remaining running works will be cancelled by
            # the other alive workers.

            # Make sure removal is after closing.
            with self._close_lock:
                # This will not cause bugs since submit are not
                # allowed after closing.
                self._worker_manager.remove_closed_worker(
                    worker_id=state_meta)

            # Putting to queue is not allowed after executor.shutdown.
            # Thus no thread-safty issue. Just use deque.
            for __, pending_future, __ in self._queue.queue:
                pending_future.cancel()
                pending_future.set_running_or_notify_cancel()
            self._queue.queue.clear()

        elif state in (WorkerStates.BROKEN, WorkerStates.ERROR):
            # Worker not closed correctly. Work raised
            # BaseException or some worker was terminated abruptly.

            # TODO(Lun): Add worker type information in the reasons.

            # Sets exception reason to executor.
            if state == WorkerStates.BROKEN:
                worker_error = AplexWorkerError(
                    'Aplex worker was terminated abruptly.')
            elif state == WorkerStates.ERROR:
                worker_error = AplexWorkerError(
                    'Aplex worker was terminated since the event '
                    'loop has crashed due to raised BaseException.')
            worker_error.__cause__ = state_meta

            executor = self._executor()
            if executor is not None:
                with executor._lock:
                    executor._exception = worker_error
                executor = None

            # Sets exception reason to undone futures.
            future_error = concurrent.futures.CancelledError()
            future_error.__cause__ = worker_error
            for __, pending_future, in self._queue.queue:
                pending_future.set_exception(future_error)
            for running_future in self._running_futures.values():
                running_future.set_exception(future_error)

            self._worker_manager.close()  # Closes all related workers.

            raise worker_error  # Forces handler thread to terminate.

    def _clear(self):
        self._queue.queue.clear()  # I.e., deque.clear().
        self._running_futures.clear()

        self._worker_manager.join()
        self._worker_manager.clear()


# It's needed to define these here since a local object is not pickable.
async def _handle_chunk_async(work, chunk: Iterable[Iterable]) -> List:
    return await asyncio.gather(*(work(*args) for args in chunk),
                                return_exceptions=False)


def _handle_chunk(work, chunk: Iterable[Iterable]) -> List:
    return [work(*args) for args in chunk]


class BaseAsyncPoolExecutor(object):
    """Executor implementation."""

    _WorkerManager = None

    def __init__(self, *,
                 pool_size: Optional[int] = DEFAULT_POOL_SIZE,
                 max_works_per_worker:
                     Optional[int] = DEFAULT_MAX_WORKS_PER_WORKER,
                 load_balancer: Optional[LoadBalancer] = RoundRobin,
                 awaitable: Optional[bool] = False,
                 future_loop: asyncio.AbstractEventLoop = None,
                 worker_loop_factory:
                     Optional[asyncio.AbstractEventLoop] = None):
        """Setups executor and adds self to executor track set.

        Args:
            pool_size: Number of workers, i.e., number of threads or processes.
            max_works_per_worker: The max number of works a worker can run at
                the same time. This does not **limit** the number of asyncio
                tasks of a worker.
            load_balancer: A subclass of aplex.LoadBalancer for submitted
                item load balancing that has implemented abstract method
                ``get_proper_worker``.
            awaitable: If it's set to True, futures returned from ``submit``
                method will be awaitable, and ``map`` will return async
                generator(async iterator if python3.5).
            future_loop: Loop instance set in awaitable futures returned from
                ``submit`` method.

                If specified, ``awaitable`` must be set to true.

                This loop can also be set in ``set_future_loop`` method.
            worker_loop_factory: A factory to generate loop instance for
                workers to run their job.

        Raises:
            ValueError:
                ``future_loop`` is specified while ``awaitable`` is False.

        """
        if future_loop and not awaitable:
            raise ValueError('Awaitable should be set to True if '
                             'future_loop is specified.')

        with _executor_track_set_lock:
            _executor_track_set.add(self)

        self._awaitable = awaitable
        self._loop = future_loop

        self._closed = False
        self._exception = None
        self._lock = threading.Lock()

        self._handler_thread = _HandlerThread(self,
                                              pool_size,
                                              max_works_per_worker,
                                              load_balancer,
                                              worker_loop_factory,
                                              self._WorkerManager)

        # Important to set thread daemon to True such that one can
        # trigger functions registerd to atexit, closing workers gracefully.
        self._handler_thread.setDaemon(True)
        self._handler_thread.start()

    def set_future_loop(self, loop: asyncio.AbstractEventLoop):
        """Sets loop for awaitable futures to await results.

        This loop can also be set in initialization.

        Args:
            loop: The Loop needed for awaitable futures.

        Raises:
            RuntimeError: If executor has been shut down, or
                executor is set to be unawaitable.
            AplexWorkerError: If some workers are broken or raise
                BaseException.

        """
        with self._lock:
            if self._closed:
                raise RuntimeError('Executor has been shut down.')
            elif self._exception:
                raise self._exception

            if not self._awaitable:
                raise RuntimeError('Assigning event loop to an unawaitable')
            self._loop = loop

    def submit(self,
               work: Callable,
               *args: Tuple,
               load_balancing_meta: Optional[Any] = None,
               **kwargs: Dict
    ) -> Union[AsyncioFuture, ConcurrentFuture]:
        """submits your work like the way in concurrent.futures.

        The work submitted will be sent to the specific worker that
        the load balancer choose.

        Note:
            The work you submit should be a callable, And a ``coroutine``
            is **not** a callable. You should submit a ``coroutine function``
            and specify its args and kwargs here instead.

        Args:
            work: The callable that will be run in a worker.
            *args: Position arguments for work.
            load_balancing_meta: This will be passed to load balancer for
                the choice of proper worker.
            **kwargs: Keyword arguments for work.

        Returns:
            A future.

            The future will be awaitable like that in asyncio if ``awaitable``
            is set to True in executor construction, otherwise, unawaitable
            like that in concurrent.futures.

        Raises:
            RuntimeError: If executor has been shut down.
            AplexWorkerError: If some workers are broken or raise
                BaseException.
            TypeError: If work is not a callable.

        """
        with self._lock:
            if self._closed:
                raise RuntimeError('Executor has been shut down.')
            elif self._exception:
                raise self._exception

            if not callable(work):
                raise TypeError('{} is not a callable to run.'.format(work))

            # For running callables, only coroutine can be cancelled.
            cancel_interface = None
            if inspect.iscoroutinefunction(work):
                cancel_interface = self._handler_thread.cancel_work

            concurrent_future = ConcurrentFuture(cancel_interface)
            if self._awaitable:
                # Chain these two futures before submit for thread-safty.
                asyncio_future = AsyncioFuture(concurrent_future,
                                               self._loop)

            item = SubmitItem.work_form(work, args, kwargs,
                                        id(concurrent_future))
            self._handler_thread.submit((item, concurrent_future,
                                         load_balancing_meta))
            return asyncio_future if self._awaitable else concurrent_future

    if compat.PY36:
        _MapReturnAsync = AsyncGenerator
    else:
        _MapReturnAsync = AsyncIterator

    def map(self,
            work: Callable,
            *iterables: Tuple[Iterable],
            timeout: Optional[float] = None,
            chunksize: int = 1,
            load_balancing_meta: Optional[Any] = None
    ) -> Union[_MapReturnAsync, Generator]:
        """map your work like the way in concurrent.futures.

        The work submitted will be sent to the specific worker that
        the load balancer choose.

        Note:
            The work you submit should be a callable, And a ``coroutine``
            is **not** a callable. You should submit a ``coroutine function``
            and specify its args and kwargs here instead.

        Args:
            work: The callable that will be run in a worker.
            *iterables: Position arguments for work. All of them are iterable
                and have same length.
            timeout: The time limit for waiting results.
            chunksize: Works are gathered, partitioned as chunks in this size,
                and then sent to workers.
            load_balancing_meta: This will be passed to load balancer for
                the choice of proper worker.

        Returns:
            A async generator yielding the map results if ``awaitable`` is set
            to True, otherwise a generator. In python3.5, async iterator is
            used to replace async generator.

            If a exception is raised in a work, it will be re-raised in the
            generator, and the remaining works will be cancelled.

        Raises:
            ValueError: If chunksize is less than 1.
            TypeError: If work is not a callable.
        """
        if chunksize < 1:
            raise ValueError('Chunksize should be at least 1.')

        if inspect.iscoroutinefunction(work):
            work_on_chunk = _handle_chunk_async
        elif callable(work):
            work_on_chunk = _handle_chunk
        else:
            raise TypeError('{} is not a callable to run.'.format(work))

        # Include internal process/thread communication time.
        end_time = timeout + time.time() if timeout is not None else None

        futures = []
        total = zip(*iterables)
        chunk = tuple(itertools.islice(total, chunksize))
        while chunk:
            futures.append(
                self.submit(work_on_chunk, work, chunk,
                            load_balancing_meta=load_balancing_meta))
            chunk = tuple(itertools.islice(total, chunksize))

        if not futures:
            return []

        if self._awaitable:
            if not compat.PY36:
                return self._MapAsyncIterator(futures, end_time, self._loop)
            return self._map_async_gen(futures, end_time)
        else:
            return self._map_gen(futures, end_time)


    if compat.PY36:
        # If statement with variable condition can not prevent Python3.5
        # from compiling this function to bytecode, causing syntax error
        # in Python3.5, which does not support async generator.

        from .py35_support import _map_async_gen
    else:
        # According to PEP525: `asynchronous generators are 2x faster
        # than an equivalent implemented as an asynchronous iterator`,
        # async iterator is used only for python3.5, which does not
        # support async generator.

        class _MapAsyncIterator(object):
            """A async iterator to replace async generator in python3.5."""

            def __init__(self,
                         futures: AsyncioFuture,
                         end_time=None,
                         loop=None):
                self._futures = futures
                self._end_time = end_time
                self._loop = loop

                self._results = None

            if compat.PY352:
                def __aiter__(self):
                    return self
            else:
                async def __aiter__(self):
                    return self

            async def __anext__(self):
                if not self._results:
                    if not self._futures:
                        raise StopAsyncIteration()

                    future = self._futures.pop(0)
                    try:
                        if self._end_time is not None:
                            self._results = await asyncio.wait_for(
                                future, self._end_time - time.time(), loop=self._loop)
                        else:
                            self._results = await future
                    except Exception:
                        future.cancel()
                        for future in self._futures:
                            future.cancel()
                        raise
                return self._results.pop(0)


    def _map_gen(self,
                 futures: List[ConcurrentFuture],
                 end_time=None) -> Generator:
        """The generator that ``map`` return when ``awaitable`` is False."""

        try:
            while futures:
                future = futures.pop(0)
                if end_time is not None:
                    yield from future.result(end_time - time.time())
                else:
                    yield from future.result()
        # Finally clause for generator exit and timeout.
        finally:
            future.cancel()  # This future may have been done.
            for future in futures:
                future.cancel()

    def shutdown(self, wait: bool = True):
        """Shuts down the executor and frees the resource.

        Args:
            wait: Whether to block until shutdown is finished.

        """
        import concurrent
        with self._lock:
            if self._closed:
                return
            self._closed = True

        # Handler may has already been closed.
        if self._exception is None:
            self._handler_thread.close()

        if wait:
            self._handler_thread.join()

        self._handler_thread = None
        self._loop = None

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.shutdown(wait=True)

    def __del__(self):
        # XXX
        if not self._closed and self._exception is None:
            warnings.warn('{} got garbage-collected. Started to '
                          'shutdown'.format(self))
            self.shutdown(wait=False)


@atexit.register  # This decorator will return _ensure_close back.
def _ensure_close():
    """Closes all workers for graceful exit."""

    # Use tuple to construct strong ref to avoid gc-collection during the loop.
    for executor in tuple(_executor_track_set):
        executor.shutdown(wait=True)


def _register_keyboard_interrupt_handler():
    origin = signal.getsignal(signal.SIGINT)

    def handler(signum, stack):
        _ensure_close()
        print(signum)
        origin(signum, stack)

    signal.signal(signal.SIGINT, handler)


_register_keyboard_interrupt_handler()
