import copy
import datetime
import json
import logging
import os
import signal
import time
from eliot import start_action
from kubernetes import client, config
from jupyterhubutils.scanrepo import ScanRepo
from threading import Thread
from ..utils import make_logger


class Prepuller(object):
    '''Class for generating and reaping the Pods for the prepuller.
    '''

    def __init__(self, args=None):
        self.logger = make_logger()
        self.repo = None
        if not args:
            raise ValueError("args must be set!")
        self.args = args
        self.debug = False
        if args.debug:
            self.debug = True
        if self.debug:
            self.logger.setLevel(logging.DEBUG)
            self.logger.debug("Debug logging enabled.")
        namespace = self.args.namespace
        if not namespace:
            self.logger.warning("Using kubernetes namespace 'default'")
            namespace = "default"
        self.namespace = namespace
        try:
            config.load_incluster_config()
        except config.ConfigException:
            config.load_kube_config()
        self.client = client.CoreV1Api()
        self.images = []
        self.nodes = []
        self.pod_specs = {}
        self.created_pods = []

        self.logger.debug("Arguments: %s" % str(args))
        self.command = self.args.command
        self.cachefile = self.args.cachefile
        if self.args.list:
            for image in self.args.list:
                # Make fully-qualified image name
                colons = image.count(':')
                if colons == 0:
                    image = image + ":latest"
                slashes = image.count('/')
                if slashes == 0:
                    image = "library/" + image
            self.images.append(image)
        # Cheap way to deduplicate lists
        self.images = list(set(self.images))
        if self.images:
            self.images.sort()
        # Not portable to non-Unixy systems.
        if self.args.timeout >= 0:
            self.logger.debug("Setting timeout to %d s." % self.args.timeout)
            signal.signal(signal.SIGALRM, self._timeout_handler)
            signal.alarm(self.args.timeout)

    def _timeout_handler(self, signum, frame):
        with start_action(action_type="_timeout_handler"):
            self.logger.error(
                "Did not complete in %d s.  Terminating." % self.args.timeout)
            self._destroy_pods(selective=False)
            raise RuntimeError("Timed out")

    def _destroy_pods(self, selective=False):
        '''Get a pod list and delete any that are still running if
        selective is False, or any in state "Succeeded" or "Failed"
        if selective is True.
        '''
        with start_action(action_type="_destroy_pods"):
            self.logger.debug("Looking for pods to delete.")
            cleanup = self._get_deletion_list(selective=selective)
            for podname in cleanup:
                self.delete_pod(podname)

    def _get_deletion_list(self, selective=True):
        with start_action(action_type="_get_deletion_list"):
            cleanup = []
            speclist = []
            v1 = self.client
            for x in self.pod_specs:
                speclist.extend(self.pod_specs[x])
            specnames = [self._derive_pod_name(x) for x in speclist]
            podlist = v1.list_namespaced_pod(self.namespace)
            for pod in podlist.items:
                podname = pod.metadata.name
                phase = pod.status.phase
                if podname in specnames:
                    if selective:
                        if (phase != "Succeeded" and phase != "Failed"):
                            self.logger.debug("Pod '%s' %s; not cleaning."
                                              % (podname, phase))
                            continue
                    self.logger.debug(
                        "Pod '%s' %s; adding to cleanup." % (podname, phase))
                    cleanup.append(podname)
            return cleanup

    def update_images_from_repo(self):
        '''Scan the repo looking for images.
        '''
        with start_action(action_type="update_images_from_repo"):
            if not self.repo:
                self.repo = ScanRepo(host=self.args.repo,
                                     path=self.args.path,
                                     owner=self.args.owner,
                                     name=self.args.name,
                                     dailies=self.args.dailies,
                                     weeklies=self.args.weeklies,
                                     releases=self.args.releases,
                                     experimentals=self.args.experimentals,
                                     recommended=self.args.recommended,
                                     json=True, insecure=self.args.insecure,
                                     sort_field=self.args.sort,
                                     cachefile=self.cachefile,
                                     debug=self.args.debug)
            if not self.args.no_scan:
                if self.args.repo:
                    self.logger.debug("Scanning images: '%s' " %
                                      self.args.repo)
                else:
                    self.logger.debug("Scanning Docker repo for images")
                self.repo.scan()
                if self.repo.debug:
                    repocopy = copy.deepcopy(self.repo.data)
                    for itype in repocopy:
                        for item in repocopy[itype]:
                            uddt = item["updated"]
                            if type(uddt) is datetime.datetime:
                                uddt = uddt.isoformat()
                                item["updated"] = uddt
                    self.logger.debug("Scan Data: " +
                                      json.dumps(repocopy,
                                                 sort_keys=True,
                                                 indent=4))
                scan_imgs = []
                sections = []
                if self.args.recommended:
                    sections.extend(["recommended"])
                sections.extend(["experimental", "daily", "weekly", "release"])
                for section in sections:
                    if section not in self.repo.data:
                        continue
                    for entry in self.repo.data[section]:
                        exhost = ''
                        if self.args.repo:
                            exhost = self.args.repo
                            if self.args.port:
                                exhost += ":" + self.args.port
                        if exhost == "hub.docker.com":
                            # We could change it to "index.docker.io" but
                            #  leaving it empty works too.
                            exhost = ""
                            if exhost and exhost[-1] != "/":
                                exhost += "/"
                            scan_imgs.append(exhost + self.args.owner + "/" +
                                             self.args.name + ":" +
                                             entry["name"])
                current_imgs = [x for x in self.images]
                # Dedupe by running the list through a set.
                current_imgs.extend(scan_imgs)
                current_imgs = list(set(current_imgs))
                if current_imgs:
                    current_imgs.sort()
                self.images = current_imgs

    def build_nodelist(self):
        '''Make a list of all schedulable nodes, respecting RESTRICT_*
        environment variables.
        '''
        with start_action(action_type="build_nodelist"):
            v1 = self.client
            logger = self.logger
            logger.debug("Getting schedulable node list.")
            v1nodelist = v1.list_node()
            nodes = []
            for thing in v1nodelist.items:
                spec = thing.spec
                if spec.unschedulable:
                    continue
                if spec.taints:
                    taints = [x.effect for x in spec.taints]
                    if "NoSchedule" in taints:
                        continue
                if self.reject_by_label(thing):
                    continue
                nodes.append(thing.metadata.name)
            logger.debug("Schedulable list: %s" % str(nodes))
            self.nodes = nodes

    def reject_by_label(self, node):
        '''If node labels are set to restrict Lab spawn, reject nodes that
        are not suitable for Lab/Dask.
        '''
        with start_action(action_type="reject_by_label"):
            logger = self.logger
            logger.debug("Checking for node labels.")
            if not os.getenv("RESTRICT_LAB_NODES"):
                logger.debug("Lab nodes are not restricted.")
                return False
            labels = node.metadata.labels
            name = node.metadata.name or 'Node Name Unknown'
            if not labels:
                logger.debug("Nodes are not labelled.")
                return False
            lab_ok = labels.get("jupyterlab")
            if lab_ok and lab_ok == "ok":
                logger.debug("Node '%s' is allowed for Lab usage." % name)
                return False
            if not os.getenv("ALLOW_DASK_SPAWN"):
                logger.debug("Lab spawn not allowed for node '%s'." % name)
                return True
            if not os.getenv("RESTRICT_DASK_NODES"):
                logger.debug("Dask allowed and unrestricted.")
                return False
            dask_ok = labels.get("dask")
            if dask_ok and dask_ok == "ok":
                logger.debug("Node '%s' is allowed for Dask usage." % name)
                return False
            logger.debug("Lab/Dask spawn not allowed for node '%s'" % name)
            return True

    def build_pod_specs(self):
        '''Build a dict of Pod specs by node, each node having a list of
        specs.
        '''
        with start_action(action_type="build_pod_specs"):
            specs = {}
            for node in self.nodes:
                specs[node] = []
                for img in self.images:
                    specs[node].append(self._build_pod_spec(img, node))
            self.pod_specs = specs
            self.logger.debug("Specs: %s" % str(self.pod_specs))

    def _build_pod_spec(self, img, node):
        with start_action(action_type="_build_pod_spec"):
            spec = client.V1PodSpec(
                containers=[
                    client.V1Container(
                        command=self.command,
                        image=img,
                        image_pull_policy="Always",
                        name=self._podname_from_image(img),
                        security_context=client.V1PodSecurityContext(
                            run_as_user=self.args.uid)
                    )
                ],
                restart_policy="Never",
                node_name=node
            )
            return spec

    def _podname_from_image(self, img):
        with start_action(action_type="_podname_from_image"):
            iname = '-'.join(img.split('/')[-2:])
            iname = iname.replace(':', '-')
            iname = iname.replace('_', '-')
            return iname

    def clean_completed_pods(self):
        '''Get a pod list and delete any that are in the speclist and have
        already run to completion.
        '''
        with start_action(action_type="clean_completed_pods"):
            self._destroy_pods(selective=True)

    def start_single_pod(self, spec):
        '''Run a pod, with a single container, on a particular node.
        (Assuming that the pod is itself tied to a node in the pod spec.)
        This has the effect of pulling the image for that pod onto that
        node.  The run itself is unimportant.  It returns the name of the
        created pod.
        '''
        with start_action(action_type="start_single_pod"):
            v1 = self.client
            name = self._derive_pod_name(spec)
            pod = client.V1Pod(spec=spec,
                               metadata=client.V1ObjectMeta(
                                   name=name)
                               )
            name = spec.containers[0].name
            self.logger.debug("Running pod %s" % name)
            made_pod = v1.create_namespaced_pod(self.namespace, pod)
            podname = made_pod.metadata.name
            return podname

    def _derive_pod_name(self, spec):
        '''Pod name is based on image and node.
        '''
        with start_action(action_type="_derive_pod_name"):
            return ("pp-" +
                    self._podname_from_image(spec.containers[0].image) +
                    "-" + spec.node_name.split('-')[-1])

    def run_pods(self):
        '''Run pods for all nodes.  Parallelize across nodes.
        '''
        with start_action(action_type="run_pods"):
            tlist = []
            for node in self.pod_specs:
                speclist = list(self.pod_specs[node])
                thd = Thread(target=self.run_pods_for_node,
                             args=(node, speclist))
                tlist.append(thd)
                thd.start()
            # Wait for all threads to return
            for thd in tlist:
                self.logger.debug("Waiting for thread '%s'" % thd.name)
                thd.join()

    def run_pods_for_node(self, node, speclist):
        '''Execute pods one at a time, so we don't overwhelm I/O.
        Execute this method in parallel across all nodes for best
        results.  Each node should have its own I/O, so they should all be
        busy at once.
        '''
        with start_action(action_type="run_pods_for_node"):
            self.logger.debug("Running pods for node %s" % node)
            for spec in speclist:
                name = spec.containers[0].name
                self.logger.debug("Running pod '%s' for node '%s'" % (name,
                                                                      node))
                podname = self.start_single_pod(spec)
                self.wait_for_pod(podname)

    def wait_for_pod(self, podname, delay=1, max_tries=3600):
        '''Wait for a particular pod to go into phase "Succeeded" or
        "Failed", and then delete the pod.
        Raise an exception if the delay timer expires.
        '''
        with start_action(action_type="wait_for_pod"):
            v1 = self.client
            namespace = self.namespace
            tries = 1
            while True:
                pod = v1.read_namespaced_pod(podname, namespace)
                phase = pod.status.phase
                if phase in ["Failed", "Succeeded"]:
                    if phase == "Failed":
                        self.logger.error("Pod '%s' failed" % podname)
                    self.delete_pod(podname)
                    return
                if tries >= max_tries:
                    errstr = ("Pod '%s' did not complete after " % podname +
                              "%d %d s iterations." % (max_tries, delay))
                    self.logger.error(errstr)
                    raise RuntimeError(errstr)
                pstr = "Wait %d s [%d/%d]for pod '%s' [%s]" % (
                    delay, tries, max_tries, podname, phase)
                self.logger.debug(pstr)
                time.sleep(delay)
                tries = tries + 1

    def delete_pod(self, podname):
        '''Delete a named pod.
        '''
        with start_action(action_type="delete_pod"):
            v1 = self.client
            self.logger.debug("Deleting pod %s" % podname)
            v1.delete_namespaced_pod(podname, self.namespace)
