import asyncio
import atexit
import sys

import dask
from dask.distributed import Client, Nanny, Scheduler
from distributed.utils import import_term
from tornado import gen
from tornado.ioloop import IOLoop


def initialize(
    interface=None,
    nthreads=1,
    local_directory="",
    memory_limit="auto",
    nanny=False,
    dashboard=True,
    dashboard_address=":8787",
    protocol=None,
    worker_class="distributed.Worker",
    worker_options=None,
    comm=None,
    exit=True,
):
    """
    Initialize a Dask cluster using mpi4py

    Using mpi4py, MPI rank 0 launches the Scheduler, MPI rank 1 passes through to the
    client script, and all other MPI ranks launch workers.  All MPI ranks other than
    MPI rank 1 block while their event loops run.

    In normal operation these ranks exit once rank 1 ends. If exit=False is set they
    instead return an bool indicating whether they are the client and should execute
    more client code, or a worker/scheduler who should not.  In this case the user is
    responsible for the client calling send_close_signal when work is complete, and
    checking the returned value to choose further actions.

    Parameters
    ----------
    interface : str
        Network interface like 'eth0' or 'ib0'
    nthreads : int
        Number of threads per worker
    local_directory : str
        Directory to place worker files
    memory_limit : int, float, or 'auto'
        Number of bytes before spilling data to disk.  This can be an
        integer (nbytes), float (fraction of total memory), or 'auto'.
    nanny : bool
        Start workers in nanny process for management (deprecated, use worker_class instead)
    dashboard : bool
        Enable Bokeh visual diagnostics
    dashboard_address : str
        Bokeh port for visual diagnostics
    protocol : str
        Protocol like 'inproc' or 'tcp'
    worker_class : str
        Class to use when creating workers
    worker_options : dict
        Options to pass to workers
    comm: mpi4py.MPI.Intracomm
        Optional MPI communicator to use instead of COMM_WORLD
    exit: bool
        Whether to call sys.exit on the workers and schedulers when the event
        loop completes.

    Returns
    -------
    is_client: bool
        Only returned if exit=False. Inidcates whether this rank should continue
        to run client code (True), or if it acts as a scheduler or worker (False).
    """
    if comm is None:
        from mpi4py import MPI

        comm = MPI.COMM_WORLD

    rank = comm.Get_rank()
    loop = IOLoop.current()

    if not worker_options:
        worker_options = {}

    if rank == 0:

        async def run_scheduler():
            async with Scheduler(
                interface=interface,
                protocol=protocol,
                dashboard=dashboard,
                dashboard_address=dashboard_address,
            ) as scheduler:
                comm.bcast(scheduler.address, root=0)
                comm.Barrier()
                await scheduler.finished()

        asyncio.get_event_loop().run_until_complete(run_scheduler())
        if exit:
            sys.exit()
        else:
            return False

    else:
        scheduler_address = comm.bcast(None, root=0)
        dask.config.set(scheduler_address=scheduler_address)
        comm.Barrier()

    if rank == 1:
        if exit:
            atexit.register(send_close_signal)
        return True
    else:

        async def run_worker():
            WorkerType = import_term(worker_class)
            if nanny:
                raise DeprecationWarning(
                    "Option nanny=True is deprectaed, use worker_class='distributed.Nanny' instead"
                )
                WorkerType = Nanny
            opts = {
                "interface": interface,
                "protocol": protocol,
                "nthreads": nthreads,
                "memory_limit": memory_limit,
                "local_directory": local_directory,
                "name": rank,
                **worker_options,
            }
            async with WorkerType(**opts) as worker:
                await worker.finished()

        asyncio.get_event_loop().run_until_complete(run_worker())
        if exit:
            sys.exit()
        else:
            return False


def send_close_signal():
    """
    The client can call this function to explicitly stop
    the event loop.

    This is not needed in normal usage, where it is run
    automatically when the client code exits python.

    You only need to call this manually when using exit=False
    in initialize.
    """

    async def stop(dask_scheduler):
        await dask_scheduler.close()
        await gen.sleep(0.1)
        local_loop = dask_scheduler.loop
        local_loop.add_callback(local_loop.stop)

    with Client() as c:
        c.run_on_scheduler(stop, wait=False)
