import asyncio
import asyncio.events as events
import os
import sys
import threading
from heapq import heappop


def apply(loop=None):
    """Patch asyncio to make its event loop reentrent."""
    loop = loop or asyncio.get_event_loop()
    if not isinstance(loop, asyncio.BaseEventLoop):
        raise ValueError('Can\'t patch loop of type %s' % type(loop))
    if getattr(loop, '_nest_patched', None):
        # already patched
        return
    _patch_asyncio()
    _patch_loop(loop)
    _patch_task()
    _patch_handle()
    _patch_tornado()


def _patch_asyncio():
    """
    Patch asyncio module to use pure Python tasks and futures,
    use module level _current_tasks, all_tasks and patch run method.
    """
    def run(future, *, debug=False):
        loop = asyncio.get_event_loop()
        loop.set_debug(debug)
        return loop.run_until_complete(future)

    if sys.version_info >= (3, 6, 0):
        asyncio.Task = asyncio.tasks._CTask = asyncio.tasks.Task = \
            asyncio.tasks._PyTask
        asyncio.Future = asyncio.futures._CFuture = asyncio.futures.Future = \
            asyncio.futures._PyFuture
    if sys.version_info < (3, 7, 0):
        asyncio.tasks._current_tasks = asyncio.tasks.Task._current_tasks  # noqa
        asyncio.all_tasks = asyncio.tasks.Task.all_tasks  # noqa
    if not hasattr(asyncio, '_run_orig'):
        asyncio._run_orig = getattr(asyncio, 'run', None)
        asyncio.run = run


def _patch_loop(loop):
    """Patch loop to make it reentrent."""

    def run_forever(self):
        if sys.version_info >= (3, 7, 0):
            set_coro_tracking = self._set_coroutine_origin_tracking
        else:
            set_coro_tracking = self._set_coroutine_wrapper

        self._check_closed()
        old_thread_id = self._thread_id
        old_running_loop = events._get_running_loop()
        set_coro_tracking(self._debug)
        self._thread_id = threading.get_ident()

        if self._asyncgens is not None:
            old_agen_hooks = sys.get_asyncgen_hooks()
            sys.set_asyncgen_hooks(
                firstiter=self._asyncgen_firstiter_hook,
                finalizer=self._asyncgen_finalizer_hook)
        try:
            events._set_running_loop(self)
            while True:
                self._run_once()
                if self._stopping:
                    break
        finally:
            self._stopping = False
            self._thread_id = old_thread_id
            events._set_running_loop(old_running_loop)
            set_coro_tracking(False)
            if self._asyncgens is not None:
                sys.set_asyncgen_hooks(*old_agen_hooks)

    def run_until_complete(self, future):
        old_running_loop = events._get_running_loop()
        try:
            self._check_closed()
            events._set_running_loop(self)
            f = asyncio.ensure_future(future, loop=self)
            if f is not future:
                f._log_destroy_pending = False
            while not f.done():
                self._run_once()
                if self._stopping:
                    break
            if not f.done():
                raise RuntimeError(
                    'Event loop stopped before Future completed.')
            return f.result()
        finally:
            events._set_running_loop(old_running_loop)

    def _run_once(self):
        """
        Simplified re-implementation of asyncio's _run_once that
        runs handles as they become ready.
        """
        now = self.time()
        ready = self._ready
        scheduled = self._scheduled
        while scheduled and scheduled[0]._cancelled:
            heappop(scheduled)

        timeout = (
            0 if ready or self._stopping
            else min(max(0, scheduled[0]._when - now), 86400) if scheduled
            else 0.01 if self._is_proactorloop
            else None)
        event_list = self._selector.select(timeout)
        self._process_events(event_list)

        end_time = self.time() + self._clock_resolution
        while scheduled and scheduled[0]._when < end_time:
            handle = heappop(scheduled)
            ready.append(handle)

        while ready:
            handle = ready.popleft()
            if not handle._cancelled:
                handle._run()
        handle = None

    def _check_running(self):
        """Do not throw exception if loop is already running."""
        pass

    cls = loop.__class__
    cls._run_once_orig = cls._run_once
    cls._run_once = _run_once
    cls._run_forever_orig = cls.run_forever
    cls.run_forever = run_forever
    cls._run_until_complete_orig = cls.run_until_complete
    cls.run_until_complete = run_until_complete
    cls._check_running = _check_running
    cls._check_runnung = _check_running  # typo in Python 3.7 source
    cls._nest_patched = True
    cls._is_proactorloop = (
            os.name == 'nt' and issubclass(cls, asyncio.ProactorEventLoop))


def _patch_task():
    """Patch the Task's step and enter/leave methods to make it reentrant."""

    def step(task, exc=None):
        curr_task = curr_tasks.get(task._loop)
        try:
            step_orig(task, exc)
        finally:
            if curr_task is None:
                curr_tasks.pop(task._loop, None)
            else:
                curr_tasks[task._loop] = curr_task

    Task = asyncio.Task
    if sys.version_info >= (3, 7, 0):

        def enter_task(loop, task):
            curr_tasks[loop] = task

        def leave_task(loop, task):
            curr_tasks.pop(loop, None)

        asyncio.tasks._enter_task = enter_task
        asyncio.tasks._leave_task = leave_task
        curr_tasks = asyncio.tasks._current_tasks
        step_orig = Task._Task__step
        Task._Task__step = step
    else:
        curr_tasks = Task._current_tasks
        step_orig = Task._step
        Task._step = step


def _patch_handle():
    """Patch Handle to allow recursive calls."""

    def update_from_context(ctx):
        """Copy context ctx to currently active context."""
        for var in ctx:
            var.set(ctx[var])

    def run(self):
        """
        Run the callback in a sub-context, then copy any sub-context vars
        over to the Handle's context.
        """
        try:
            ctx = self._context.copy()
            ctx.run(self._callback, *self._args)
            if ctx:
                self._context.run(update_from_context, ctx)
        except (SystemExit, KeyboardInterrupt):
            raise
        except BaseException as exc:
            cb = format_helpers._format_callback_source(
                self._callback, self._args)
            msg = 'Exception in callback {}'.format(cb)
            context = {
                'message': msg,
                'exception': exc,
                'handle': self,
            }
            if self._source_traceback:
                context['source_traceback'] = self._source_traceback
            self._loop.call_exception_handler(context)
        self = None

    if sys.version_info >= (3, 7, 0):
        from asyncio import format_helpers
        events.Handle._run = run


def _patch_tornado():
        """
        If tornado is imported before nest_asyncio, make tornado aware of
        the pure-Python asyncio Future.
        """
        if 'tornado' in sys.modules:
            import tornado.concurrent as tc
            tc.Future = asyncio.Future
            if asyncio.Future not in tc.FUTURES:
                tc.FUTURES += (asyncio.Future,)
