# -*- coding: utf-8 -*-
#
# Copyright 2012-2015 Spotify AB
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import collections
import importlib
import inspect
import logging
import os
import pickle
import re
import shutil
import sys
import tarfile
import tempfile

from luigi import configuration
from luigi.contrib.external_program import ExternalProgramTask

logger = logging.getLogger("luigi-interface")


class SparkSubmitTask(ExternalProgramTask):
    """
    Template task for running a Spark job

    Supports running jobs on Spark local, standalone, Mesos or Yarn

    See http://spark.apache.org/docs/latest/submitting-applications.html
    for more information

    """

    # Application (.jar or .py file)
    name = None
    entry_class = None
    app = None

    # Only log stderr if spark fails (since stderr is normally quite verbose)
    always_log_stderr = False

    # Spark applications write its logs into stderr
    stream_for_searching_tracking_url = "stderr"

    @property
    def tracking_url_pattern(self):
        if self.deploy_mode == "cluster":
            # in cluster mode client only receives application status once a period of time
            return r"tracking URL: (https?://.*)\s"
        else:
            return r"Bound (?:.*) to (?:.*), and started at (https?://.*)\s"

    def app_options(self):
        """
        Subclass this method to map your task parameters to the app's arguments

        """
        return []

    @property
    def pyspark_python(self):
        return None

    @property
    def pyspark_driver_python(self):
        return None

    @property
    def hadoop_user_name(self):
        return None

    @property
    def spark_version(self):
        return "spark"

    @property
    def spark_submit(self):
        return configuration.get_config().get(self.spark_version, "spark-submit", "spark-submit")

    @property
    def master(self):
        return configuration.get_config().get(self.spark_version, "master", None)

    @property
    def deploy_mode(self):
        return configuration.get_config().get(self.spark_version, "deploy-mode", None)

    @property
    def jars(self):
        return self._list_config(configuration.get_config().get(self.spark_version, "jars", None))

    @property
    def packages(self):
        return self._list_config(configuration.get_config().get(self.spark_version, "packages", None))

    @property
    def py_files(self):
        return self._list_config(configuration.get_config().get(self.spark_version, "py-files", None))

    @property
    def files(self):
        return self._list_config(configuration.get_config().get(self.spark_version, "files", None))

    @property
    def _conf(self):
        conf = collections.OrderedDict(self.conf or {})
        if self.pyspark_python:
            conf["spark.pyspark.python"] = self.pyspark_python
        if self.pyspark_driver_python:
            conf["spark.pyspark.driver.python"] = self.pyspark_driver_python
        return conf

    @property
    def conf(self):
        return self._dict_config(configuration.get_config().get(self.spark_version, "conf", None))

    @property
    def properties_file(self):
        return configuration.get_config().get(self.spark_version, "properties-file", None)

    @property
    def driver_memory(self):
        return configuration.get_config().get(self.spark_version, "driver-memory", None)

    @property
    def driver_java_options(self):
        return configuration.get_config().get(self.spark_version, "driver-java-options", None)

    @property
    def driver_library_path(self):
        return configuration.get_config().get(self.spark_version, "driver-library-path", None)

    @property
    def driver_class_path(self):
        return configuration.get_config().get(self.spark_version, "driver-class-path", None)

    @property
    def executor_memory(self):
        return configuration.get_config().get(self.spark_version, "executor-memory", None)

    @property
    def driver_cores(self):
        return configuration.get_config().get(self.spark_version, "driver-cores", None)

    @property
    def supervise(self):
        return bool(configuration.get_config().get(self.spark_version, "supervise", False))

    @property
    def total_executor_cores(self):
        return configuration.get_config().get(self.spark_version, "total-executor-cores", None)

    @property
    def executor_cores(self):
        return configuration.get_config().get(self.spark_version, "executor-cores", None)

    @property
    def queue(self):
        return configuration.get_config().get(self.spark_version, "queue", None)

    @property
    def num_executors(self):
        return configuration.get_config().get(self.spark_version, "num-executors", None)

    @property
    def archives(self):
        return self._list_config(configuration.get_config().get(self.spark_version, "archives", None))

    @property
    def hadoop_conf_dir(self):
        return configuration.get_config().get(self.spark_version, "hadoop-conf-dir", None)

    def get_environment(self):
        env = os.environ.copy()
        for prop in ("HADOOP_CONF_DIR", "HADOOP_USER_NAME"):
            var = getattr(self, prop.lower(), None)
            if var:
                env[prop] = var
        return env

    def program_environment(self):
        return self.get_environment()

    def program_args(self):
        return self.spark_command() + self.app_command()

    def spark_command(self):
        command = [self.spark_submit]
        command += self._text_arg("--master", self.master)
        command += self._text_arg("--deploy-mode", self.deploy_mode)
        command += self._text_arg("--name", self.name)
        command += self._text_arg("--class", self.entry_class)
        command += self._list_arg("--jars", self.jars)
        command += self._list_arg("--packages", self.packages)
        command += self._list_arg("--py-files", self.py_files)
        command += self._list_arg("--files", self.files)
        command += self._list_arg("--archives", self.archives)
        command += self._dict_arg("--conf", self._conf)
        command += self._text_arg("--properties-file", self.properties_file)
        command += self._text_arg("--driver-memory", self.driver_memory)
        command += self._text_arg("--driver-java-options", self.driver_java_options)
        command += self._text_arg("--driver-library-path", self.driver_library_path)
        command += self._text_arg("--driver-class-path", self.driver_class_path)
        command += self._text_arg("--executor-memory", self.executor_memory)
        command += self._text_arg("--driver-cores", self.driver_cores)
        command += self._flag_arg("--supervise", self.supervise)
        command += self._text_arg("--total-executor-cores", self.total_executor_cores)
        command += self._text_arg("--executor-cores", self.executor_cores)
        command += self._text_arg("--queue", self.queue)
        command += self._text_arg("--num-executors", self.num_executors)
        return command

    def app_command(self):
        if not self.app:
            raise NotImplementedError("subclass should define an app (.jar or .py file)")
        return [self.app] + self.app_options()

    def _list_config(self, config):
        if config and isinstance(config, str):
            return list(map(lambda x: x.strip(), config.split(",")))

    def _dict_config(self, config):
        if config and isinstance(config, str):
            return dict(map(lambda i: i.split("=", 1), config.split("|")))

    def _text_arg(self, name, value):
        if value:
            return [name, value]
        return []

    def _list_arg(self, name, value):
        if value and isinstance(value, (list, tuple)):
            return [name, ",".join(value)]
        return []

    def _dict_arg(self, name, value):
        command = []
        if value and isinstance(value, dict):
            for prop, value in value.items():
                command += [name, "{0}={1}".format(prop, value)]
        return command

    def _flag_arg(self, name, value):
        if value:
            return [name]
        return []


