sublime_async.py

import sublime
import sublime_plugin

try:
    import sublime_aio
except:
    sublime_aio = None

import sys
from contextvars import Context
import inspect
import functools
import os
import threading
import logging
import platform
import heapq
import asyncio
import contextlib
import io
import faulthandler
 
from contextlib import contextmanager, redirect_stderr, redirect_stdout
from concurrent.futures import ThreadPoolExecutor, Future as ConcurrentFuture
from typing import IO, Callable, ClassVar, Concatenate, Dict, Generator, Generic, Iterable, List, Optional, ParamSpec, Protocol, Sequence, Set, TextIO, Tuple, Type, TypeVar, Coroutine, Any, TypeVarTuple, Unpack, cast, overload
from asyncio import AbstractEventLoop, EventLoop, Handle, ProactorEventLoop, Future as AsyncIoFuture, SelectorEventLoop, TimerHandle, queues

from asyncio.base_events import _MIN_SCHEDULED_TIMER_HANDLES # type: ignore
from asyncio.base_events import _MIN_CANCELLED_TIMER_HANDLES_FRACTION # type: ignore
from asyncio.base_events import MAXIMUM_SELECT_TIMEOUT # type: ignore

T = TypeVar('T')
Ts = TypeVarTuple("Ts")
U = TypeVar('U')
P = ParamSpec('P')
R = TypeVar('R')
Self = TypeVar('Self')

asyncio_logger = logging.getLogger("asyncio")

logging.basicConfig()
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

if platform.system() != "Windows":
    logger.warning("SublimeAsync has only been tested on Windows. Expect dragons.")

__all__ = [
    "plugin_loaded",

    "get_ui_event_loop",
    "get_background_event_loop",

    "run_in_executor",
    "run_as_task",
    "run_as_task_ui",
    "run_as_task_background",
    "fire_and_forget",
    "with_event_loop",

    "AsyncWindowCommand",
    "AsyncApplicationCommand",

    "run_as_text_command",
    "run_as_window_command",
    "run_as_application_command",

    "SublimeAsyncRunAsTextCommandCommand",
    "SublimeAsyncRunAsWindowCommandCommand",
    "SublimeAsyncRunAsApplicationCommandCommand"
]

_debug_stream: Optional[TextIO] = None
def _enable_debug_logging():
    global _debug_stream
    try:
        _debug_stream = open(r"\\.\pipe\debugpipe", "w", buffering=1, encoding="utf-8")
        _debug_handler = logging.StreamHandler(_debug_stream)
        _debug_handler.setLevel(logging.DEBUG)
        _debug_handler.setFormatter(logging.Formatter(logging.BASIC_FORMAT))
        sys.stderr = _debug_stream
        sys.stdout = _debug_stream
        logging.getLogger().addHandler(_debug_handler)
    except:
        pass

# _enable_debug_logging()

def _enable_fault_handler():
    if _debug_stream is None:
        return
    
    def _reset_faulthandler_watchdog():
        faulthandler.cancel_dump_traceback_later()
        faulthandler.dump_traceback_later(5, repeat=False, file=cast(io.StringIO, _debug_stream))
        sublime.set_timeout(_reset_faulthandler_watchdog, 4000)

    _reset_faulthandler_watchdog()


