
import logging
import shlex
from typing import Dict, List

from magnus import defaults, integration, utils
from magnus.executor import BaseExecutor
from magnus.graph import Graph
from magnus.nodes import BaseNode
from ruamel.yaml import YAML

logger = logging.getLogger(defaults.NAME)

try:
    from kubernetes import client
    from kubernetes import config as k8s_config
    from kubernetes.client import (V1EnvVar, V1EnvVarSource,
                                   V1PersistentVolumeClaimVolumeSource,
                                   V1SecretKeySelector)
except ImportError as _e:
    msg = (
        "Kubernetes Dependencies have not been installed!!"
    )
    raise Exception(msg) from _e


class K8sJobExecutor(BaseExecutor):

    service_name = "k8s-job"

    class Config(BaseExecutor.Config):
        config_path: str
        docker_image: str
        namespace: str = "default"
        cpu_limit: str = "250m"
        memory_limit: str = "1G"
        gpu_limit: int = 0
        gpu_vendor: str = "nvidia.com/gpu"
        cpu_request: str = ""
        memory_request: str = ""
        active_deadline_seconds: int = 60 * 60 * 2  # 2 hours
        ttl_seconds_after_finished: int = 60  #  1 minute
        image_pull_policy: str = "Always"
        secrets_from_k8s: dict = {}  # EnvVar=SecretName:Key
        persistent_volumes: dict = {}  # volume-name:mount_path
        labels: Dict[str, str] = {}

    def __init__(self, config: dict = None):
        super().__init__(config)
        self.persistent_volumes = {}

        for i, (claim, mount_path) in enumerate(self.config.persistent_volumes.items()):
            self.persistent_volumes[f"executor-{i}"] = (claim, mount_path)

    def prepare_for_graph_execution(self):
        """
        This method would be called prior to calling execute_graph.
        Perform any steps required before doing the graph execution.

        The most common implementation is to prepare a run log for the run if the run uses local interactive compute.

        But in cases of actual rendering the job specs (eg: AWS step functions, K8's) we need not do anything.
        """

        integration.validate(self, self.run_log_store)
        integration.configure_for_traversal(self, self.run_log_store)

        integration.validate(self, self.catalog_handler)
        integration.configure_for_traversal(self, self.catalog_handler)

        integration.validate(self, self.secrets_handler)
        integration.configure_for_traversal(self, self.secrets_handler)

        integration.validate(self, self.experiment_tracker)
        integration.configure_for_traversal(self, self.experiment_tracker)

    def prepare_for_node_execution(self):
        """
        Perform any modifications to the services prior to execution of the node.

        Args:
            node (Node): [description]
            map_variable (dict, optional): [description]. Defaults to None.
        """

        integration.validate(self, self.run_log_store)
        integration.configure_for_execution(self, self.run_log_store)

        integration.validate(self, self.catalog_handler)
        integration.configure_for_execution(self, self.catalog_handler)

        integration.validate(self, self.secrets_handler)
        integration.configure_for_execution(self, self.secrets_handler)

        integration.validate(self, self.experiment_tracker)
        integration.configure_for_execution(self, self.experiment_tracker)

        self._set_up_run_log(exists_ok=True)

    @property
    def _client(self):
        k8s_config.load_kube_config(config_file=self.config.config_path)
        return client

    def execute_job(self, node: BaseNode):
        command = utils.get_job_execution_command(self, node)
        logger.info(f'Triggering a kubernetes job with : {command}')

        self.config.labels['job_name'] = self.run_id

        k8s_batch = self._client.BatchV1Api()

        cpu_limit = self.config.cpu_limit
        memory_limit = self.config.memory_limit

        cpu_request = self.config.cpu_request or cpu_limit
        memory_request = self.config.memory_request or memory_limit

        gpu_limit = str(self.config.gpu_limit)  # Should be something like nvidia -etc

        limits = {"cpu": cpu_limit, "memory": memory_limit, self.config.gpu_vendor: gpu_limit}
        requests = {"cpu": cpu_request, "memory": memory_request}
        resources = {"limits": limits, "requests": requests}

        environment_variables = []
        for secret_env, k8_secret in self.config.secrets_from_k8s.items():
            try:
                secret_name, key = k8_secret.split(':')
            except Exception as _e:
                msg = (
                    "K8's secret should be of format EnvVar=SecretName:Key"
                )
                raise Exception(msg) from _e
            secret_as_env = V1EnvVar(name=secret_env, value_from=V1EnvVarSource(
                secret_key_ref=V1SecretKeySelector(name=secret_name, key=key)))
            environment_variables.append(secret_as_env)

        overridden_params = utils.get_user_set_parameters()
        # The parameters present in the environment override the parameters present in the parameters file
        # The values are coerced to be strings, hopefully they will be fine on the other side.
        for k, v in overridden_params.items():
            environment_variables.append(V1EnvVar(name=defaults.PARAMETER_PREFIX+k, value=str(v)))

        pod_volumes = []
        volume_mounts = []
        for claim_name, (claim, mount_path) in self.persistent_volumes.items():
            pod_volumes.append(self._client.V1Volume(
                name=claim_name, persistent_volume_claim=V1PersistentVolumeClaimVolumeSource(claim_name=claim)))
            volume_mounts.append(self._client.V1VolumeMount(name=claim_name, mount_path=mount_path))

        base_container = self._client.V1Container(
            name=self.run_id,
            image=self.config.docker_image,
            command=shlex.split(command),
            resources=resources,
            env=environment_variables,
            image_pull_policy="Always",
            volume_mounts=volume_mounts or None
        )

        pod_spec = self._client.V1PodSpec(volumes=pod_volumes or None,
                                          restart_policy='Never',
                                          containers=[base_container])

        pod_template = self._client.V1PodTemplateSpec(metadata=client.V1ObjectMeta(
            labels=self.config.labels, annotations={"sidecar.istio.io/inject": "false"}), spec=pod_spec)

        job_spec = client.V1JobSpec(template=pod_template, backoff_limit=2,
                                    ttl_seconds_after_finished=self.config.ttl_seconds_after_finished)
        job_spec.active_deadline_seconds = self.config.active_deadline_seconds

        job = client.V1Job(
            api_version='batch/v1',
            kind='Job',
            metadata=client.V1ObjectMeta(name=self.run_id),
            spec=job_spec)

        logger.debug(f'Submitting kubernetes job: {job}')

        try:
            response = k8s_batch.create_namespaced_job(
                body=job, namespace=self.config.namespace, _preload_content=False, pretty=True)
            print(f'Kubernetes job {self.run_id} created')
            logger.debug(f'Kubernetes job response: {response}')
        except Exception as e:
            logger.exception(e)
            raise

    def execute_node(self, node: BaseNode, map_variable: dict = None, **kwargs):
        step_log = self.run_log_store.create_step_log(node.name, node._get_step_log_name(map_variable))

        self.add_code_identities(node=node, step_log=step_log)

        step_log.step_type = node.node_type
        step_log.status = defaults.PROCESSING
        self.run_log_store.add_step_log(step_log, self.run_id)

        super()._execute_node(node, map_variable=map_variable, **kwargs)

        step_log = self.run_log_store.get_step_log(node._get_step_log_name(map_variable), self.run_id)
        if step_log.status == defaults.FAIL:
            raise Exception(f'Step {node.name} failed')

    def execute_graph(self, dag: Graph, map_variable: dict = None, **kwargs):
        msg = (
            "This executor is not supported to execute any graphs but only jobs (functions or notebooks)"
        )
        raise NotImplementedError(msg)

    def send_return_code(self, stage='traversal'):
        """
        Convenience function used by pipeline to send return code to the caller of the cli

        Raises:
            Exception: If the pipeline execution failed
        """
        if stage != 'traversal':  # traversal does no actual execution, so return code is pointless
            run_id = self.run_id

            run_log = self.run_log_store.get_run_log_by_id(run_id=run_id, full=False)
            if run_log.status == defaults.FAIL:
                raise Exception('Pipeline execution failed')
