# coding: utf-8

"""
Custom luigi base task definitions.
"""


__all__ = ["Task", "WrapperTask", "ExternalTask"]


import sys
import socket
import time
import math
import logging
from collections import OrderedDict
from contextlib import contextmanager
from abc import abstractmethod
from inspect import getargspec

import luigi
import luigi.util
import six

from law.config import Config
from law.parameter import NO_STR, CSVParameter
from law.parser import global_cmdline_values
from law.target.file import FileSystemTarget, localize_file_targets
from law.target.collection import TargetCollection
from law.parser import root_task
from law.util import (
    abort, colored, uncolored, make_list, query_choice, multi_match, flatten, check_bool_flag,
    BaseStream, human_duration, quote_cmd, patch_object,
)


logger = logging.getLogger(__name__)


class BaseRegister(luigi.task_register.Register):

    def __new__(metacls, classname, bases, classdict):
        # default attributes
        classdict.setdefault("exclude_index", False)

        # union "exclude_params_*" sets with those of all base classes
        for base in bases:
            for attr, base_params in vars(base).items():
                if isinstance(base_params, set) and attr.startswith("exclude_params_"):
                    params = classdict.setdefault(attr, set())
                    if isinstance(params, set):
                        params.update(base_params)

        return super(BaseRegister, metacls).__new__(metacls, classname, bases, classdict)


@six.add_metaclass(BaseRegister)
class BaseTask(luigi.Task):

    exclude_index = True
    exclude_params_index = set()
    exclude_params_req = set()
    exclude_params_req_set = set()
    exclude_params_req_get = set()

    @staticmethod
    def resource_name(name, host=None):
        if host is None:
            host = socket.gethostname().partition(".")[0]
        return "{}_{}".format(host, name)

    @classmethod
    def get_param_values(cls, *args, **kwargs):
        values = super(BaseTask, cls).get_param_values(*args, **kwargs)
        if callable(cls.modify_param_values):
            return cls.modify_param_values(OrderedDict(values)).items()
        else:
            return values

    # method that can be implemented to update parameters via get_param_values before instantiation
    # example:
    # @classmethod
    # def modify_param_values(cls, params):
    #     return params
    modify_param_values = None

    @classmethod
    def req(cls, *args, **kwargs):
        return cls(**cls.req_params(*args, **kwargs))

    @classmethod
    def req_params(cls, inst, _exclude=None, _prefer_cli=None, **kwargs):
        # common/intersection params
        params = luigi.util.common_params(inst, cls)

        # determine parameters to exclude
        _exclude = set() if _exclude is None else set(make_list(_exclude))

        # always exclude interactive parameters
        _exclude |= set(inst.interactive_params)

        # also use this class' req and req_get sets
        # and the req and req_set sets of the instance's class
        _exclude.update(cls.exclude_params_req, cls.exclude_params_req_get)
        _exclude.update(inst.exclude_params_req, inst.exclude_params_req_set)

        # remove excluded parameters
        for name in list(params.keys()):
            if multi_match(name, _exclude, any):
                del params[name]

        # add kwargs
        params.update(kwargs)

        # remove params that are preferably set via cli class arguments
        if _prefer_cli:
            cls_args = []
            prefix = cls.task_family + "_"
            if luigi.cmdline_parser.CmdlineParser.get_instance():
                for key in global_cmdline_values().keys():
                    if key.startswith(prefix):
                        cls_args.append(key[len(prefix):])
            for name in make_list(_prefer_cli):
                if name in params and name in cls_args:
                    del params[name]

        return params

    def complete(self):
        outputs = [t for t in flatten(self.output()) if not t.optional]

        if len(outputs) == 0:
            logger.warning("task {!r} has either no non-optional outputs or no custom complete() "
                "method".format(self))
            return False

        return all(t.exists() for t in outputs)

    @property
    def live_task_id(self):
        """
        The task id depends on the task family and parameters, and is generated by luigi once in the
        constructor. As the latter may change, this property returns to the id with the current set
        of parameters.
        """
        # create a temporary dictionary of param_kwargs that is patched for the duration of the
        # call to create the string representation of the parameters
        param_kwargs = {attr: getattr(self, attr) for attr in self.param_kwargs}
        # only_public was introduced in 2.8.0, so check if that arg exists
        str_params_kwargs = {"only_significant": True}
        if "only_public" in getargspec(self.to_str_params).args:
            str_params_kwargs["only_public"] = True
        with patch_object(self, "param_kwargs", param_kwargs):
            str_params = self.to_str_params(**str_params_kwargs)

        # create the task id
        task_id = luigi.task.task_id_str(self.get_task_family(), str_params)

        return task_id

    def walk_deps(self, max_depth=-1, order="level"):
        # see https://en.wikipedia.org/wiki/Tree_traversal
        if order not in ("level", "pre"):
            raise ValueError("unknown traversal order '{}', use 'level' or 'pre'".format(order))

        tasks = [(self, 0)]
        while len(tasks):
            task, depth = tasks.pop(0)
            if max_depth >= 0 and depth > max_depth:
                continue
            deps = flatten(task.requires())

            yield (task, deps, depth)

            deps = ((d, depth + 1) for d in deps)
            if order == "level":
                tasks[len(tasks):] = deps
            elif order == "pre":
                tasks[:0] = deps

    def cli_args(self, exclude=None, replace=None):
        if exclude is None:
            exclude = set()
        if replace is None:
            replace = {}

        # always exclude interactive parameters
        exclude |= set(self.interactive_params)

        args = []
        for name, param in self.get_params():
            if multi_match(name, exclude, any):
                continue
            raw = replace.get(name, getattr(self, name))
            val = param.serialize(raw)
            arg = "--{}".format(name.replace("_", "-"))
            args.extend([arg, quote_cmd([val])])

        return args

    @abstractmethod
    def run(self):
        return