"""Event loop implementation that runs on Sublime's UI/main thread."""
class SublimeUiEventLoop(ProactorEventLoop):
    def __init__(self) -> None:
        super().__init__()
        self._event_list = None
        self._executor: Optional[ThreadPoolExecutor] = None
        self._next: Optional[Callable[[], None]] = None
        self._scheduling_stopped = True
        self._check_thread_disabled = False
        self._executor_thread_id: int
        self._select_fut: ConcurrentFuture = ConcurrentFuture()
        self._select_fut.set_result(None)

    @property
    def _is_waiting_for_select(self):
        if self._select_fut.done():
            return False
        return True

    @contextmanager
    def _disable_thread_check(self):
        prev = self._check_thread_disabled
        self._check_thread_disabled = True
        try: 
            yield
        finally:
            self._check_thread_disabled = prev

    def _wait_for_wakeup(self) -> None:
        print(f"waiting for wakeupg {self._select_fut}")
        self._select_fut.exception()

    def _wake_up(self) -> None: 
        # NOTE: The wake up is only actually required when waiting for a select(...) in the executor.
        if not self._is_waiting_for_select:
            return

        # NOTE: Do not wake up, if called by the executor thread.
        #       When a self write is received, _poll schedules the 
        #       self read future callback via call_soon, which registers
        #       a new future for the next self write.
        #       If we issue another self write during call_soon, before the callback
        #       has executed, there is no future to handle the self write and 
        #       _poll errors, because it received an unexpected IOCP completion for 
        #       the self write.
        if threading.get_ident() == self._executor_thread_id:
            return

        print("waking up")

        try:
            self._write_to_self() # type: ignore
        except:
            # NOTE: If _write_to_self is not available,
            #       wake up by scheduling a lambda function.
            self.call_soon_threadsafe(lambda: ())   

        print("waiting for wakeup")
        self._wait_for_wakeup()

    def _start_scheduling(self) -> None:
        self._scheduling_stopped = False
        self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="UiLoopSelectThread")
        fut = self._executor.submit(threading.get_ident)
        self._executor_thread_id = fut.result()
        self._schedule_next()

    def _stop_scheduling(self) -> None:
        # NOTE: Shuts down the executor and prevents further automatic 
        #       scheduling of steps via sublime.set_timeout(...).
        #       This *MUST* only be called on the UI thread.
        self._scheduling_stopped = True
        self._wake_up()
        if (self._executor is not None):
            self._executor.shutdown(True, cancel_futures=True)
            self._executor = None

    def call_soon(self, callback: Callable[[Unpack[Ts]], object], *args: Unpack[Ts], context: Context | None = None) -> Handle:
        self._wake_up()
        res = super().call_soon(callback, *args, context=context)
        return res

    def call_at(self, when: float, callback: Callable[[Unpack[Ts]], object], *args: Unpack[Ts], context: Context | None = None) -> TimerHandle:
        self._wake_up()
        res = super().call_at(when, callback, *args, context=context)
        return res

    def _run_next(self) -> None:
        if (self._next is not None):
            logger.debug(f"Running next step: {self._next}.")
            self._next()

    def _check_thread(self) -> None:
        # NOTE: Skip the thread check if disabled (e.g. while running select on a different thread).
        if (not self._check_thread_disabled):
            super()._check_thread() # type: ignore

    def _run_until_run_forever_complete(self) -> None:
        # NOTE: Runs until run_forever completes.
        #       This *MUST* only be called on the UI thread and the loop *MUST* have
        #       been signaled to stop or a pending coroutine or callback that stops the loop
        #       *MUST* have already been scheduled to the loop.
        self._stop_scheduling()
        while (self._next is not None):
            self._run_next()

    def _schedule_next(self) -> None:
        if (self._next is not None and not self._scheduling_stopped):
            sublime.set_timeout(self._run_next, 0)

    def _run_once_begin(self) -> None:
        # NOTE: This is effectively the first part of BaseEventLoop.run_once.
        try:
            sched_count = len(self._scheduled)
            if (sched_count > _MIN_SCHEDULED_TIMER_HANDLES and
                self._timer_cancelled_count / sched_count >
                    _MIN_CANCELLED_TIMER_HANDLES_FRACTION):
                new_scheduled = []
                for handle in self._scheduled:
                    if handle._cancelled:
                        handle._scheduled = False
                    else:
                        new_scheduled.append(handle)

                heapq.heapify(new_scheduled)
                self._scheduled = new_scheduled
                self._timer_cancelled_count = 0
            else:
                while self._scheduled and self._scheduled[0]._cancelled:
                    self._timer_cancelled_count -= 1
                    handle = heapq.heappop(self._scheduled)
                    handle._scheduled = False

            timeout = None
            if self._ready or self._stopping: # type: ignore
                timeout = 0
            elif self._scheduled:
                timeout = self._scheduled[0]._when - self.time()
                if timeout > MAXIMUM_SELECT_TIMEOUT:
                    timeout = MAXIMUM_SELECT_TIMEOUT
                elif timeout < 0:
                    timeout = 0

            self._select(timeout)
        except:
            logger.exception("An exception occured in the UI event loop.")
            self._run_forever_cleanup() # type: ignore

    def _select(self, timeout):
        # NOTE: This wraps _selector.select(...).
        #       select(...) blocks until a coroutine becomes ready to run, or until the timeout expires.
        #       This is run on a separate thread to not block the UI thread.
        #       Once it returns, a single iteration of the event loop is scheduled to the UI thread.
        def do_select(timeout):
            try:
                event_list = None
                # NOTE: When debugging is enabled, call_soon verifies that the caller
                #       is on the same thread  as the event loop and throws otherwise.
                #       Disable the check here since we are running select on a different thread.

                with self._disable_thread_check():
                    event_list = self._selector.select(timeout) # type: ignore

                self._next = lambda: self._run_once_end(event_list)
            except Exception as e:
                logger.exception("An exception occured in the UI event loop.")
                self._next = self._run_forever_cleanup

            self._schedule_next()
        
        if self._select_fut and not self._select_fut.done():
            logger.error("Skip scheduling new select, because the previous one has not completed. Future: %s. Traceback:\n", self._select_fut)
            faulthandler.dump_traceback(_debug_stream, True)
            return

        try:
            if (self._executor is None or self._scheduling_stopped):
                raise RuntimeError("UI event loop executor has been shut down.")
            self._select_fut = self._executor.submit(lambda: do_select(timeout))
            self._select_fut.add_done_callback(lambda t: print(f"{t} is done"))
            logger.debug("Scheduled new select. Future: %s. Traceback:\n", self._select_fut)
            faulthandler.dump_traceback(_debug_stream, True)
        except Exception as ex:
            # NOTE: executor has been shutdown, schedule on the UI thread instead.
            #       This should only happen during shutdown.
            #       When running on the UI thread, we cannot block.
            #       Set the timeout to 0, regardless of what was requested.
            self._next = lambda: do_select(0)
            self._schedule_next()

    def _run_once_end(self, event_list) -> None:
        # NOTE: This is the final part of BaseEventLoop.run_once.

        try:
            self._process_events(event_list) # type: ignore
            # Needed to break cycles when an exception occurs.
            event_list = None

            # Handle 'later' callbacks that are ready.
            end_time = self.time() + self._clock_resolution # type: ignore
            while self._scheduled:
                handle = self._scheduled[0]
                if handle._when >= end_time:
                    break
                handle = heapq.heappop(self._scheduled)
                handle._scheduled = False
                self._ready.append(handle) # type: ignore

            ntodo = len(self._ready) # type: ignore
            for i in range(ntodo):
                handle = self._ready.popleft() # type: ignore
                if handle._cancelled:
                    continue
                if self._debug: # type: ignore
                    try:
                        self._current_handle = handle
                        t0 = self.time()
                        handle._run()
                        dt = self.time() - t0
                        if dt >= self.slow_callback_duration:
                            asyncio_logger.warning('Executing %s took %.3f seconds',
                                           _format_handle(handle), dt) # type: ignore
                    finally:
                        self._current_handle = None
                else:
                    handle._run()
            handle = None

            if not self._stopping: # type: ignore
                self._next = self._run_once_begin
            else:
                self._run_forever_cleanup() # type: ignore
        except:
            logger.exception("An exception occured in the UI event loop.")
            self._run_forever_cleanup() # type: ignore

        self._schedule_next()

    def _run_forever_cleanup(self):
        # NOTE: No more steps to execute.
        logger.info("run forever cleanup")
        self._next = None
        super()._run_forever_cleanup() # type: ignore

    def run_forever(self):
        self.start_run_forever()

        if (self._scheduling_stopped):
            # NOTE: If scheduling via sublime.set_timeout is disabled (e.g. during shutdown),
            #       just run run_forever synchronously.
            #       This is required for run_until_completion to work properly.
            self._run_until_run_forever_complete()
        else:
            self._schedule_next()

    def start_run_forever(self):
        self._run_forever_setup() # type: ignore
        self._next = self._run_once_begin