class PySparkTask(SparkSubmitTask):
    """
    Template task for running an inline PySpark job

    Simply implement the ``main`` method in your subclass

    You can optionally define package names to be distributed to the cluster
    with ``py_packages`` (uses luigi's global py-packages configuration by default)

    """

    # Path to the pyspark program passed to spark-submit
    app = os.path.join(os.path.dirname(__file__), "pyspark_runner.py")

    @property
    def name(self):
        return self.__class__.__name__

    @property
    def py_packages(self):
        packages = configuration.get_config().get("spark", "py-packages", None)
        if packages:
            return map(lambda s: s.strip(), packages.split(","))

    @property
    def files(self):
        if self.deploy_mode == "cluster":
            return [self.run_pickle]

    @property
    def pickle_protocol(self):
        return configuration.get_config().getint("spark", "pickle-protocol", pickle.DEFAULT_PROTOCOL)

    def setup(self, conf):
        """
        Called by the pyspark_runner with a SparkConf instance that will be used to instantiate the SparkContext

        :param conf: SparkConf
        """

    def setup_remote(self, sc):
        self._setup_packages(sc)

    def main(self, sc, *args):
        """
        Called by the pyspark_runner with a SparkContext and any arguments returned by ``app_options()``

        :param sc: SparkContext
        :param args: arguments list
        """
        raise NotImplementedError("subclass should define a main method")

    def app_command(self):
        if self.deploy_mode == "cluster":
            pickle_loc = os.path.basename(self.run_pickle)
        else:
            pickle_loc = self.run_pickle
        return [self.app, pickle_loc] + self.app_options()

    def run(self):
        path_name_fragment = re.sub(r"[^\w]", "_", self.name)
        self.run_path = tempfile.mkdtemp(prefix=path_name_fragment)
        self.run_pickle = os.path.join(self.run_path, ".".join([path_name_fragment, "pickle"]))
        with open(self.run_pickle, "wb") as fd:
            # Copy module file to run path.
            module_path = os.path.abspath(inspect.getfile(self.__class__))
            shutil.copy(module_path, os.path.join(self.run_path, "."))
            self._dump(fd)
        try:
            super(PySparkTask, self).run()
        finally:
            shutil.rmtree(self.run_path)

    def _dump(self, fd):
        with self.no_unpicklable_properties():
            if self.__module__ == "__main__":
                d = pickle.dumps(self, protocol=self.pickle_protocol)
                module_name = os.path.basename(sys.argv[0]).rsplit(".", 1)[0]
                d = d.replace(b"c__main__", b"c" + module_name.encode("ascii"))
                fd.write(d)
            else:
                pickle.dump(self, fd, protocol=self.pickle_protocol)

    def _setup_packages(self, sc):
        """
        This method compresses and uploads packages to the cluster

        """
        packages = self.py_packages
        if not packages:
            return
        for package in packages:
            mod = importlib.import_module(package)
            try:
                mod_path = mod.__path__[0]
            except AttributeError:
                mod_path = mod.__file__
            os.makedirs(self.run_path, exist_ok=True)
            tar_path = os.path.join(self.run_path, package + ".tar.gz")
            tar = tarfile.open(tar_path, "w:gz")
            tar.add(mod_path, os.path.basename(mod_path))
            tar.close()
            sc.addPyFile(tar_path)