class Register(BaseRegister):

    def __call__(cls, *args, **kwargs):
        inst = super(Register, cls).__call__(*args, **kwargs)

        # check for interactive parameters
        for param in inst.interactive_params:
            value = getattr(inst, param)
            if value:
                skip_abort = False
                try:
                    logger.debug("evaluating interactive parameter '{}' with value '{}'".format(
                        param, value))
                    skip_abort = getattr(inst, "_" + param)(value)
                except KeyboardInterrupt:
                    print("\naborted")

                # abort the process if not explicitly skipped
                if not skip_abort:
                    abort(exitcode=0)

        return inst


@six.add_metaclass(Register)
class Task(BaseTask):

    log_file = luigi.Parameter(default=NO_STR, significant=False, description="a custom log file, "
        "default: <task.default_log_file>")
    print_deps = CSVParameter(default=[], significant=False, description="print task dependencies, "
        "do not run any task, the passed numbers set the recursion depth (0 means non-recursive)")
    print_status = CSVParameter(default=[], significant=False, description="print the task status, "
        "do not run any task, the passed numbers set the recursion depth (0 means non-recursive) "
        "and optionally the collection depth")
    print_output = CSVParameter(default=[], significant=False, description="print a flat list of "
        "output targets, do not run any task, the passed number sets the recursion depth (0 means "
        "non-recursive")
    remove_output = CSVParameter(default=[], significant=False, description="remove all outputs, "
        "do not run any task, the passed number sets the recursion depth (0 means non-recursive)")

    interactive_params = ["print_deps", "print_status", "print_output", "remove_output"]

    message_cache_size = 10

    exclude_index = True
    exclude_params_req = set(interactive_params)
    exclude_params_repr = set()

    def __init__(self, *args, **kwargs):
        super(Task, self).__init__(*args, **kwargs)

        # cache for messages published to the scheduler
        self._message_cache = []

        # cache for the last progress published to the scheduler
        self._last_progress_percentage = None

    @property
    def default_log_file(self):
        return "-"

    def is_root_task(self):
        return root_task() == self

    def publish_message(self, *args):
        msg = " ".join(str(arg) for arg in args)
        print(msg)
        sys.stdout.flush()

        self._publish_message(*args)

    def _publish_message(self, *args):
        msg = " ".join(str(arg) for arg in args)

        # add to message cache and handle overflow
        msg = uncolored(msg)
        self._message_cache.append(msg)
        if self.message_cache_size >= 0:
            end = max(len(self._message_cache) - self.message_cache_size, 0)
            del self._message_cache[:end]

        # set status message using the current message cache
        self.set_status_message("\n".join(self._message_cache))

    def create_message_stream(self, *args, **kwargs):
        return TaskMessageStream(self, *args, **kwargs)

    @contextmanager
    def publish_step(self, msg, success_message="done", fail_message="failed", runtime=False):
        self.publish_message(msg)
        success = False
        t0 = time.time()
        try:
            yield
            success = True
        finally:
            msg = success_message if success else fail_message
            if runtime:
                diff = time.time() - t0
                msg = "{} (took {})".format(msg, human_duration(seconds=diff))
            self.publish_message(msg)

    def publish_progress(self, percentage):
        percentage = int(math.floor(percentage))
        if percentage != self._last_progress_percentage:
            self._last_progress_percentage = percentage
            self.set_progress_percentage(percentage)

    def create_progress_callback(self, n_total, reach=(0, 100)):
        def make_callback(n, start, end):
            def callback(i):
                self.publish_progress(start + (i + 1) / float(n) * (end - start))
            return callback

        if isinstance(n_total, (list, tuple)):
            width = 100. / len(n_total)
            reaches = [(width * i, width * (i + 1)) for i in range(len(n_total))]
            return n_total.__class__(make_callback(n, *r) for n, r in zip(n_total, reaches))
        else:
            return make_callback(n_total, *reach)

    def __repr__(self):
        return self.repr(color=False)

    def repr(self, all_params=False, color=None):
        if color is None:
            color = Config.instance().get_expanded_boolean("task", "colored_repr")

        family = self._repr_family(self.task_family, color=color)

        parts = [
            self._repr_param(*pair, color=color)
            for pair in self._repr_params(all_params=all_params)
        ] + [
            self._repr_flag(flag, color=color)
            for flag in self._repr_flags()
        ]

        return "{}({})".format(family, ", ".join(parts))

    def colored_repr(self, all_params=False):
        # deprecation warning until v0.1
        logger.warning("the use of {0}.colored_repr() is deprecated, please use "
            "{0}.repr(color=True) instead".format(self.__class__.__name__))

        return self.repr(all_params=all_params, color=True)

    def _repr_params(self, all_params=False):
        # build key value pairs of all significant parameters
        params = self.get_params()
        param_values = self.get_param_values(params, [], self.param_kwargs)
        param_objs = dict(params)

        if all_params:
            exclude = set()
        else:
            exclude = self.exclude_params_repr | self.inst_exclude_params_repr()

        pairs = []
        for param_name, param_value in param_values:
            if param_objs[param_name].significant and not multi_match(param_name, exclude):
                pairs.append((param_name, param_objs[param_name].serialize(param_value)))

        return pairs

    def _repr_flags(self):
        return []

    def inst_exclude_params_repr(self):
        return set()

    @classmethod
    def _repr_family(cls, family, color=False):
        return colored(family, "green") if color else family

    @classmethod
    def _repr_param(cls, name, value, color=False):
        return "{}={}".format(colored(name, color="blue", style="bright") if color else name, value)

    @classmethod
    def _repr_flag(cls, name, color=False):
        return colored(name, color="magenta") if color else name

    def _print_deps(self, args):
        return print_task_deps(self, *args)

    def _print_status(self, args):
        return print_task_status(self, *args)

    def _print_output(self, args):
        return print_task_output(self, *args)

    def _remove_output(self, args):
        return remove_task_output(self, *args)

    def localize_input(self, *args, **kwargs):
        return localize_file_targets(self.input(), *args, **kwargs)

    def localize_output(self, *args, **kwargs):
        return localize_file_targets(self.output(), *args, **kwargs)