class TaskSetBase:
    def __init__(self, tag: str = "[unknown]") -> None:
        self._tag = tag

    def add_impl(self, task: Future) -> bool:
        raise NotImplementedError()

    def remove_impl(self, task: Future) -> bool:
        raise NotImplementedError()

    @property
    def pending_tasks(self) -> Sequence[Future]:
        raise NotImplementedError()
    
    def tasks(self) -> List[Future]:
        raise NotImplementedError()

    def ensure_empty(self) -> None:
        tasks = self.pending_tasks
        if len(tasks) > 0:
            logger.warning("Task set %s is not empty.", self._tag)
            if logger.isEnabledFor(logging.DEBUG):
                for t in tasks:
                    logger.debug("Task %s in task set %s is still pending.", t, self._tag)
            raise RuntimeError(f"Task set {self._tag} is not empty.")

    def add(self, task: Future) -> None:
        res = self.add_impl(task)
        if res:
            logger.debug(f"Added future %s to task set %s.", task, self._tag)
        elif logger.isEnabledFor(logging.WARNING):
            logger.warning(f"Failed to add future %s to task set %s.", task, self._tag)

    def remove(self, task: Future) -> None:
        res = self.remove_impl(task)
        if res:
            logger.debug(f"Removed future %s from task set %s.", task, self._tag)
        else:
            logger.warning(f"Failed to remove future %s from task set %s.", task, self._tag)

class ThreadSafeTaskSet(TaskSetBase):
    def __init__(self, tag: str):
        super().__init__(tag)
        self._lock = threading.Lock()
        self._tasks: Set[Future] = set()

    @property
    def pending_tasks(self) -> Sequence[Future]:
        with self._lock:
            return tuple(self._tasks.copy())

    def add_impl(self, task: Future) -> bool:
        with self._lock:
            if task in self._tasks:
                return False
            self._tasks.add(task)
            return True

    def remove_impl(self, task: Future) -> bool:
        with self._lock:
            try:
                self._tasks.remove(task)
                return True
            except KeyError:
                return False

class TaskSet(TaskSetBase):
    def __init__(self, tag: str):
        super().__init__(tag)
        self._tasks: Set[Future] = set()

    @property
    def pending_tasks(self) -> Sequence[Future]:
        return tuple(self._tasks.copy())
        
    def add_impl(self, task: Future) -> bool:
        if task in self._tasks:
            return False
        self._tasks.add(task)
        return True

    def remove_impl(self, task: Future) -> bool:
        try:
            self._tasks.remove(task)
            return True
        except KeyError:
            return False

class EventLoopManagerBase:
    _threadsafe_tasks: ClassVar[TaskSetBase] = ThreadSafeTaskSet("fallback")

    def start(self) -> None:
        pass

    def stop(self) -> None:
        pass

    def plugin_loaded(self) -> None:
        pass

    @classmethod
    def get_manager(cls, loop: Optional[AbstractEventLoop]) -> Optional[EventLoopManagerBase]:
        return getattr(loop, "_owning_manager", None)

    @classmethod
    def get_tasks(cls, loop: Optional[AbstractEventLoop]) -> TaskSetBase:
        mgr = cls.get_manager(loop)
        return mgr.tasks if mgr else cls._threadsafe_tasks

    @property
    def loop(self) -> AbstractEventLoop:
        raise NotImplementedError()

    @property
    def tasks(self) -> TaskSetBase:
        return self._threadsafe_tasks


# NOTE: FIXME: This is a hack to reuse sublime_aio's event loop as the background loop if present
class SublimeAioEventLoopManager(EventLoopManagerBase):
    def __init__(self):
        self.loop._owning_manager = self # type: ignore
        self._tasks = TaskSet("sublime aio background")

    def stop(self):
        # NOTE: This is the shutdown procedure of sublime_aio's event loop.
        #       In sublime_aio, this is part of their custom on_exit handler.
        assert(sublime_aio)
        sublime_aio.ExitEvent.wait()
        if sublime_aio._loop is not None and sublime_aio._thread is not None:
            loop = sublime_aio._loop
            sublime_aio._loop = None

            def shutdown():
                for task in asyncio.all_tasks(loop):
                    task.cancel()
                loop.stop()

            loop.call_soon_threadsafe(shutdown)
            sublime_aio._thread.join()
            loop.run_until_complete(loop.shutdown_asyncgens())
            loop.close()
        self._tasks.ensure_empty()

    @property
    def loop(self) -> AbstractEventLoop:
        assert(sublime_aio)
        if (sublime_aio._loop is None):
            raise RuntimeError("sublime_aio event loop is not initialized.")
        return sublime_aio._loop

    @property
    def tasks(self) -> TaskSetBase:
        return self._tasks
    
class AsyncioEventLoopManager(EventLoopManagerBase):
    def __init__(self):
        self._executor = ThreadPoolExecutor(max_workers=16)
        self._loop = asyncio.EventLoop()
        self._loop.set_debug(True)
        self._loop.set_default_executor(self._executor)
        self._thread = threading.Thread(target=self._loop.run_forever)
        self._loop._owning_manager = self # type: ignore
        self._tasks = TaskSet("background")

    def start(self):
        self._thread.start()

    def stop(self):
        def shutdown():
            for task in asyncio.all_tasks(self._loop):
                task.cancel()
            self._loop.stop()
        self._loop.call_soon_threadsafe(shutdown)
        self._thread.join()
        self._loop.run_until_complete(self._loop.shutdown_asyncgens())
        self._loop.close()
        self._executor.shutdown()
        self._tasks.ensure_empty()

    @property
    def loop(self) -> AbstractEventLoop:
        return self._loop

    @property
    def tasks(self) -> TaskSetBase:
        return self._tasks
    
