# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Threading utilities.
"""

from __future__ import absolute_import  # so we can import standard 'threading'

import sys
import itertools
import multiprocessing
from threading import (Thread, Lock)
from Queue import (Queue, Full, Empty)

from .exceptions import print_exception

class ExecutorException(Exception):
    pass


class DaemonThread(Thread):
    def __init__(self, *args, **kwargs):
        super(DaemonThread, self).__init__(*args, **kwargs)
        self.daemon = True

    def run(self):
        """
        We're overriding ``Thread.run`` in order to avoid annoying (but harmless) error messages
        during shutdown. The problem is that CPython nullifies the global state _before_ shutting
        down daemon threads, so that exceptions might happen, and then ``Thread.__bootstrap_inner``
        prints them out.

        Our solution is to swallow these exceptions here.

        The side effect is that uncaught exceptions in our own thread code will _not_ be printed out
        as usual, so it's our responsibility to catch them in our code.
        """

        try:
            super(DaemonThread, self).run()
        except SystemExit as e:
            # This exception should be bubbled up
            raise e
        except BaseException:
            # Exceptions might occur in daemon threads during interpreter shutdown
            pass


class Executor(object):
    """
    Executes tasks.

    Makes sure to gather all returned results and thrown exceptions in one place, in order of task
    submission.

    Example::

        def sum(arg1, arg2):
            return arg1 + arg2

        executor = FixedThreadPoolExecutor(10)
        try:
            for value in range(100):
                executor.submit(sum, value, value)
            executor.drain()
        except:
            executor.close()
        executor.raise_first()
        print executor.returns

    You can also use it with the Python ``with`` keyword, in which case you don't need to call
    ``close`` explicitly::

        with FixedThreadPoolExecutor(10) as executor:
            for value in range(100):
                executor.submit(sum, value, value)
            executor.drain()
            executor.raise_first()
            print executor.returns
    """

    def __init__(self, print_exceptions=False):
        self.print_exceptions = print_exceptions

    def submit(self, func, *args, **kwargs):
        """
        Submit a task for execution.

        The task will be called ASAP on the next available worker thread in the pool.

        :raises ExecutorException: if cannot be submitted
        """
        raise NotImplementedError

    def close(self):
        """
        Blocks until all current tasks finish execution and all worker threads are dead.

        You cannot submit tasks anymore after calling this.

        This is called automatically upon exit if you are using the ``with`` keyword.
        """
        pass

    def drain(self):
        """
        Blocks until all current tasks finish execution, but leaves the worker threads alive.
        """
        pass

    @property
    def returns(self):
        """
        The returned values from all tasks, in order of submission.
        """
        return ()

    @property
    def exceptions(self):
        """
        The raised exceptions from all tasks, in order of submission.
        """
        return ()

    def raise_first(self):
        """
        If exceptions were thrown by any task, then the first one will be raised.

        This is rather arbitrary: proper handling would involve iterating all the exceptions.
        However, if you want to use the "raise" mechanism, you are limited to raising only one of
        them.
        """

        exceptions = self.exceptions
        if exceptions:
            raise exceptions[0]

    def __enter__(self):
        return self

    def __exit__(self, the_type, value, traceback):
        pass


class BlockingExecutor(Executor):
    """
    Executes tasks in the current thread.
    """

    def __init__(self, print_exceptions=False):
        super(BlockingExecutor, self).__init__(print_exceptions=print_exceptions)
        self._returns = []
        self._exceptions = []

    def submit(self, func, *args, **kwargs):
        try:
            result = func(*args, **kwargs)
            self._returns.append(result)
        except Exception as e:
            self._exceptions.append(e)
            if self.print_exceptions:
                print_exception(e)

    @property
    def returns(self):
        return self._returns

    @property
    def exceptions(self):
        return self._exceptions


# https://gist.github.com/tliron/81dd915166b0bfc64be08b4f8e22c835
class FixedThreadPoolExecutor(Executor):
    """
    Executes tasks in a fixed thread pool.
    """

    _CYANIDE = object() # special task marker used to kill worker threads

    def __init__(self,
                 size=None,
                 timeout=None,
                 print_exceptions=False):
        """
        :param size: number of threads in the pool; if ``None`` will use an optimal number for the
         platform
        :param timeout: timeout in seconds for all blocking operations (``None`` means no timeout)
        :param print_exceptions: set to ``True`` in order to print exceptions from tasks
        """
        super(FixedThreadPoolExecutor, self).__init__(print_exceptions=print_exceptions)

        if not size:
            try:
                size = multiprocessing.cpu_count() * 2 + 1
            except NotImplementedError:
                size = 3

        self.size = size
        self.timeout = timeout

        self._tasks = Queue()
        self._returns = {}
        self._exceptions = {}
        self._id_creator = itertools.count()
        self._lock = Lock() # for console output

        self._workers = []
        for index in range(size):
            worker = DaemonThread(
                name='{0}{1:d}'.format(self.__class__.__name__, index),
                target=self._thread_worker)
            worker.start()
            self._workers.append(worker)

    def submit(self, func, *args, **kwargs):
        try:
            self._tasks.put((self._id_creator.next(), func, args, kwargs), timeout=self.timeout)
        except Full:
            raise ExecutorException('cannot submit task: queue is full')

    def close(self):
        self.drain()
        while self.is_alive:
            try:
                self._tasks.put(self._CYANIDE, timeout=self.timeout)
            except Full:
                raise ExecutorException('cannot close executor: a thread seems to be hanging')
        self._workers = None

    def drain(self):
        self._tasks.join() # oddly, the API does not support a timeout parameter

    @property
    def is_alive(self):
        """
        True if any of the worker threads are alive.
        """

        for worker in self._workers:
            if worker.is_alive():
                return True
        return False

    @property
    def returns(self):
        return [self._returns[k] for k in sorted(self._returns)]

    @property
    def exceptions(self):
        return [self._exceptions[k] for k in sorted(self._exceptions)]

    def _thread_worker(self):
        while True:
            if not self._execute_next_task():
                break

    def _execute_next_task(self):
        try:
            task = self._tasks.get(timeout=self.timeout)
        except Empty:
            # Happens if timeout is reached
            return True
        if task == self._CYANIDE:
            # Time to die :(
            return False
        self._execute_task(*task)
        return True

    def _execute_task(self, task_id, func, args, kwargs):
        try:
            result = func(*args, **kwargs)
            self._returns[task_id] = result
        except Exception as e:
            self._exceptions[task_id] = e
            if self.print_exceptions:
                with self._lock:
                    print_exception(e)
        self._tasks.task_done()

    def __enter__(self):
        return self

    def __exit__(self, the_type, value, traceback):
        self.close()


class LockedList(list):
    """
    A list that supports the ``with`` keyword with a built-in lock.

    Though Python lists are thread-safe in that they will not raise exceptions during concurrent
    access, they do not guarantee atomicity. This class will let you gain atomicity when needed.
    """

    def __init__(self, *args, **kwargs):
        super(LockedList, self).__init__(*args, **kwargs)
        self.lock = Lock()

    def __enter__(self):
        return self.lock.__enter__()

    def __exit__(self, the_type, value, traceback):
        return self.lock.__exit__(the_type, value, traceback)


class ExceptionThread(Thread):
    """
    A thread from which top level exceptions can be retrieved or re-raised.
    """
    def __init__(self, *args, **kwargs):
        Thread.__init__(self, *args, **kwargs)
        self.exception = None
        self.daemon = True

    def run(self):
        try:
            super(ExceptionThread, self).run()
        except BaseException:
            self.exception = sys.exc_info()

    def is_error(self):
        return self.exception is not None

    def raise_error_if_exists(self):
        if self.is_error():
            type_, value, trace = self.exception
            raise type_, value, trace