class WrapperTask(Task):
    """
    Use for tasks that only wrap other tasks and that by definition are done
    if all their requirements exist.
    """

    exclude_index = True

    def _repr_flags(self):
        return super(WrapperTask, self)._repr_flags() + ["wrapper"]

    def complete(self):
        return all(task.complete() for task in flatten(self.requires()))

    def run(self):
        return


class ExternalTask(Task):

    exclude_index = True

    run = None

    def _repr_flags(self):
        return super(ExternalTask, self)._repr_flags() + ["external"]


class TaskMessageStream(BaseStream):

    def __init__(self, task, stdout=True):
        super(TaskMessageStream, self).__init__()
        self.task = task
        self.stdout = stdout

    def _write(self, *args):
        if self.stdout:
            self.task.publish_message(*args)
        else:
            self.task._publish_message(*args)


def print_task_deps(task, max_depth=1):
    max_depth = int(max_depth)

    print("print task dependencies with max_depth {}\n".format(max_depth))

    ind = "|   "
    for dep, _, depth in task.walk_deps(max_depth=max_depth, order="pre"):
        print(depth * ind + "> " + dep.repr(color=True))


def print_task_status(task, max_depth=0, target_depth=0, flags=None):
    max_depth = int(max_depth)
    target_depth = int(target_depth)
    if flags:
        flags = tuple(flags.lower().split("-"))

    print("print task status with max_depth {} and target_depth {}".format(
        max_depth, target_depth))

    done = []
    ind = "|   "
    for dep, _, depth in task.walk_deps(max_depth=max_depth, order="pre"):
        offset = depth * ind
        print(offset)
        print("{}> check status of {}".format(offset, dep.repr(color=True)))
        offset += ind

        if dep in done:
            print(offset + "- " + colored("outputs already checked", "yellow"))
            continue

        done.append(dep)

        for outp in flatten(dep.output()):
            print("{}- {}".format(offset, outp.repr(color=True)))

            status_text = outp.status_text(max_depth=target_depth, flags=flags, color=True)
            status_lines = status_text.split("\n")
            status_text = status_lines[0]
            for line in status_lines[1:]:
                status_text += "\n" + offset + "     " + line
            print("{}  -> {}".format(offset, status_text))