class SublimeUiEventLoopManager(EventLoopManagerBase):
    def __init__(self):
        self._loop = SublimeUiEventLoop()
        self._loop.set_debug(True)
        self._loop._owning_manager = self # type: ignore
        self._tasks = TaskSet("ui")

    def start(self):
        # NOTE: This will also set the event loop as the running event loop
        #       for the calling thread. This *MUST* be called on the UI thread.
        self._loop.start_run_forever()

    def plugin_loaded(self):
        logger.debug("Plugin loaded. Start scheduling UI loop steps.")
        self._loop._start_scheduling()

    def stop(self):
        # NOTE: At this point, sublime API is no longer available.
        #       Remainder of the loop is run synchronously instead of through sublime.set_timeout.

        def shutdown():
            for task in asyncio.all_tasks(self._loop):
                task.cancel()
            # self._loop.stop()
            # logging.info("shutdown proc complete")

        print("stop scheduling")
        self._loop._stop_scheduling()
        print("schedule shutdown")
        # self._loop.call_soon(shutdown)
        # logging.info("run until complete")
        # self._loop._run_until_run_forever_complete()
        print(f"{self._loop}")
        # self._loop.run_until_complete(self._loop.shutdown_asyncgens())
        # self._loop.close()
        # self._tasks.ensure_empty()

    @property
    def loop(self) -> AbstractEventLoop:
        return self._loop

    @property
    def tasks(self) -> TaskSetBase:
        return self._tasks

class EventLoopManager:
    def __init__(self):
        self._ui_loop_mgr = SublimeUiEventLoopManager()
        self._background_loop_mgr: EventLoopManagerBase
        # if sublime_aio:
        #     logger.info("sublime_aio installed. Use its event loop as background loop.")
        #     self._background_loop_mgr = SublimeAioEventLoopManager()
        # else:
        #     self._background_loop_mgr = AsyncioEventLoopManager()

    @property
    def ui_event_loop(self):
        return self._ui_loop_mgr.loop

    @property
    def background_event_loop(self):
        pass
        # return self._background_loop_mgr.loop

    def start(self):
        self._ui_loop_mgr.start()
        # self._background_loop_mgr.start()

    def plugin_loaded(self):
        self._ui_loop_mgr.plugin_loaded()
        # self._background_loop_mgr.plugin_loaded()

    def stop(self):
        try:
            self._ui_loop_mgr.stop()
            logger.info(f"UI event loop shut down.")
        except Exception as ex:
            logger.exception(f"An exception occured while shutting down the UI event loop: {ex}.")

        # NOTE: Background loop must be stopped last.
        #       We can't call run_until_complete otherwise, because the UI loop is still running.
        # try:
        #     self._background_loop_mgr.stop()
        #     logger.info(f"Background event loop shut down.")
        # except Exception as ex:
        #     logger.exception(f"An exception occured while shutting down the background event loop: {ex}.")


_event_loop_manager: EventLoopManager

if globals().get("_event_loop_manager"):
    logger.info("Event loops already started. Shutting down first.")
    _event_loop_manager.stop() # type: ignore
    logger.info("Event loops shut down.")

# _event_loop_manager = EventLoopManager()
# _event_loop_manager.start()

_counter = 0

def plugin_loaded():
    global _counter
    print(f"count {_counter}")
    _counter += 1
    _enable_fault_handler()
    # _event_loop_manager.plugin_loaded()

class TeeStream(io.TextIOBase, TextIO):
    def __init__(self, *streams: IO[str], close_streams: bool = True) -> None:
        super().__init__()
        self._streams: List[IO[str]] = list(streams)
        self._close_streams = close_streams

    def write(self, data: str) -> int:
        for s in self._streams:
            s.write(data)
        return len(data)

    def flush(self) -> None:
        for s in self._streams:
            s.flush()

    def writable(self) -> bool:
        return True

    def readable(self) -> bool:
        return False

    def seekable(self) -> bool:
        return False

    def read(self, *args, **kwargs):
        raise io.UnsupportedOperation("TeeStream is write-only.")

    def close(self) -> None:
        if self._close_streams:
            for s in self._streams:
                s.close()
        super().close()

    def __add__(self, stream: IO[str]) -> TeeStream:
        if not stream.writable:
            raise io.UnsupportedOperation("Stream is not writable.")
        return TeeStream(*self._streams, stream)

    def __iadd__(self, stream: IO[str]) -> TeeStream:
        if not stream.writable:
            raise io.UnsupportedOperation("Stream is not writable.")
        self._streams.append(stream)
        return self

def on_exit(log_path: str):
    # NOTE: Custom exit handler to properly shut down the event loops.
    #       This overrides the one registered by sublime_aio, if present.
    #       The sublime_aio event loop will be shut down via _background_loop_manager.stop(), see above.

    stdout = io.StringIO()
    redirect_stream = TeeStream(stdout)

    if _debug_stream:
        redirect_stream += _debug_stream

    root = logging.getLogger()
    handler = logging.StreamHandler(stdout)
    handler.setFormatter(logging.Formatter(logging.BASIC_FORMAT))
    root.addHandler(handler)

    with contextlib.redirect_stdout(redirect_stream), contextlib.redirect_stderr(redirect_stream):
        logger.info("Executing on_exit.")

        for callback in sublime_plugin.el_callbacks('on_exit'):
            callback()

        _event_loop_manager.stop()

        handler.flush()
        redirect_stream.flush()

        if len(stdout.getvalue()):
            with open(log_path, "w", encoding="utf-8") as f:
                f.write(stdout.getvalue())
        else:
            os.unlink(log_path)

# Register the custom exit handler.
sublime_plugin.on_exit = on_exit

def get_ui_event_loop() -> AbstractEventLoop:
    """
    Returns the UI event loop.

    Returns:
      AbstractEventLoop:
        The UI event loop.
    """
    return _event_loop_manager.ui_event_loop

def get_background_event_loop() -> AbstractEventLoop:
    """
    Returns the background event loop.

    Returns:
      AbstractEventLoop:
        The background event loop.
    """
    return _event_loop_manager.background_event_loop

