import asyncio
from contextlib import suppress
import logging
import threading
import warnings
from tornado.ioloop import PeriodicCallback

import dask.config
from dask.utils import format_bytes

from .adaptive import Adaptive

from ..utils import (
    log_errors,
    sync,
    Log,
    Logs,
    thread_state,
    format_dashboard_link,
    parse_timedelta,
)


logger = logging.getLogger(__name__)


class Cluster:
    """ Superclass for cluster objects

    This class contains common functionality for Dask Cluster manager classes.

    To implement this class, you must provide

    1.  A ``scheduler_comm`` attribute, which is a connection to the scheduler
        following the ``distributed.core.rpc`` API.
    2.  Implement ``scale``, which takes an integer and scales the cluster to
        that many workers, or else set ``_supports_scaling`` to False

    For that, you should get the following:

    1.  A standard ``__repr__``
    2.  A live IPython widget
    3.  Adaptive scaling
    4.  Integration with dask-labextension
    5.  A ``scheduler_info`` attribute which contains an up-to-date copy of
        ``Scheduler.identity()``, which is used for much of the above
    6.  Methods to gather logs
    """

    _supports_scaling = True

    def __init__(self, asynchronous):
        self.scheduler_info = {"workers": {}}
        self.periodic_callbacks = {}
        self._asynchronous = asynchronous

        self.status = "created"

    async def _start(self):
        comm = await self.scheduler_comm.live_comm()
        await comm.write({"op": "subscribe_worker_status"})
        self.scheduler_info = await comm.read()
        self._watch_worker_status_comm = comm
        self._watch_worker_status_task = asyncio.ensure_future(
            self._watch_worker_status(comm)
        )
        self.status = "running"

    async def _close(self):
        if self.status == "closed":
            return

        await self._watch_worker_status_comm.close()
        await self._watch_worker_status_task

        for pc in self.periodic_callbacks.values():
            pc.stop()
        await self.scheduler_comm.close_rpc()

        self.status = "closed"

    def close(self, timeout=None):
        with suppress(RuntimeError):  # loop closed during process shutdown
            return self.sync(self._close, callback_timeout=timeout)

    def __del__(self):
        if self.status != "closed":
            with suppress(AttributeError, RuntimeError):  # during closing
                self.loop.add_callback(self.close)

    async def _watch_worker_status(self, comm):
        """ Listen to scheduler for updates on adding and removing workers """
        while True:
            try:
                msgs = await comm.read()
            except OSError:
                break

            with log_errors():
                for op, msg in msgs:
                    self._update_worker_status(op, msg)

        await comm.close()

    def _update_worker_status(self, op, msg):
        if op == "add":
            workers = msg.pop("workers")
            self.scheduler_info["workers"].update(workers)
            self.scheduler_info.update(msg)
        elif op == "remove":
            del self.scheduler_info["workers"][msg]
        else:
            raise ValueError("Invalid op", op, msg)

    def adapt(self, Adaptive=Adaptive, **kwargs) -> Adaptive:
        """ Turn on adaptivity

        For keyword arguments see dask.distributed.Adaptive