from concurrent.futures import Executor
from threading import Thread, Lock
from collections import namedtuple, deque
from functools import partial
import logging
import weakref

from .common import MAX_TIMEOUT
from .wrap import CanCustomizeBind
from .map import MapFuture
from .helpers import executor_loop, ShutdownHelper
from .event import get_event, is_shutdown
from .logwrap import LogWrapper

from .metrics import metrics, track_future


class ThrottleFuture(MapFuture):
    def __init__(self, executor):
        self._executor = executor
        super(ThrottleFuture, self).__init__(delegate=None, map_fn=lambda x: x)
        self.add_done_callback(self._clear_executor)

    def _me_cancel(self):
        if self._delegate:
            return self._delegate.cancel()
        executor = self._executor
        return executor and executor._do_cancel(self)

    @classmethod
    def _clear_executor(cls, future):
        future._executor = None


ThrottleJob = namedtuple("ThrottleJob", ["future", "fn", "args", "kwargs"])


class AtomicInt(object):
    def __init__(self):
        self.value = 0
        self.lock = Lock()

    def decr(self):
        with self.lock:
            self.value -= 1

    def incr(self):
        with self.lock:
            self.value += 1


class ThrottleExecutor(CanCustomizeBind, Executor):
    """An executor which delegates to another executor while enforcing
    a limit on the number of futures running concurrently.

    - Callables are submitted to the delegate executor, from a different
      thread than the calling thread.

    - Where `count` is used to initialize this executor, if there
      are already `count` futures submitted to the delegate executor and not
      yet :meth:`~concurrent.futures.Future.done`, additional callables will
      either be queued or will block on submit, and will only be submitted
      to the delegate executor once there are less than `count` futures in
      progress.

    .. versionadded:: 1.9.0
    """

    def __init__(self, delegate, count, logger=None, name="default", block=False):
        """
        Parameters:
            delegate (~concurrent.futures.Executor):
                an executor to which callables will be submitted

            count (int, callable):
                int:
                    maximum number of concurrently running futures.
                callable:
                    a callable which returns an ``int`` (or ``None``, to indicate
                    no throttling).

                    The callable will be invoked each time this executor needs
                    to decide whether to throttle futures; this may be used
                    to implement dynamic throttling.

                    .. versionadded:: 2.5.0

            block (bool)
                If ``True``, calls to ``submit()`` on this executor may block if
                there are already ``count`` futures in progress.

                Otherwise, calls to ``submit()`` will always return immediately
                and callables will be queued internally.

            logger (~logging.Logger):
                a logger used for messages from this executor

            name (str):
                a name for this executor

        .. versionchanged:: 2.7.0
            Introduced ``name``.

        .. versionchanged:: 2.11.0
            Introduced ``block``.
        """
        self._log = LogWrapper(
            logger if logger else logging.getLogger("ThrottleExecutor")
        )
        self._block = block
        self._name = name
        self._delegate = delegate
        self._to_submit = deque()
        self._lock = Lock()
        self._event = get_event()
        self._running_count = AtomicInt()
        self._throttle = count if callable(count) else lambda: count
        self._last_throttle = self._throttle()
        self._shutdown = ShutdownHelper()

        event = self._event
        self_ref = weakref.ref(self, lambda _: event.set())

        metrics.EXEC_INPROGRESS.labels(type="throttle", executor=self._name).inc()
        metrics.EXEC_TOTAL.labels(type="throttle", executor=self._name).inc()

        self._thread = Thread(
            name="ThrottleExecutor-%s" % name, target=_submit_loop, args=(self_ref,)
        )
        self._thread.daemon = True
        self._thread.start()

    def submit(self, fn, *args, **kwargs):  # pylint: disable=arguments-differ
        with self._shutdown.ensure_alive():
            self._block_until_ready(self._eval_throttle())

            out = ThrottleFuture(self)
            track_future(out, type="throttle", executor=self._name)

            job = ThrottleJob(out, fn, args, kwargs)
            with self._lock:
                self._to_submit.append(job)
                metrics.THROTTLE_QUEUE.labels(executor=self._name).inc()
                self._log.debug("Enqueued: %s", job)
            self._event.set()
            return out

    def shutdown(self, wait=True, **_kwargs):
        if self._shutdown():
            self._log.debug("Shutting down")
            metrics.EXEC_INPROGRESS.labels(type="throttle", executor=self._name).dec()
            self._delegate.shutdown(wait, **_kwargs)
            self._event.set()
            if wait:
                self._thread.join(MAX_TIMEOUT)

    def _block_until_ready(self, throttle_val):
        while self._block and not self._shutdown.is_shutdown:
            if len(self._to_submit) < throttle_val:
                return
            self._log.debug("%s: throttling on submit", self._name)
            self._event.wait(30.0)

    def _eval_throttle(self):
        try:
            self._last_throttle = self._throttle()
        except Exception:
            self._log.exception(
                "Error evaluating throttle count via %r", self._throttle
            )

        return self._last_throttle

    def _do_submit(self, job):
        delegate_future = self._delegate.submit(job.fn, *job.args, **job.kwargs)
        self._log.debug("Submitted %s yielding %s", job, delegate_future)

        delegate_future.add_done_callback(
            partial(
                self._delegate_future_done, self._log, self._running_count, self._event
            )
        )
        job.future._set_delegate(delegate_future)

    def _do_cancel(self, future):
        with self._lock:
            for job in self._to_submit:
                if job.future is future:
                    self._to_submit.remove(job)
                    self._log.debug("Cancelled %s", job)
                    return True
        self._log.debug("Could not find for cancel: %s", future)
        return False

    @classmethod
    def _delegate_future_done(cls, log, running_count, event, future):
        # Whenever an inner future completes, the thread should wake up
        # in case there's something to be submitted
        log.debug("Delegate future done: %s", future)
        running_count.decr()
        event.set()


def _submit_loop_iter(executor):
    if not executor:
        return

    if executor._shutdown.is_shutdown or is_shutdown():
        return

    throttle = executor._eval_throttle()
    to_submit = []
    with executor._lock:
        while executor._to_submit:
            if throttle is not None and (executor._running_count.value >= throttle):
                executor._log.debug("Throttled")
                break
            job = executor._to_submit.popleft()
            executor._log.debug("Will submit: %s", job)
            to_submit.append(job)

            # While not actually running yet, we've committed to running it, so...
            executor._running_count.incr()
            metrics.THROTTLE_QUEUE.labels(executor=executor._name).dec()

        executor._log.debug(
            "Submitting %s, throttling %s", len(to_submit), len(executor._to_submit)
        )

    for job in to_submit:
        executor._do_submit(job)

    # Because the throttle count is dynamic, we should wake up at some point in
    # the future to re-check throttle, even if no other events occurred.
    # If there's nothing running at all, we should do this sooner, because
    # there won't be any events from completing futures.
    return executor._event, 30.0 if executor._running_count.value else 2.0


@executor_loop
def _submit_loop(executor_ref):
    while True:
        result = _submit_loop_iter(executor_ref())
        if not result:
            break

        (event, wait_time) = result
        event.wait(wait_time)
        event.clear()