class Future[T]:
    """
    Wrapper for `AsyncIoFuture` and `ConcurrentFuture` objects to provide a unified interface for asynchronous operations.
    
    Raises:
      RuntimeError:
        If an operation is attempted that is not available on the wrapped Future.
    """
    def __init__(self, future: AsyncIoFuture[T] | ConcurrentFuture[T]):
        setattr(future, "_owner", self)
        self._future: AsyncIoFuture[T] | ConcurrentFuture[T] = future

    def __repr__(self) -> str:
        return self._future.__repr__()

    def __await__(self) -> Generator[Any, None, T]:
        if isinstance(self._future, AsyncIoFuture):
            return self._future.__await__()
        raise RuntimeError("Future Object is not an asyncio Future.")

    def __iter__(self) -> Generator[Any, None, T]:
        if isinstance(self._future, AsyncIoFuture):
            return self._future.__iter__()
        raise RuntimeError("Future Object is not an asyncio Future.")

    def cancel(self) -> None:
        self._future.cancel()

    def get_loop(self) -> AbstractEventLoop:
        if isinstance(self._future, AsyncIoFuture):
            return self._future.get_loop()
        raise RuntimeError("Future Object is not an asyncio Future.")

    def try_get_loop(self) -> Optional[AbstractEventLoop]:
        if isinstance(self._future, AsyncIoFuture):
            return self._future.get_loop()
        return None

    def cancelled(self) -> bool:
        return self._future.cancelled()

    def running(self) -> bool:
        if isinstance(self._future, ConcurrentFuture):
            return self._future.running()
        raise RuntimeError("Future Object is an asyncio Future.")

    def done(self) -> bool:
        return self._future.done()

    def result(self, timeout=None) -> T:
        if isinstance(self._future, AsyncIoFuture):
            return self._future.result()
        return self._future.result(timeout=timeout)

    def exception(self, timeout=None) -> Optional[BaseException]:
        if isinstance(self._future, AsyncIoFuture):
            return self._future.exception()
        return self._future.exception(timeout=timeout)

    def add_done_callback(self, fn: Callable[[Future], object], *, context=None) -> None:
        if isinstance(self._future, AsyncIoFuture):
            self._future.add_done_callback(lambda f: fn(getattr(f, "_owner")), context=context)
        else:
            self._future.add_done_callback(lambda f: fn(getattr(f, "_owner")))

    def remove_done_callback(self, fn) -> None:
        # NOTE: Because added callbacks are wrapped in a lambda, we can't really remove them.
        raise RuntimeError("remove_done_callback is not supported.")

    def set_result(self, result: T) -> None:
        self._future.set_result(result)

    def set_exception(self, exception: Exception) -> None:
        self._future.set_exception(exception)

@overload
def run_as_task(loop_or_factory: Callable[[], AbstractEventLoop] = asyncio.get_running_loop) -> \
    Callable[[Callable[P, Coroutine[Any, Any, T] | T]], Callable[P, Future[T]]]:
    ...

@overload
def run_as_task(loop_or_factory: AbstractEventLoop) -> \
    Callable[[Callable[P, Coroutine[Any, Any, T] | T]], Callable[P, Future[T]]]:
    ...

def run_as_task(loop_or_factory: Callable[[], AbstractEventLoop] | AbstractEventLoop = asyncio.get_running_loop) -> \
    Callable[[Callable[P, Coroutine[Any, Any, T] | T]], Callable[P, Future[T]]]:
    """
    Decorator to run a function or coroutine as an asynchronous task in a an event loop.

    Parameters:
      loop_or_factory (Callable[[], AbstractEventLoop] | AbstractEventLoop):
        An event loop or a factory function that returns an event loop. 
        Defaults to the currently running loop.
    
    Returns:
      Callable[[Callable[..., Coroutine[Any, Any, T] | T]], Callable[..., Future[T]]]:
        A decorator that wraps the input function to run it as an asynchronous task, returning a `Future`.
    
    Usage Example:
      ```python
      @run_as_task()
      async def my_async_function() -> int:
          await asyncio.sleep(1)
          print("Hello from my_async_function!")
          return 1

      async def main():
        task = my_async_function()
        print("Hello from main!")
        return await task
    
      asyncio.runt(main())
    """
    if isinstance(loop_or_factory, AbstractEventLoop):
        loop_factory = lambda: loop_or_factory
    else:
        loop_factory = loop_or_factory
    def decorator(func: Callable[P, Coroutine[Any, Any, T] | T]) -> Callable[P, Future[T]]:
        coro_func: Callable[..., Coroutine[Any, Any, T]]
        if (asyncio.iscoroutinefunction(func)):
            coro_func = func
        else:
            async def coro_wrapper(*args, **kwargs) -> T:
                return cast(Callable[P, T], func)(*args, **kwargs)
            coro_func = coro_wrapper

        @functools.wraps(func)
        def wrapper(*args, **kwargs) -> Future[T]:
            loop = loop_factory()
            running_loop = asyncio._get_running_loop()
            coro = coro_func(*args, **kwargs)

            if (running_loop is not None):
                if (loop is running_loop):
                    logger.debug
                    res = Future(loop.create_task(coro))
                else:
                    res = Future(asyncio.wrap_future(asyncio.run_coroutine_threadsafe(coro, loop)))
            else: 
                res = Future(asyncio.run_coroutine_threadsafe(coro, loop))
            logger.debug("Running %s as task on event loop %s, scheduled from %s. Future: %s", func, loop, running_loop, res)
            return res
        return wrapper
    return decorator

def run_as_task_ui(func: Callable[P, Coroutine[Any, Any, T]]) -> Callable[P, Future[T]]:
    return run_as_task(get_ui_event_loop)(func)

def run_as_task_background(func: Callable[P, Coroutine[Any, Any, T]]) -> Callable[P, Future[T]]:
    return run_as_task(get_background_event_loop)(func)

def fire_and_forget(func: Callable[P, Future[T]]) -> Callable[P, None]:
    """
    A decorator to run a coroutine in a fire and forget manner.
    The function's result (or exception) is ignored.

    Parameters:
      func (Callable[..., Future[T]]):
        The async function to be wrapped. It must return a Future object.

    Returns:
      Callable[..., None]:
        A wrapper function that executes the original function in a fire-and-forget manner.

    Example:
      ```python
      @fire_and_forget
      async def background_task():
          await asyncio.sleep(5)
          print("Task completed")

      background_task()  # Runs in the background without blocking
      """
    @functools.wraps(func)
    def wrapper(*args, **kwargs) -> None:
        fut = func(*args, **kwargs)
        loop = fut.try_get_loop()
        tasks = EventLoopManagerBase.get_tasks(loop)
        
        logger.debug("Running task %s as fire and forget.", fut)

        tasks.add(fut)

        def done_callback(f: Future):
            exc = fut.exception()
            if exc is not None:
                logger.error("Fire and forget task %s has failed: %s.", fut, exc, exc_info=(type(exc), exc, exc.__traceback__))
            else:
                logger.debug("Fire and forget task %s has completed successfully.", fut)
            tasks.remove(fut)

        fut.add_done_callback(done_callback)
    return wrapper

