# The spawner is based on PRP's https://gitlab.nrp-nautilus.io/vpeddu/jupyterlab-west.git
# This file is managed by [youngsu.kim|dung.vu]@csusb.edu
from kubespawner import KubeSpawner


class MySpawner(KubeSpawner):
    with open('/etc/jupyterhub/custom/my_spawner.html', 'r') as f:
        profile_form_template = f.readlines()
    profile_form_template = "".join(profile_form_template)

    def options_from_form(self, formdata):
 #       if not self.profile_list or not hasattr(self, '_profile_list'):
 #           return formdata

        profile_slug = formdata.get('profile', [None])[0]

          # initialize a dictionary to return
        user_options = {}

          # if a profile is declared, add a dictionary key for the profile, and
          # dictionary keys for the formdata related to the profile's
          # profile_options, as recognized by being named like:
          #
          #     profile-option-{profile_slug}--{profile_option_slug}
          #

        if profile_slug:
            user_options["profile"] = profile_slug
            prefix = f'profile-option-{profile_slug}--'
            for k, v in formdata.items():
                if k.startswith(prefix):
                    profile_option_slug = k[len(prefix) :]
                    user_options[profile_option_slug] = v[0]
                else:
                    user_options[prefix+k] = v[0]

          # self.log.info("Applying KubeSpawner override for profile '%s'", user_options)
          # kubespawner_override = user_options.get('kubespawner_override', {})

        gpus = int(formdata.get('gpus', [0])[0])

        profile = self._get_profile(profile_slug, self.profile_list)
        if (("image_choice" in profile) and (isinstance(profile["image_choice"], dict))):
            if gpus > 0:
              image = profile["image_choice"]["cuda"]
            else:
              image = profile["image_choice"]["cpu"]
            
        setattr(self, "image", image)
        
        setattr(self, "extra_resource_limits", {"nvidia.com/gpu": gpus})
        setattr(self, "mem_guarantee", formdata.get('ram', [0])[0]+"G")
        setattr(self, "cpu_guarantee", float(formdata.get('cores', [0])[0]))
        setattr(self, "mem_limit", formdata.get('ram', [0])[0]+"G")
        setattr(self, "cpu_limit", float(formdata.get('cores', [0])[0]))

        nodeSelectorTermsExpressions = [
            {
            'key': 'topology.kubernetes.io/region',
            'operator': 'In',
            'values': ["us-west"]
            }
        ]

        if formdata.get('gputype', [0])[0]:
            nodeSelectorTermsExpressions.append({
                'key': 'nvidia.com/gpu.product',
                'operator': 'In',
                'values': formdata.get('gputype', [0])
            })
        if gpus == 0:
            nodeSelectorTermsExpressions.append({
                'key': 'nvidia.com/gpu.product',
                'operator': 'NotIn',
                'values': [ "NVIDIA-GeForce-RTX-3090", "NVIDIA-A100-80GB-PCIe-MIG-1g.10gb", "NVIDIA-L40", "NVIDIA-A10",
                           "NVIDIA-RTX-A4000", "NVIDIA-RTX-A5000", "NVIDIA-RTX-A5000", "NVIDIA-A100-SXM4-80GB", "NVIDIA-GeForce-GTX-1080-Ti",
                           "NVIDIA-GeForce-RTX-2080-Ti"]
            })


#        tolerationsExpression = [{
#            'key': 'nautilus.io/nrp-testing',
#            'operator': 'Exists',
#            'effect': 'NoSchedule',
#            }]

        setattr(self, 'extra_pod_config', {
            'affinity': {
                'nodeAffinity': {
                   'requiredDuringSchedulingIgnoredDuringExecution': {
                       'nodeSelectorTerms': [{
                            'matchExpressions': nodeSelectorTermsExpressions,
                        }],
                    },
                },
            },
#            'tolerations': tolerationsExpression,
        })
        
#        self.volume_mounts = [
#            {
#                'name': 'volume-{username}',
#                'mountPath': '/home/jovyan',
#            },
#            {
#                'name': 'csusb-share',
#                'mountPath': '/home/jovyan/shared',
#            },
#            ]
#        self.volumes = [
#            {
#                'name': 'volume-{username}',
#                'persistentVolumeClaim': {
#                    'claimName': 'claim-{username}'
#                    }
#            },
#            {
#                'name': 'csusb-share',
#                'persistentVolumeClaim': {
#                    'claimName': 'csusb-khan-share'
#                    }
#            },
#            ]

        return options

    profile_list = [
        {
            'display_name': 'Stack Minimal',
        },
        {
            'display_name': 'Tensorflow',
            'slug': 'tf',
            'image_choice': {
              'cuda': 'quay.io/jupyter/tensorflow-notebook:cuda-2024-04-22',
              'cpu': 'quay.io/jupyter/tensorflow-notebook:2024-04-22'
            }
        },
        {
            'display_name': "Stack R-Studio",
            'kubespawner_override': {
                'image': 'gitlab-registry.nrp-nautilus.io/csusb-prp/csusb-jupyter-stacks/r-studio-notebook:v1.3.1',
#                'image_spec': 'gitlab-registry.nrp-nautilus.io/youngsu_kim/csusb-jupyter-stack/stack-rstudio-becerra-class',
            }
        },
        {
            'display_name': 'Stack All-Spark',
            'kubespawner_override': {
#                'image_spec': 'jupyter/all-spark-notebook',
                'image_spec': 'gitlab-registry.nrp-nautilus.io/prp/jupyter-stack/all-spark:v1.3.1',
            }
        },
        {
            'display_name': 'Stack pySpark',
            'kubespawner_override': {
                'image_spec': 'gitlab-registry.nrp-nautilus.io/prp/jupyter-stack/pyspark:v1.3.1',
            }
        },
        {
            'display_name': 'Stack Tensorflow + PRP Packages',
            'kubespawner_override': {
                'image_spec': 'gitlab-registry.nrp-nautilus.io/prp/jupyter-stack/prp:v1.3.1',
            }
        }
        ,
        {
            'display_name': 'Stack Scipy',
            'kubespawner_override': {
                'image': 'gitlab-registry.nrp-nautilus.io/prp/jupyter-stack/scipy:v1.3.1',
                'imagePullPolicy': 'IfNotPresent',
            }
        },
        {
            'display_name': 'Stack PyTorch2',
            'kubespawner_override': {

                'image': 'gitlab-registry.nrp-nautilus.io/prp/jupyter-stack/prp-pytorch2:v1.0',
                'imagePullPolicy': 'IfNotPresent',
            }
        },


    ]

c.JupyterHub.spawner_class = MySpawner