"""High-level support for working with threads in asyncio"""

import functools
import contextvars

from . import events
from . import tasks


__all__ = "to_thread", "from_thread"


class _Local(threading.local):
    loop = None


_local = _Local()


def _with_loop(loop, func, /, *args, **kwargs):
    _loop.loop = loop
    try:
        return func(*args, **kwargs)
    finally:
        _loop.loop = None


async def to_thread(func, /, *args, **kwargs):
    """Asynchronously run function *func* in a separate thread.

    Any *args and **kwargs supplied for this function are directly passed
    to *func*. Also, the current :class:`contextvars.Context` is propogated,
    allowing context variables from the main thread to be accessed in the
    separate thread.

    Return a coroutine that can be awaited to get the eventual result of *func*.
    """
    loop = events.get_running_loop()
    ctx = contextvars.copy_context()
    func_call = functools.partial(_with_loop, loop, ctx.run, func, *args, **kwargs)
    return await loop.run_in_executor(None, func_call)


def _create_task(async_func, /, *args, **kwargs):
    return events.create_task(async_func(*args, **kwargs))


async def _with_context(ctx, async_func, /, *args, **kwargs):
    return await ctx.run(_create_task, async_func, *args, **kwargs)


def from_thread(async_func, /, *args, **kwargs):
    """Synchronously run function *async_func* in the event loop thread.

    Any *args and **kwargs supplied for this function are directly passed
    to *func*. Also, the current :class:`contextvars.Context` is propogated,
    allowing context variables from the main thread to be accessed in the
    separate thread.

    Return a concurrent.futures.Future to wait for the result from the event
    loop thread.
    """
    loop = _loop.loop
    if loop is None:
        raise RuntimeError(
            "asyncio.from_thread can only be run in a thread started by "
            "asyncio.to_thread"
        )

    ctx = contextvars.copy_context()
    return tasks.run_coroutine_threadsafe(loop, _with_context(ctx, async_func, *args, **kwargs))