import json
import os
from kubernetes.client import V1ResourceQuotaSpec
from kubernetes.client.rest import ApiException
from kubernetes import client

from ..utils import make_logger


class LSSTQuotaManager(object):
    '''Quota support for LSST LSP Jupyterlab and Dask pods.
    '''
    quota = {}
    custom_resources = {}
    resourcemap = None

    def __init__(self, *args, **kwargs):
        self.log = make_logger()
        self.log.debug("Creating LSSTQuotaManager")
        self.parent = kwargs.pop('parent')

    def define_resource_quota_spec(self):
        '''We're going to return a resource quota spec that checks whether we
        have a custom resource map and uses that information.  If we do not
        then we use the quota from our parent's config object.

        Note that you could get a lot fancier, and check the user group
        memberships to determine what class a user belonged to, or some other
        more-sophisticated-than-one-size-fits-all quota mechanism.
        '''
        self.log.debug("Calculating default resource quotas.")
        om = self.parent.optionsform_mgr
        sizemap = om.sizemap
        # sizemap is an ordered dict, and we want the last-inserted one,
        #  which is the biggest
        big = list(sizemap.keys())[-1]
        cpu = sizemap[big]['cpu']
        cfg = self.parent.config
        max_dask_workers = 0
        if cfg.allow_dask_spawn:
            max_dask_workers = cfg.max_dask_workers
        mem_per_cpu = cfg.mb_per_cpu
        total_cpu = (1 + max_dask_workers) * cpu
        total_mem = str(int(total_cpu * mem_per_cpu + 0.5)) + "Mi"
        total_cpu = str(int(total_cpu + 0.5))
        self.log.debug("Default quota sizes: CPU %r, mem %r" % (
            total_cpu, total_mem))
        self._set_custom_user_resources()
        if self.custom_resources:
            self.log.debug("Have custom resources.")
            cpuq = self.custom_resources.get("cpu_quota")
            if cpuq:
                self.log.debug("Overriding CPU quota.")
                total_cpu = str(cpuq)
            memq = self.custom_resources.get("mem_quota")
            if memq:
                self.log.debug("Overriding memory quota.")
                total_mem = str(memq) + "Mi"
        self.log.debug("Determined quota sizes: CPU %r, mem %r" % (
            total_cpu, total_mem))
        qs = V1ResourceQuotaSpec(
            hard={"limits.cpu": total_cpu,
                  "limits.memory": total_mem})
        self.quota = qs.hard

    def _set_custom_user_resources(self):
        '''Create custom resource definitions for user.
        '''
        if not self.resourcemap:
            self.log.debug("No resource map found; generating.")
            self._create_resource_map()
        if not self.resourcemap:
            self.log.warning("No resource map; cannot generate quota.")
            return
        resources = {
            "size_index": 0,
            "cpu_quota": 0,
            "mem_quota": 0
        }
        gnames = self.parent.user.groups
        uname = self.parent.user.escaped_name
        for resdef in self.resourcemap:
            apply = False
            if resdef.get("disabled"):
                continue
            candidate = resdef.get("resources")
            if not candidate:
                continue
            self.log.debug(
                "Considering candidate resource map {}".format(resdef))
            ruser = resdef.get("user")
            rgroup = resdef.get("group")
            if ruser and ruser == uname:
                self.log.debug("User resource map match.")
                apply = True
            if rgroup and rgroup in gnames:
                self.log.debug("Group resource map match.")
                apply = True
            if apply:
                for fld in ["size_index", "cpu_quota", "mem_quota"]:
                    vv = candidate.get(fld)
                    if vv and vv > resources[fld]:
                        resources[fld] = vv
                    self.log.debug(
                        "Setting custom resources '{}'".format(resources) +
                        "for user '{}'".format(uname))
                    self.custom_resources = resources

    def _create_resource_map(self):
        resource_file = self.parent.config.resource_map
        if not os.path.exists(resource_file):
            nf_msg = ("Could not find resource definition file" +
                      " at '{}'".format(resource_file))
            self.log.error(nf_msg)
            return
        with open(resource_file, "r") as rf:
            resmap = json.load(rf)
        self.resourcemap = resmap

    # Brought in from namespacedkubespawner
    def ensure_namespaced_resource_quota(self, quotaspec):
        '''Create K8s quota object if necessary.
        '''
        self.log.debug("Entering ensure_namespaced_resource_quota()")
        namespace = self.parent.namespace_mgr.namespace
        api = self.parent.api
        if namespace == "default":
            self.log.error("Will not create quota for default namespace!")
            return
        quota = client.V1ResourceQuota(
            metadata=client.V1ObjectMeta(
                name="quota",
            ),
            spec=quotaspec
        )
        self.log.info("Creating quota: %r" % quota)
        try:
            api.create_namespaced_resource_quota(namespace, quota)
        except ApiException as e:
            if e.status != 409:
                self.log.exception("Create resourcequota '%s'" % quota +
                                   "in namespace '%s' " % namespace +
                                   "failed: %s", str(e))
                raise
            else:
                self.log.debug("Resourcequota '%r' " % quota +
                               "already exists in '%s'." % namespace)

    def destroy_namespaced_resource_quota(self):
        '''Destroys the Kubernetes namespaced resource quota.
        You don't usually have to call this, since it will get
        cleaned up as part of namespace deletion.
        '''
        namespace = self.parent.namespace_mgr.namespace
        api = self.parent.api
        qname = "quota-" + namespace
        dopts = client.V1DeleteOptions()
        self.log.info("Deleting resourcequota '%s'" % qname)
        api.delete_namespaced_resource_quota(
            qname, namespace, dopts)

    def dump(self):
        '''Return contents dict for pretty-printing/aggregation.
        '''
        qd = {"parent": str(self.parent),
              "quota": self.quota,
              "custom_resources": self.custom_resources,
              "resourcemap": self.resourcemap}
        return qd
