import json

from kubeflow.kubeflow.crud_backend import logging
from werkzeug.exceptions import BadRequest

from . import utils

log = logging.getLogger(__name__)

SERVER_TYPE_ANNOTATION = "notebooks.kubeflow.org/server-type"
HEADERS_ANNOTATION = "notebooks.kubeflow.org/http-headers-request-set"
URI_REWRITE_ANNOTATION = "notebooks.kubeflow.org/http-rewrite-uri"


def get_form_value(body, defaults, body_field, defaults_field=None,
                   optional=False):
    """
    Get the value to set by respecting the readOnly configuration for
    the field.
    If the field does not exist in the configuration then just use the form
    value.
    """
    # The field in the defaults json not be the same in the request body
    if defaults_field is None:
        defaults_field = body_field

    # If no default value exists then just return value in the request body.
    # This is also useful if we add a new field and the configmap isn't updated
    # yet
    user_value = body.get(body_field, None)
    if defaults_field not in defaults:
        return user_value

    readonly = defaults[defaults_field].get("readOnly", False)
    default_value = defaults[defaults_field]["value"]

    # if the value of a field is readonly then the request/form should not
    # contain this field
    if readonly:
        if body_field in body:
            raise BadRequest(
                "'%s' is readonly but a value was provided: %s"
                % (body_field, user_value),
            )

        log.info("Using default value for '%s': %s", body_field, default_value)
        return default_value

    # field is not readonly and no value was provided
    if user_value is None:
        if not optional:
            raise BadRequest("No value provided for: %s" % body_field)

        # no value for field, but it was optional
        log.info("No value provided for '%s'", defaults_field)
        return None

    log.info("Using provided value for '%s': %s", body_field, user_value)
    return user_value


# Volume handling functions
def is_config_volume(vol):
    if "name" not in vol:
        return False

    if not isinstance(vol["name"], dict):
        return False

    return True


# Notebook YAML processing
def set_notebook_image(notebook, body, defaults):
    """
    If the image is set to readOnly, use only the value from the config
    """
    image_body_field = "image"
    is_custom_image = body.get("customImage", False)
    if is_custom_image:
        image_body_field = "customImage"

    image = get_form_value(body, defaults, image_body_field, "image")
    container = notebook["spec"]["template"]["spec"]["containers"][0]
    container["image"] = image.strip()


def set_notebook_image_pull_policy(notebook, body, defaults):
    container = notebook["spec"]["template"]["spec"]["containers"][0]

    container["imagePullPolicy"] = get_form_value(
        body, defaults, "imagePullPolicy"
    )


def set_server_type(notebook, body, defaults):
    valid_server_types = ["jupyter", "group-one", "group-two"]
    notebook_annotations = notebook["metadata"]["annotations"]
    server_type = get_form_value(body, defaults, "serverType")
    if server_type == "":
        server_type == "jupyter"
    if server_type not in valid_server_types:
        raise BadRequest("'%s' is not a valid server type" % server_type)

    nb_name = get_form_value(body, defaults, "name")
    nb_ns = get_form_value(body, defaults, "namespace")
    rstudio_header = '{"X-RStudio-Root-Path":"/notebook/%s/%s/"}' % (nb_ns,
                                                                     nb_name)
    notebook_annotations[SERVER_TYPE_ANNOTATION] = server_type
    if server_type == "group-one" or server_type == "group-two":
        notebook_annotations[URI_REWRITE_ANNOTATION] = "/"
    if server_type == "group-two":
        notebook_annotations[HEADERS_ANNOTATION] = rstudio_header


def set_notebook_cpu(notebook, body, defaults):
    container = notebook["spec"]["template"]["spec"]["containers"][0]

    cpu = get_form_value(body, defaults, "cpu")
    if cpu and 'nan' in cpu.lower():
        raise BadRequest("Invalid value for cpu: %s" % cpu)

    cpu_limit = get_form_value(body, defaults, "cpuLimit")
    if cpu_limit and 'nan' in cpu_limit.lower():
        raise BadRequest("Invalid value for cpu limit: %s" % cpu_limit)

    limit_factor = utils.load_spawner_ui_config()["cpu"].get("limitFactor")
    if not cpu_limit and limit_factor != "none":
        cpu_limit = str(round((float(cpu) * float(limit_factor)), 1))

    container["resources"]["requests"]["cpu"] = cpu

    if cpu_limit is None or cpu_limit == "":
        # user explicitly asked for no limits
        return

    if float(cpu_limit) < float(cpu):
        raise BadRequest("CPU limit must be greater than the request")

    limits = container["resources"].get("limits", {})
    limits["cpu"] = cpu_limit
    container["resources"]["limits"] = limits


