from distributed.compatibility import MutableMapping
from distributed.utils import log_errors, tokey


class PublishExtension(object):
    """ An extension for the scheduler to manage collections

    *  publish-list
    *  publish-put
    *  publish-get
    *  publish-delete
    """

    def __init__(self, scheduler):
        self.scheduler = scheduler
        self.datasets = dict()

        handlers = {
            "publish_list": self.list,
            "publish_put": self.put,
            "publish_get": self.get,
            "publish_delete": self.delete,
        }

        self.scheduler.handlers.update(handlers)
        self.scheduler.extensions["publish"] = self

    def put(self, stream=None, keys=None, data=None, name=None, client=None):
        with log_errors():
            if name in self.datasets:
                raise KeyError("Dataset %s already exists" % name)
            self.scheduler.client_desires_keys(keys, "published-%s" % tokey(name))
            self.datasets[name] = {"data": data, "keys": keys}
            return {"status": "OK", "name": name}

    def delete(self, stream=None, name=None):
        with log_errors():
            out = self.datasets.pop(name, {"keys": []})
            self.scheduler.client_releases_keys(
                out["keys"], "published-%s" % tokey(name)
            )

    def list(self, *args):
        with log_errors():
            return list(sorted(self.datasets.keys(), key=str))

    def get(self, stream, name=None, client=None):
        with log_errors():
            return self.datasets.get(name, None)


class Datasets(MutableMapping):
    """A dict-like wrapper around :class:`Client` dataset methods.

    Parameters
    ----------
    client : distributed.client.Client

    """

    def __init__(self, client):
        self.__client = client

    def __getitem__(self, key):
        return self.__client.get_dataset(key)

    def __setitem__(self, key, value):
        self.__client.publish_dataset(value, name=key)

    def __delitem__(self, key):
        self.__client.unpublish_dataset(key)

    def __iter__(self):
        for key in self.__client.list_datasets():
            yield key

    def __len__(self):
        return len(self.__client.list_datasets())