import pytest
from tornado.gen import sleep

import dask
from distributed.utils_test import gen_test
from distributed.metrics import time

from dask_labextension.manager import DaskClusterManager


config = {
    "labextension": {
        "initial": [],
        "default": {},
        "factory": {
            "module": "dask.distributed",
            "class": "LocalCluster",
            "kwargs": {"processes": False},
            "args": [],
        },
    }
}


@gen_test()
async def test_start():
    with dask.config.set(config):
        async with DaskClusterManager() as manager:
            # add cluster
            model = await manager.start_cluster()
            assert not model.get("adapt")

            # close cluster
            assert len(manager.list_clusters()) == 1
            await manager.close_cluster(model["id"])

            # add cluster with adaptive configuration
            model = await manager.start_cluster(
                configuration={"adapt": {"minimum": 1, "maximum": 3}}
            )
            assert model["adapt"] == {"minimum": 1, "maximum": 3}

            await manager.close()


@gen_test()
async def test_close():
    with dask.config.set(config):
        async with DaskClusterManager() as manager:
            # start a cluster
            model = await manager.start_cluster()

            # return None if a nonexistent cluster is closed
            assert not await manager.close_cluster("fake")

            # close the cluster
            await manager.close_cluster(model["id"])
            assert not manager.list_clusters()


@gen_test()
async def test_get():
    with dask.config.set(config):
        async with DaskClusterManager() as manager:
            # start a cluster
            model = await manager.start_cluster()

            # return None if a nonexistent cluster is requested
            assert not manager.get_cluster("fake")

            # get the cluster by id
            assert model == manager.get_cluster(model["id"])


@pytest.mark.filterwarnings("ignore")
@gen_test()
async def test_list():
    with dask.config.set(config):
        async with DaskClusterManager() as manager:
            # start with an empty list
            assert not manager.list_clusters()
            # start clusters
            model1 = await manager.start_cluster()
            model2 = await manager.start_cluster()

            models = manager.list_clusters()
            assert len(models) == 2
            assert model1 in models
            assert model2 in models


@gen_test()
async def test_scale():
    with dask.config.set(config):
        async with DaskClusterManager() as manager:
            # add cluster with number of workers configuration
            model = await manager.start_cluster(configuration={"workers": 3})
            start = time()
            while model["workers"] != 3:
                await sleep(0.01)
                model = manager.get_cluster(model["id"])
                assert time() < start + 10, model["workers"]

            await sleep(0.2)  # let workers settle # TODO: remove need for this

            # rescale the cluster
            model = await manager.scale_cluster(model["id"], 6)
            start = time()
            while model["workers"] != 6:
                await sleep(0.01)
                model = manager.get_cluster(model["id"])
                assert time() < start + 10, model["workers"]


@gen_test()
async def test_adapt():
    with dask.config.set(config):
        async with DaskClusterManager() as manager:
            # add a new cluster
            model = await manager.start_cluster()
            assert not model.get("adapt")
            model = manager.adapt_cluster(model["id"], 0, 4)
            adapt = model.get("adapt")
            assert adapt
            assert adapt["minimum"] == 0
            assert adapt["maximum"] == 4


@gen_test()
async def test_initial():
    with dask.config.set(
        {
            "labextension": {
                "initial": [{"name": "foo"}],
                "default": {},
                "factory": {
                    "module": "dask.distributed",
                    "class": "LocalCluster",
                    "kwargs": {"processes": False},
                    "args": [],
                },
            }
        }
    ):
        # Test asynchronous starting of clusters via a context
        async with DaskClusterManager() as manager:
            clusters = manager.list_clusters()
            assert len(clusters) == 1
            assert clusters[0]["name"] == "foo"

        # Test asynchronous starting of clusters outside of a context
        manager = DaskClusterManager()
        assert len(manager.list_clusters()) == 0
        await manager
        clusters = manager.list_clusters()
        assert len(clusters) == 1
        assert clusters[0]["name"] == "foo"
        await manager.close()

        manager = await DaskClusterManager()
        clusters = manager.list_clusters()
        assert len(clusters) == 1
        assert clusters[0]["name"] == "foo"
        await manager.close()
