import ipaddress
import logging
from .forwarder import Forwarder
from .utils import get_proxy_url, get_hostname
from dask.distributed import Client


def _remove_prefix(s, prefix):
    return s[len(prefix):] if s.startswith(prefix) else s


class ClusterProxy(object):
    """Provides a proxy service to map a local port to a worker node's
    dashboard, which should be on its "bokeh" service.

    This allows us to proxy to the worker even though the k8s network is not
    accessible externally.

    It must be created with an instance of a dask.distributed.Client
    as its argument.

    This would work better if k8s workers had a bokeh service, and if the
    bokeh service told you how to get to its dashboard page once you have
    the host and port.
    """
    client = None
    cluster = None
    ioloop = None
    workers = {}

    def __init__(self, client):
        if not isinstance(client, Client):
            estr = "'client' argument must be dask.distributed.Client!"
            raise RuntimeError(estr)
        self.client = client
        self.logger = logging.getLogger(__name__)
        self.cluster = client.cluster
        self.ioloop = client.io_loop
        self.scheduler_url = None
        port = client.cluster.scheduler.identity()["services"].get('bokeh')
        if port:
            self.scheduler_url = get_proxy_url(port) + "/status"
        self.refresh_workers()

    def refresh_workers(self):
        """Rebuild current worker map from actual state.
        """
        current_workers = self.cluster.scheduler.identity().get('workers')
        current_workerlist = [_remove_prefix(x, "tcp://")
                              for x in current_workers]
        removed_workers = []
        for worker_id in self.workers:
            if worker_id not in current_workerlist:
                # Get rid of removed workers
                self._remove_worker_proxy(worker_id)
                removed_workers.append(worker_id)
        for worker_id in removed_workers:
            del self.workers[worker_id]
        for worker_id in current_workerlist:
            if worker_id not in self.workers:
                worker_record = current_workerlist[worker_id]
                self.workers[worker_id] = self._create_worker_proxy(
                    worker_record)
            # Otherwise it hasn't changed.

    def _remove_worker_proxy(self, worker_id):
        worker = self.workers.get(worker_id)
        if not worker:
            return
        forwarder = worker["forwarder"]
        forwarder.stop()

    def _create_worker_proxy(self, worker_record):
        host = worker_record["host"]
        ipaddr = ipaddress.IPAddress(host)
        port = worker_record["services"].get("bokeh")
        if not port:
            return None
        if ipaddr.is_local():
            proxy = None
            local_port = port
        else:
            proxy = Forwarder(host, port, ioloop=self.ioloop)
            proxy.start()
            local_port = proxy.get_port()
        url = get_proxy_url(local_port)
        worker = {"forwarder": proxy,
                  "url": url,
                  "local_port": local_port}
        return worker

    def get_proxies(self, workers):
        """Returns a dict of worker endpoints as keys, mapped to a dict
        containing the worker proxy url and local port it's mapped to.
        """
        rval = {}
        if not workers:
            workers = list(self.workers.keys())
        if workers:
            if type(workers) is str:
                workers = [workers]
        for worker in workers:
            rval[worker] = {"url": self.workers[worker]["url"],
                            "local_port": self.workers[worker]["local_port"]}
        return rval

    def __repr__(self):
        s = "ClusterProxy {name}:".format(name=get_hostname())
        s += "\n  Scheduler: {url}".format(url=self.scheduler_url)
        sw = self.workers
        if sw:
            s = s+"\n  Workers:"
        for worker in sw:
            s += "\n    {worker}: {url}".format(worker=worker,
                                                url=sw[worker]["url"])
        return s

    def _repr_html_(self):
        s = "<h3>ClusterProxy {name}</h3>".format(name=get_hostname())
        s += "<h4>Scheduler: {url}</h4>".format(url=self.scheduler_url)
        if len(self.workers) > 0:
            s += "<h4>Workers</h4><dl>"
            sw = self.workers
            for worker in sw:
                s += "<dt>{w}</dt><dd>{u}</dd>".format(w=worker,
                                                       u=sw[worker]["url"])
            s += "</dl>"
        return s