def run_in_executor(loop_or_factory: Callable[[], AbstractEventLoop] | AbstractEventLoop = asyncio.get_running_loop) -> \
    Callable[[Callable[P, T]], Callable[P, Future[T]]]:
    """
    A decorator to run a synchronous function in an executor and return a Future.

    Parameters:
        loop_or_factory (Callable[[], AbstractEventLoop] | AbstractEventLoop): 
            Either an event loop or a factory function that returns an event loop. 
            Defaults to the currently running event loop.

    Returns:
        Callable[[Callable[..., T]], Callable[..., Future[T]]]: 
            A decorator that wraps the input function to run it in an executor.

    Raises:
        RuntimeError: If the decorated function is a coroutine.

    Example:
        ```python
        @run_in_executor()
        def blocking_io_task():
            time.sleep(2)
            return "Done"

        async def main():
            future = blocking_io_task()
            result = await future
            print(result)  # Output: "Done"

        asyncio.run(main())
    """
    if isinstance(loop_or_factory, AbstractEventLoop):
        loop_factory = lambda: loop_or_factory
    else:
        loop_factory = loop_or_factory
    def decorator(func: Callable[P, T]) -> Callable[P, Future[T]]:
        if (asyncio.iscoroutinefunction(func)):
            raise RuntimeError("run_in_executor does not support coroutines. Use run_as_task instead.")
        @functools.wraps(func)
        def wrapper(*args, **kwargs) -> Future[T]:
            loop = loop_factory()
            running_loop = asyncio._get_running_loop()

            if (running_loop is not None):
                res = Future(asyncio.wrap_future(loop.run_in_executor(None, lambda: func(*args, **kwargs))))
            else:
                res = Future(loop.run_in_executor(None, lambda: func(*args, **kwargs)))
            logger.debug("Running %s in executor of loop %s. Future: %s.", func, loop, res)
            return res
        return wrapper
    return decorator

class CommandHandlerDict(Generic[U]):
    def __init__(self, tag: str):
        self._dict: Dict[int, U] = dict()
        self._lock = threading.Lock()
        self._tag = tag

    def pop(self, key: int, default: Optional[U] = None) -> Optional[U]:
        with self._lock:
            return self._dict.pop(key, default)

    def get(self, key: int, default: Optional[U] = None) -> Optional[U]:
        with self._lock:
            return self._dict.get(key, default)

    @property
    def pending_handlers(self) -> Dict[int, U]:
        with self._lock:
            return self._dict.copy()

    def ensure_empty(self) -> None:
        pending_handlers = self.pending_handlers.items()

        if len(pending_handlers) > 0:
            logger.warning("Command handler dict %s is not empty.", self._tag)
            if logger.isEnabledFor(logging.DEBUG):
                for k, v in pending_handlers:
                    logger.debug("Pending handler %s (handler id %s) in command handler dict %s was never executed", v, k, self._tag)

    def __setitem__(self, key: int, value: U) -> None:
        with self._lock:
            if key in self._dict:
                logger.warning("Failed to add command handler %s (handler id %s) to command handler dict %s.", value, key, self._tag)
            else:
                self._dict[key] = value
                logger.debug("Added command handler %s (handler id %s) to command handler dict %s.", value, key, self._tag)

    def __getitem__(self, key: int) -> U:
        with self._lock:
            return self._dict[key]

class HasViewProperty(Protocol):
    @property
    def view(self) -> sublime.View:
        ...

HasViewPropertyT = TypeVar('HasViewPropertyT', bound="HasViewProperty")

@overload
def run_as_text_command(func: Callable[Concatenate[HasViewPropertyT, sublime.Edit, sublime.View, P], R]) -> Callable[Concatenate[HasViewPropertyT, P], Future[R]]:
    ...

@overload
def run_as_text_command(func: Callable[Concatenate[Self, sublime.Edit, sublime.View, P], R]) -> Callable[Concatenate[Self, sublime.View, P], Future[R]]:
    ...

@overload
def run_as_text_command(func: Callable[Concatenate[sublime.Edit, sublime.View, P], R]) -> Callable[Concatenate[sublime.View, P], Future[R]]:
    ...

