'''This is a JupyterHub KubeSpawner, extended with the ability to manipulate
namespaces, and with an lsst_mgr attribute.
'''
import json
from .multispawner import MultiNamespacedKubeSpawner
from eliot import start_action
from kubespawner.objects import make_pod
from tornado import gen
from traitlets import Bool
from ..utils import make_logger, sanitize_dict


class LSSTSpawner(MultiNamespacedKubeSpawner):
    '''This, plus the LSST Manager class structure, implements the
    LSST-specific parts of our spawning requirements.
    '''
    lsst_mgr = None
    wf_api = None
    delete_grace_period = 5
    # In our LSST setup, there is a "provisionator" user, uid/gid 769,
    #  that is who we should start as.
    uid = 769
    gid = 769
    # The fields need to be defined; we don't use them.
    fs_gid = None
    supplemental_gids = []
    extra_labels = {}
    extra_annotations = []
    image_pull_secrets = None
    privileged = False
    working_dir = None
    lifecycle_hooks = {}  # This one will be useful someday.
    init_containers = []
    lab_service_account = None
    extra_container_config = None
    extra_pod_config = None
    extra_containers = []

    delete_namespace_on_stop = Bool(
        True,
        config=True,
        help='''
        If True, the entire namespace will be deleted when the lab pod stops.
        '''
    ).tag(config=True)

    enable_namespace_quotas = Bool(
        True,
        config=True,
        help='''
        If True, will create a ResourceQuota object by calling
        `self.quota_mgr.get_resource_quota_spec()` and create a quota with
        the resulting specification within the namespace.

        A subclass should override the quota manager's
        define_resource_quota_spec() to build a
        situationally-appropriate resource quota spec.
        '''
    ).tag(config=True)

    def __init__(self, *args, **kwargs):
        self.log = make_logger()
        super().__init__(*args, **kwargs)
        self.log.debug("Creating LSSTSpawner.")
        # Our API and our RBAC API are set in the super() __init__()
        # We assume that we're using an LSST Authenticator, which will
        #  therefore have an LSST MiddleManager.
        #
        # This might change with Argo Workflow.
        # self.log.debug("Spawner: {}".format(json.dumps(self.dump())))
        self.log.debug("Setting LSST Manager from authenticated user.")
        user = self.user
        if not user:
            self.log.error("No user found!")
            raise RuntimeError("Could not create spawner!")
        auth = user.authenticator
        if not auth:
            self.log.error("No authenticator found!")
            raise RuntimeError("Could not create spawner!")
        lm = auth.lsst_mgr
        if not lm:
            self.log.error("No LSST Manager found!")
            raise RuntimeError("Could not create spawner!")
        self.lsst_mgr = lm
        lm.spawner = self
        lm.user = user
        self.log.debug("Initialized {}".format(__name__))

    def auth_state_hook(self, spawner, auth_state):
        # Turns out this is in the wrong place.  It should be called
        #  _before_ get_options_form()
        super().auth_state_hook(spawner, auth_state)
        self.log.debug("{} auth_state_hook firing.".format(__name__))

    @gen.coroutine
    def get_options_form(self):
        '''Present an LSST-tailored options form; delegate to options
        form manager.

        This really is stuff that should get set from auth_state_hook...
        But as it happens, that doesn't run until after get_options_form.
        '''
        # If this is run with start_action from eliot, you get an error that
        #  the token was created in a different context.
        # I suspect this to be somewhere near the heart of our concurrency
        #  problems.
        self.log.debug("Entering get_options_form()")
        lm = self.lsst_mgr
        # Take the APIs from the api_manager
        self.api = lm.api_mgr.api
        self.rbac_api = lm.api_mgr.rbac_api
        self.wf_api = lm.api_mgr.wf_api
        om = lm.optionsform_mgr
        self.log.debug("Requesting auth state")
        auth_state = yield self.user.get_auth_state()
        if not auth_state:
            raise ValueError(
                "Auth state empty for user {}".format(self.user))
        self.log.debug("Setting auth_mgr's auth_state from user.")
        lm.auth_mgr.auth_state = auth_state
        self.log.debug("Parsing auth_state.")
        _ = yield self.asynchronize(lm.auth_mgr.parse_auth_state)
        # Will throw error if no UID
        self.log.debug("Requesting options_form from manager.")
        form = yield self.asynchronize(om.get_options_form)
        self.log.debug("Form received: {}".format(form))
        return form

    def set_user_namespace(self):
        '''Get namespace and store it here (for spawning) and in
        namespace_mgr.'''
        with start_action(action_type="set_user_namespace"):
            ns = self.get_user_namespace()
            self.namespace = ns
            self.lsst_mgr.namespace_mgr.set_namespace(ns)

    def get_user_namespace(self):
        '''Return namespace for user pods (and ancillary objects).
        '''
        with start_action(action_type="get_user_namespace"):
            defname = self._namespace_default()
            # We concatenate the default namespace and the name so that we
            #  can continue having multiple Jupyter instances in the same
            #  k8s cluster in different namespaces.  The user namespaces must
            #  themselves be namespaced, as it were.
            if defname == "default":
                raise ValueError("Won't spawn into default namespace!")
            return "{}-{}".format(defname, self.user.escaped_name)

    def start(self):
        with start_action(action_type="start"):
            # All we need to do is ensure the namespace and K8s ancillary
            #  resources before we run the superclass method to spawn a pod,
            #  so we have the namespace to spawn into, and the service account
            #  with appropriate roles and rolebindings.
            self.log.debug("Starting; creating namespace/ancillary objects.")
            self.set_user_namespace()  # Set namespace here/in namespace_mgr
            self.lsst_mgr.ensure_resources()
            retval = super().start()
            return retval

    @gen.coroutine
    def stop(self, now=False):
        '''After stopping pod, delete the namespace if that option is set.
        '''
        with start_action(action_type="stop"):
            _ = yield super().stop(now)
            if self.delete_namespace_on_stop:
                nsm = self.lsst_mgr.namespace_mgr
                self.log.debug("Attempting to delete namespace.")
                self.asynchronize(nsm.maybe_delete_namespace)
            else:
                self.log.debug("'delete_namespace_on_stop' not set.")

    def options_from_form(self, formdata=None):
        '''Delegate to form manager.
        '''
        # LSST Manager sometimes unset when we make it here.  I don't
        #  understand why that is.  Let's try moving the hooking up
        #  lsst_mgr into __init__ rather than get_form_options.
        with start_action(action_type="options_from_form"):
            return self.lsst_mgr.optionsform_mgr.options_from_form(formdata)

    @gen.coroutine
    def get_pod_manifest(self):
        # Extend pod manifest.  This is a monster method.
        # Run the superclass version, and then extract the fields
        with start_action(action_type="get_pod_manifest"):
            self.log.debug("Creating pod manifest.")
            orig_pod = yield super().get_pod_manifest()
            labels = orig_pod.metadata.labels.copy()
            annotations = orig_pod.metadata.annotations.copy()
            ctrs = orig_pod.spec.containers
            cmd = None
            if ctrs and len(ctrs) > 0:
                cmd = ctrs[0].args or ctrs[0].command
            # That should be it from the standard get_pod_manifest
            # Now we finally need all that data we have been managing.
            # Add label and annotations for ArgoCD management.
            labels['argocd.argoproj.io/instance'] = 'nublado-users'
            annotations['argocd.argoproj.io/compare-options'] = \
                'IgnoreExtraneous'
            annotations['argocd.argoproj.io/sync-options'] = 'Prune=false'
            cfg = self.lsst_mgr.config
            em = self.lsst_mgr.env_mgr
            am = self.lsst_mgr.auth_mgr
            nm = self.lsst_mgr.namespace_mgr
            vm = self.lsst_mgr.volume_mgr
            om = self.lsst_mgr.optionsform_mgr
            # Get the standard env and then update it with the environment
            # from our environment manager, except that we want the tokens from
            # the standard env
            tokens = {}
            pod_env = self.get_env()
            for fld in ['JUPYTERHUB_API_TOKEN', 'JPY_API_TOKEN']:
                val = pod_env.get(fld)
                if val:
                    tokens[fld] = val
            em.create_pod_env()
            pod_env.update(em.get_env())
            pod_env.update(tokens)
            # And now glue in environment information that was stashed at
            #  authentication time (UID/GIDs, authenticator-specific settings)
            #  in the auth_mgr.
            pod_env.update(am.pod_env)
            # Set some constants
            # First pulls can be really slow for the LSST stack containers,
            #  so let's give it a big timeout (this is in seconds)
            self.http_timeout = 60 * 15
            self.start_timeout = 60 * 15
            # We are running the Lab at the far end, not the old Notebook
            self.default_url = '/lab'
            # We always want to check for refreshed images.
            self.image_pull_policy = 'Always'
            # If we can spawn other pods from the Lab, we need a svcaccount.
            self.lab_service_account = nm.service_account
            # Get image name
            pod_name = self.pod_name
            # Get default image name; we will try to replace from options form.
            image = self.image or self.orig_pod.image or cfg.lab_default_image
            # Same with tag.
            tag = "latest"
            # Parse options form result.
            size = None
            image_size = None
            clear_dotlocal = False
            if not self.user_options:
                raise ValueError("No options form data!")
            self.log.debug("User options from form:\n" +
                           json.dumps(self.user_options,
                                      sort_keys=True, indent=4))
            size = self.user_options.get('size')
            if not size:
                # Use the deployment defaults
                size_idx = cfg.size_index
                size = cfg.form_sizelist[size_idx]
            image_size = om.sizemap[size]
            self.log.debug("Image size: {}".format(
                json.dumps(image_size, sort_keys=True, indent=4)))
            mem_limit = image_size["mem"]
            cpu_limit = image_size["cpu"]
            mem_guar = image_size["mem_guar"]
            cpu_guar = image_size["cpu_guar"]
            self.mem_guarantee = mem_guar
            self.cpu_guarantee = cpu_guar
            pod_env['MEM_GUARANTEE'] = mem_guar
            pod_env['MEM_LIMIT'] = mem_limit
            pod_env['CPU_GUARANTEE'] = str(cpu_guar)
            pod_env['CPU_LIMIT'] = str(cpu_limit)
            if self.user_options.get('kernel_image'):
                image = self.user_options.get('kernel_image')
                colon = image.find(':')
                if colon > -1:
                    imgname = image[:colon]
                    tag = image[(colon + 1):]
                    if tag == "recommended" or tag.startswith("latest"):
                        # Resolve convenience tags to real build tags.
                        self.log.debug("Resolving tag '{}'".format(tag))
                        qtag = om.resolve_tag(tag)
                        if qtag:
                            tag = qtag
                            image = imgname + ":" + tag
                        else:
                            self.log.warning(
                                "Failed to resolve tag '{}'".format(tag))
                            self.log.debug(
                                "Image name: %s ; tag: %s" % (imgname, tag))
                    if tag == "__custom":
                        self.log.debug("Tag is __custom: retrieving real " +
                                       "value from drop-down list.")
                        cit = self.user_options.get('image_tag')
                        if cit:
                            tag = cit
                            image = imgname + ":" + cit
                self.log.debug("Replacing image from options form: %s" % image)
                self.image = image
                pod_env['JUPYTER_IMAGE_SPEC'] = image
                pod_env['JUPYTER_IMAGE'] = image
                # Set flag to clear .local if indicated
            clear_dotlocal = self.user_options.get('clear_dotlocal')
            if clear_dotlocal:
                pod_env['CLEAR_DOTLOCAL'] = "TRUE"
            # Set up Lab pod resource constraints (not namespace quotas)
            # These are the defaults from the config
            # We don't care about the image name anymore: the user pod will
            #  be named "nb" plus the username and tag, to keep the pod name
            #  short.
            rt_tag = tag.replace('_', '-')
            pn_template = "nb-{username}-" + rt_tag
            pod_name = self._expand_user_properties(pn_template)
            self.pod_name = pod_name
            self.log.debug(
                "Replacing pod name from options form: %s" % pod_name)
            # Get quota definitions from quota manager.
            if self.enable_namespace_quotas:
                qmq = self.lsst_mgr.quota_mgr.quota
                if qmq:
                    if "limits.cpu" in qmq:
                        pod_env['NAMESPACE_CPU_LIMIT'] = qmq["limits.cpu"]
                    if "limits.memory" in qmq:
                        nmlimit = qmq["limits.memory"]
                        if nmlimit[-2:] == "Mi":
                            nmlimit = nmlimit[:-2] + "M"
                        pod_env['NAMESPACE_MEM_LIMIT'] = nmlimit
            # Get volume definitions from volume manager.
            vm.make_volumes_from_config()
            pod_env['DASK_VOLUME_B64'] = vm.get_dask_volume_b64()
            self.volumes = vm.k8s_volumes
            self.volume_mounts = vm.k8s_vol_mts
            # Add multus annotations if requested
            if cfg.enable_multus:
                annotations['k8s.v1.cni.cncf.io/networks'] = \
                    'kube-system/macvlan-conf'
            # Generate the pod definition.
            sanitized_env = sanitize_dict(pod_env, [
                'ACCESS_TOKEN', 'GITHUB_ACCESS_TOKEN',
                'JUPYTERHUB_API_TOKEN', 'JPY_API_TOKEN'])
            self.log.debug("Pod environment: {}".format(json.dumps(
                sanitized_env,
                sort_keys=True,
                indent=4)))
            self.log.debug("About to run make_pod()")
            pod = make_pod(
                name=self.pod_name,
                cmd=cmd,
                port=self.port,
                image=self.image,
                image_pull_policy=self.image_pull_policy,
                image_pull_secret=self.image_pull_secrets,
                node_selector=self.node_selector,
                run_as_uid=self.uid,
                run_as_gid=self.gid,
                fs_gid=self.fs_gid,
                supplemental_gids=self.supplemental_gids,
                run_privileged=self.privileged,
                env=pod_env,
                volumes=self._expand_all(self.volumes),
                volume_mounts=self._expand_all(self.volume_mounts),
                working_dir=self.working_dir,
                labels=labels,
                annotations=annotations,
                cpu_limit=self.cpu_limit,
                cpu_guarantee=self.cpu_guarantee,
                mem_limit=self.mem_limit,
                mem_guarantee=self.mem_guarantee,
                extra_resource_limits=self.extra_resource_limits,
                extra_resource_guarantees=self.extra_resource_guarantees,
                lifecycle_hooks=self.lifecycle_hooks,
                init_containers=self._expand_all(self.init_containers),
                service_account=self.lab_service_account,
                extra_container_config=self.extra_container_config,
                extra_pod_config=self.extra_pod_config,
                extra_containers=self.extra_containers,
                node_affinity_preferred=self.node_affinity_preferred,
                node_affinity_required=self.node_affinity_required,
                pod_affinity_preferred=self.pod_affinity_preferred,
                pod_affinity_required=self.pod_affinity_required,
                pod_anti_affinity_preferred=self.pod_anti_affinity_preferred,
                pod_anti_affinity_required=self.pod_anti_affinity_required,
                priority_class_name=self.priority_class_name,
                logger=self.log,
            )
            return pod

    def dump(self):
        '''Return dict representation suitable for pretty-printing.'''
        sd = {"namespace": self.namespace,
              "uid": self.uid,
              "gid": self.gid,
              "fs_gid": self.fs_gid,
              "supplemental_gids": self.supplemental_gids,
              "extra_labels": self.extra_labels,
              "extra_annotations": self.extra_annotations,
              "delete_grace_period": self.delete_grace_period,
              "privileged": self.privileged,
              "working_dir": self.working_dir,
              "lifecycle_hooks": self.lifecycle_hooks,
              "init_containers": self.init_containers,
              "lab_service_account": self.lab_service_account,
              "extra_container_config": self.extra_container_config,
              "extra_pod_config": self.extra_pod_config,
              "extra_containers": self.extra_containers,
              "delete_namespace_on_stop": self.delete_namespace_on_stop,
              "enable_namespace_quotas": self.enable_namespace_quotas,
              "lsst_mgr": self.lsst_mgr.dump()}
        return sd

    def toJSON(self):
        return json.dumps(self.dump())