def print_task_output(task, max_depth=0):
    max_depth = int(max_depth)

    print("print task output with max_depth {}\n".format(max_depth))

    def print_target(target):
        if isinstance(target, FileSystemTarget):
            print(target.uri())
        else:
            logger.warning("target listing not yet implemented for {}".format(target.__class__))

    done = []
    for dep, _, depth in task.walk_deps(max_depth=max_depth, order="pre"):
        done.append(dep)

        for outp in flatten(dep.output()):
            if isinstance(outp, TargetCollection):
                for t in outp._flat_target_list:
                    print_target(t)
            else:
                print_target(outp)


def remove_task_output(task, max_depth=0, mode=None, include_external=False):
    max_depth = int(max_depth)

    print("remove task output with max_depth {}".format(max_depth))

    include_external = check_bool_flag(include_external)
    if include_external:
        print("include external tasks")

    # determine the mode, i.e., interactive, dry, all
    modes = ["i", "d", "a"]
    mode_names = ["interactive", "dry", "all"]
    if mode and mode not in modes:
        raise Exception("unknown removal mode '{}'".format(mode))
    if not mode:
        mode = query_choice("removal mode?", modes, default="i", descriptions=mode_names)
    mode_name = mode_names[modes.index(mode)]
    print("selected " + colored(mode_name + " mode", "blue", style="bright"))

    done = []
    ind = "|   "
    for dep, _, depth in task.walk_deps(max_depth=max_depth, order="pre"):
        offset = depth * ind
        print(offset)
        print("{}> remove output of {}".format(offset, dep.repr(color=True)))
        offset += ind

        if not include_external and isinstance(dep, ExternalTask):
            print(offset + "- " + colored("task is external, skip", "yellow"))
            continue

        if dep in done:
            print(offset + "- " + colored("outputs already removed", "yellow"))
            continue

        if mode == "i":
            task_mode = query_choice(offset + "  remove outputs?", ["y", "n", "a"], default="y",
                descriptions=["yes", "no", "all"])
            if task_mode == "n":
                continue

        done.append(dep)

        for outp in flatten(dep.output()):
            print("{}- {}".format(offset, outp.repr(color=True)))

            if mode == "d":
                print(offset + "  " + colored("dry removed", "yellow"))
                continue

            elif mode == "i" and task_mode != "a":
                if query_choice(offset + "  remove?", ("y", "n"), default="n") == "n":
                    print(offset + "  " + colored("skipped", "yellow"))
                    continue

            outp.remove()
            print(offset + "  " + colored("removed", "red", style="bright"))