def run_as_text_command(func: Callable[..., R]) -> Callable[..., Future[R]]:
    """
    Decorator that converts a function or method into a text command.
    Executes the decorated function within in a text command context, providing a valid
    `sublime.Edit` instance, and returns a `Future` that resolves with the function's return value once
    the text command has completed.

    Parameters:
      `func` (`Callable[[sublime.Edit, sublime.View, ...], R]`):
        The function or method to wrap. Expects the first two parameters to be `sublime.Edit`
        and `sublime.View`.

    Returns:
      `Callable[[sublime.View, ...], Future[R]]`:
        A wrapper function that accepts the `sublime.View` instance on which to run the text command and remaining arguments,
        executes the the orignal function or method in a text command context and returns a Future that resolves 
        once the text command has completed.
        If the wrapped function is a method of a class that has a `view` field of type `sublime.View`, the view
        parameter is not required.

    Example:
    ```python
      @run_as_text_command
      def insert_text(edit: sublime.Edit, view: sublime.View, text: str) -> int:
          return view.insert(edit, 0, text)
      future = insert_text(view, "Hello World")

      class Foo:
        @property
        def view(self) -> sublime.View:
            return self._view

        @run_as_text_command
        def insert_text(self, edit: sublime.Edit, view: sublime.View, text: str) -> int:
            return view.insert(edit, 0, text)
      future = Foo().insert_text("Hello World")
    """
    sig = inspect.signature(func)
    params = list(sig.parameters.values())
    param0 = params[0]
    is_method = param0.name == "self"

    def _run(f: Callable[[sublime.Edit, sublime.View], R], view: sublime.View, *args, **kwargs) -> Tuple[Future[R], int]:
        running_loop = asyncio.get_running_loop()

        fut: AsyncIoFuture[R] | ConcurrentFuture[R]
        res: Future[R]
        if (running_loop is not get_ui_event_loop()):
            # NOTE: Called without running event loop or from a thread different from the UI thread.
            fut = ConcurrentFuture()
            if (running_loop is None):
                res = Future(fut)
            else:
                res = Future(asyncio.wrap_future(fut))
        else:
            # NOTE: Command will be run on the UI thread.
            #       If called from UI thread, we can safely use an asyncio Future.
            #       This also prevents synchronously waiting for the future to
            #       avoid deadlocks.
            fut = AsyncIoFuture()
            res = Future(fut)

        def handler(edit: sublime.Edit, view: sublime.View) -> None:
            try: 
                res = f(edit, view, *args, **kwargs)
                fut.set_result(res)
            except Exception as ex:
                fut.set_exception(ex)

        handler_id = SublimeAsyncRunAsTextCommandCommand.register_handler(handler)
        view.run_command("sublime_async_run_as_text_command", {"_handler_id": handler_id})
        return (res, handler_id)

    if (is_method):
        def method_wrapper(self, *args, **kwargs) -> Future[R]:
            if (hasattr(self, "view") and isinstance(self.view, sublime.View)):
                res = _run(functools.partial(func, self), self.view, *args, **kwargs)
            else:
                res = _run(functools.partial(func, self), *args, **kwargs)
            logger.debug("Running %s as text command (handler id %s). Future: %s.", func, res[1], res[0])
            return res[0]
        return method_wrapper
    else:
        def wrapper(view: sublime.View, *args, **kwargs) -> Future[R]:
            res = _run(func, view, *args, **kwargs)
            logger.debug("Running %s as text command (handler id %s). Future: %s.", func, res[1], res[0])
            return res[0]
        return wrapper

class HasWindowProperty(Protocol):
    @property
    def window(self) -> sublime.Window:
        ...

HasWindowPropertyT = TypeVar("HasWindowPropertyT", bound=HasWindowProperty)

@overload
def run_as_window_command(func: Callable[Concatenate[HasWindowPropertyT, sublime.Window, P], R]) -> Callable[Concatenate[HasWindowPropertyT, P], Future[R]]:
    ...

@overload
def run_as_window_command(func: Callable[Concatenate[Self, sublime.Window, P], R]) -> Callable[Concatenate[Self, sublime.Window, P], Future[R]]:
    ...

@overload
def run_as_window_command(func: Callable[Concatenate[sublime.Window, P], R]) -> Callable[Concatenate[sublime.Window, P], Future[R]]:
    ...

def run_as_window_command(func: Callable[..., R]) -> Callable[..., Future[R]]:
    """
    Decorator that converts a function or method into a window command.
    Executes the decorated function within in a window command context, 
    and returns a `Future` that resolves with the function's return value once
    the text command has completed.

    Parameters:
      `func` (`Callable[[sublime.Window, ...], R]`):
        The function or method to wrap. Expects the first parameters to be `sublime.Window`.

    Returns:
      `Callable[[sublime.Window, ...], Future[R]]`:
        A wrapper function that accepts the `sublime.Window` instance on which to run the window command and remaining arguments,
        executes the the orignal function or method in a window command context and returns a Future that resolves 
        once the text command has completed.
        If the wrapped function is a method of a class that has a `window` field of type `sublime.Window`, the window
        parameter is not required.

    Example:
    ```python
      @run_as_window_command
      def to_front(window: sublime.Window, text: str) -> None:
          window.bring_to_front()
      future = to_front(window)

      class Foo:
        @property
        def window(self) -> sublime.Window:
            return self._window

        @run_as_window_command
        def to_front(self, window: sublime.Window) -> None:
            window.bring_to_front()
      future = Foo().to_front()
    """
    sig = inspect.signature(func)
    params = list(sig.parameters.values())
    param0 = params[0]
    is_method = param0.name == "self"

    def _run(f: Callable[[sublime.Window], R], window: sublime.Window, *args, **kwargs) -> Tuple[Future[R], int]:
        running_loop = asyncio._get_running_loop()
        fut: AsyncIoFuture[R] | ConcurrentFuture[R]
        res: Future[R]
        if (running_loop is not get_ui_event_loop()):
            # NOTE: Called without running event loop or from a thread different from the UI thread.
            fut = ConcurrentFuture()
            if (running_loop is None):
                res = Future(fut)
            else:
                res = Future(asyncio.wrap_future(fut))
        else:
            # NOTE: Command will be run on the UI thread.
            #       If called from UI thread, we can safely use an asyncio Future.
            #       This also prevents synchronously waiting for the future to
            #       avoid deadlocks.
            fut = AsyncIoFuture()
            res = Future(fut)

        def handler(window: sublime.Window) -> None:
            try: 
                res = f(window, *args, **kwargs)
                fut.set_result(res)
            except Exception as ex:
                fut.set_exception(ex)

        handler_id = SublimeAsyncRunAsWindowCommandCommand.register_handler(handler)
        window.run_command("sublime_async_run_as_window_command", {"_handler_id": handler_id})
        return (res, handler_id)

    if (is_method):
        def method_wrapper(self, *args, **kwargs) -> Future[R]:
            if (hasattr(self, "window") and isinstance(self.window, sublime.Window)):
                res = _run(functools.partial(func, self), self.window, *args, **kwargs)
            else:
                res = _run(functools.partial(func, self), *args, **kwargs)
            logger.debug("Running %s as window command (handler id %s). Future: %s.", func, res[1], res[0])
            return res[0]
        return method_wrapper
    else:
        def wrapper(window: sublime.Window, *args, **kwargs) -> Future[R]:
            res = _run(func, window, *args, **kwargs)
            logger.debug("Running %s as window command (handler id %s). Future: %s.", func, res[1], res[0])
            return res[0]
        return wrapper