def set_notebook_memory(notebook, body, defaults):
    container = notebook["spec"]["template"]["spec"]["containers"][0]

    memory = get_form_value(body, defaults, "memory")
    if memory and 'nan' in memory.lower():
        raise BadRequest("Invalid value for memory: %s" % memory)

    memory_limit = get_form_value(body, defaults, "memoryLimit")
    if memory_limit and 'nan' in memory_limit.lower():
        raise BadRequest("Invalid value for memory limit: %s" % memory_limit)

    limit_factor = utils.load_spawner_ui_config()["memory"].get("limitFactor")
    if not memory_limit and limit_factor != "none":
        memory_limit = str(
            round((
                float(memory.replace('Gi', '')) * float(
                    limit_factor)), 1)) + "Gi"

    container["resources"]["requests"]["memory"] = memory

    if memory_limit is None or memory_limit == "":
        # user explicitly asked for no limits
        return

    if float(memory_limit.replace('Gi', '')) < float(
            memory.replace('Gi', '')):
        raise BadRequest("Memory limit must be greater than the request")

    limits = container["resources"].get("limits", {})
    limits["memory"] = memory_limit
    container["resources"]["limits"] = limits


def set_notebook_tolerations(notebook, body, defaults):
    tolerations_group_key = get_form_value(body, defaults, "tolerationGroup")

    if tolerations_group_key == "none":
        return

    notebook_tolerations = notebook["spec"]["template"]["spec"]["tolerations"]
    config = utils.load_spawner_ui_config()
    toleration_groups = config.get("tolerationGroup", {}).get("options", [])

    for group in toleration_groups:
        if group["groupKey"] != tolerations_group_key:
            continue

        log.info("Appending Notebook tolerations: %s", group["tolerations"])
        notebook_tolerations.extend(group["tolerations"])
        return

    log.warning(
        "Didn't find any Toleration Group with key '%s' in the config",
        tolerations_group_key,
    )


def set_notebook_affinity(notebook, body, defaults):
    affinity_config_key = get_form_value(body, defaults, "affinityConfig")

    if affinity_config_key == "none":
        return

    notebook_spec = notebook["spec"]["template"]["spec"]
    config = utils.load_spawner_ui_config()
    affinity_configs = config.get("affinityConfig", {}).get("options", [])

    for affinity_config in affinity_configs:
        if affinity_config["configKey"] != affinity_config_key:
            continue

        log.info("Setting Notebook affinity: %s", affinity_config["affinity"])
        notebook_spec["affinity"] = affinity_config["affinity"]
        return

    log.warning(
        "Didn't find any Affinity Config with key '%s' in the config",
        affinity_config_key,
    )


def set_notebook_gpus(notebook, body, defaults):
    gpus = get_form_value(body, defaults, "gpus")

    # Make sure the GPUs value is properly formatted
    if "num" not in gpus:
        raise BadRequest("'gpus' must have a 'num' field")

    if gpus["num"] == "none":
        return

    if "vendor" not in gpus:
        raise BadRequest("'gpus' must have a 'vendor' field")

    # set the gpus annotation
    container = notebook["spec"]["template"]["spec"]["containers"][0]
    vendor = gpus["vendor"]
    try:
        num = str(gpus["num"])
    except ValueError:
        raise BadRequest("gpus.num is not a valid number: %s" % gpus["num"])

    limits = container["resources"].get("limits", {})
    limits[vendor] = num

    container["resources"]["limits"] = limits


def set_notebook_configurations(notebook, body, defaults):
    notebook_labels = notebook["metadata"]["labels"]
    labels = get_form_value(body, defaults, "configurations")

    if not isinstance(labels, list):
        raise BadRequest("Labels for PodDefaults are not list: %s" % labels)

    for label in labels:
        notebook_labels[label] = "true"


def set_notebook_shm(notebook, body, defaults):
    shm = get_form_value(body, defaults, "shm")
    if not shm:
        return

    notebook_spec = notebook["spec"]["template"]["spec"]
    notebook_cont = notebook["spec"]["template"]["spec"]["containers"][0]

    shm_volume = {"name": "dshm", "emptyDir": {"medium": "Memory"}}
    notebook_spec["volumes"].append(shm_volume)

    shm_mnt = {"mountPath": "/dev/shm", "name": "dshm"}
    notebook_cont["volumeMounts"].append(shm_mnt)


def set_notebook_environment(notebook, body, defaults):
    env = get_form_value(body, defaults, "environment")

    # FIXME: Validate the environment?
    env = json.loads(env) if env else {}

    env = [{"name": name, "value": str(value)} for name, value in env.items()]
    notebook["spec"]["template"]["spec"]["containers"][0]["env"] += env


# Volume add functions
def add_notebook_volume(notebook, vol_name, claim, mnt_path):
    spec = notebook["spec"]["template"]["spec"]
    container = notebook["spec"]["template"]["spec"]["containers"][0]

    volume = {"name": vol_name, "persistentVolumeClaim": {"claimName": claim}}
    spec["volumes"].append(volume)

    # Container Mounts
    mnt = {"mountPath": mnt_path, "name": vol_name}
    container["volumeMounts"].append(mnt)