def run_as_application_command(func: Callable[P, R]) -> Callable[P, Future[R]]:
    """
    Decorator that converts a function or method into an application command.
    Executes the decorated function within in an application command context, 
    and returns a `Future` that resolves with the function's return value once
    the application command has completed.

    Parameters:
      `func` (`Callable[..., R]`):
        The function or method to wrap.

    Returns:
      `Callable[..., Future[R]]`:
        A wrapper function that executes the the orignal function or method in an application
        command context and returns a Future that resolves once the application command has completed.

    Example:
    ```python
      @run_as_application_command
      def hello_world() -> None:
        print("Hello world")
      future = hello_world()
    """
    @functools.wraps(func)
    def wrapper(*args, **kwargs) -> Future[R]:
        running_loop = asyncio._get_running_loop()

        fut: AsyncIoFuture[R] | ConcurrentFuture[R]
        res: Future[R]
        if (running_loop is not get_ui_event_loop()):
            # NOTE: Called without running event loop or from a thread different from the UI thread.
            fut = ConcurrentFuture()
            if (running_loop is None):
                res = Future(fut)
            else:
                res = Future(asyncio.wrap_future(fut))
        else:
            # NOTE: Command will be run on the UI thread.
            #       If called from UI thread, we can safely use an asyncio Future.
            #       This also prevents synchronously waiting for the future to
            #       avoid deadlocks.
            fut = AsyncIoFuture()
            res = Future(fut)

        def handler():
            try:
                res = func(*args, **kwargs)
                fut.set_result(res)
            except Exception as ex:
                fut.set_exception(ex)

        handler_id = SublimeAsyncRunAsApplicationCommandCommand.register_handler(handler)
        sublime.run_command("sublime_async_run_as_application_command", {"_handler_id": handler_id})
        logger.debug("Running %s as application command (handler id %s). Future: %s.", func, handler_id, res)
        return res
    return wrapper

class SublimeAsyncRunAsTextCommandCommand(sublime_plugin.TextCommand):
    _pending: ClassVar[CommandHandlerDict[Callable[[sublime.Edit, sublime.View], Any]]] = CommandHandlerDict("text command")

    @classmethod
    def register_handler(cls, handler: Callable[[sublime.Edit, sublime.View], Any]) -> int:
        id = int.from_bytes(os.urandom(4))
        cls._pending[id] = handler
        return id

    def run(self, edit: sublime.Edit, *, _handler_id: sublime.Value, **kwargs: sublime.Value):
        handler = self._pending.pop(cast(int, _handler_id))
        if(handler is not None):
            logger.debug("Executing text command handler %s for handler id %s.", _handler_id)
            handler(edit, self.view)
        else:
            logger.warning("No text command handler registered for handler id %s.", _handler_id)

class SublimeAsyncRunAsWindowCommandCommand(sublime_plugin.WindowCommand):
    _pending: ClassVar[CommandHandlerDict[Callable[[sublime.Window], Any]]] = CommandHandlerDict("window command")

    @classmethod
    def register_handler(cls, handler: Callable[[sublime.Window], Any]) -> int:
        id = int.from_bytes(os.urandom(4))
        cls._pending[id] = handler
        return id

    def run(self, *, _handler_id: sublime.Value, **kwargs: sublime.Value):
        handler = self._pending.pop(cast(int, _handler_id))
        if (handler is not None):
            logger.debug("Executing window command handler %s for handler id %s.", _handler_id)
            handler(self.window)
        else:
            logger.warning("No window command handler registered for handler id %s.", _handler_id)

class SublimeAsyncRunAsApplicationCommandCommand(sublime_plugin.ApplicationCommand):
    _pending: ClassVar[CommandHandlerDict[Callable[[], Any]]] = CommandHandlerDict("application command")

    @classmethod
    def register_handler(cls, handler: Callable[[], Any]) -> int:
        id = int.from_bytes(os.urandom(4))
        cls._pending[id] = handler
        return id

    def run(self, *, _handler_id: sublime.Value, **kwargs: sublime.Value):
        handler = self._pending.pop(cast(int, _handler_id))
        if (handler is not None):
            logger.debug("Executing application command handler %s for handler id %s.", _handler_id)
            handler()
        else:
            logger.warning("No application command handler registered for handler id %s.", _handler_id)

def with_event_loop(loop_or_factory: Callable[[], AbstractEventLoop] | AbstractEventLoop):
    if isinstance(loop_or_factory, AbstractEventLoop):
        loop_factory = lambda: loop_or_factory
    else:
        loop_factory = loop_or_factory

    def decorator(cls: Type[AsyncCommandT]) -> Type:
        setattr(cls, "_loop_factory", staticmethod(loop_factory))
        return cls
    return decorator

@with_event_loop(asyncio.get_running_loop)
class AsyncCommand:
    _loop_factory: ClassVar[Callable[[], AbstractEventLoop]] = staticmethod(asyncio.get_running_loop)

    async def run(self, **kwargs) -> None:
        raise NotImplementedError()

    def __init_subclass__(cls, **kwargs) -> None:
        super().__init_subclass__(**kwargs)

        original_run: Optional[Callable[..., Coroutine[Any, Any, None]]] = cls.__dict__.get("run")
        if original_run is None:
            return

        @functools.wraps(original_run)
        def wrapper(self: AsyncCommand, **kwargs) -> None:
            # loop: AbstractEventLoop = type(self)._loop_factory()

            # fut = run_as_task(self._loop_factory)(original_run)(self, **kwargs)

            fut = Future(asyncio.run_coroutine_threadsafe(original_run(self, **kwargs), self._loop_factory()))

            def handle_exception(fut: Future) -> None:
                print("handle exception 11")
                if fut.cancelled():
                    return
                exc = fut.exception()
                if (exc is not None):
                    logger.error(f"Async command failed: {exc}", exc_info=(type(exc), exc, exc.__traceback__))

            fut.add_done_callback(handle_exception)
            print(fut.done)
            fire_and_forget(lambda: fut)

        setattr(cls, "run", wrapper)
        setattr(cls, "_original_run", original_run)

# NOTE: For some reason, with these command classes \/ \/ \/ plugin_loaded is called twice. If commented out, its fine?

class AsyncWindowCommand(AsyncCommand, sublime_plugin.WindowCommand): # type: ignore
    pass

class AsyncApplicationCommand(AsyncCommand, sublime_plugin.ApplicationCommand): # type: ignore
    pass

AsyncCommandT = TypeVar('AsyncCommandT', bound=AsyncCommand)