"""Extensions of :mod:`luigi.Task`
"""
from builtins import str
from builtins import object
import luigi

from ozelot.etl.targets import ORMTarget
from ozelot import client
from ozelot import config


def get_task_name(task):
    """Return the task name (= class name)
    """
    return task.__class__.__name__


def get_task_param_string(task):
    """Get all parameters of a task as one string

    Returns:
        str: task parameter string
    """
    # get dict str -> str from luigi
    param_dict = task.to_str_params()

    # sort keys, serialize
    items = []
    for key in sorted(param_dict.keys()):
        items.append("'{:s}': '{:s}'".format(key, param_dict[key]))

    return "{" + ", ".join(items) + "}"


class TaskBase(luigi.Task):
    """Common functionality for :mod:`ozelot` tasks.
    """

    def clear(self):
        """Placeholder for clear method.

        This method should clear all output generated by the task, bringing the content database into the
        same state as just before running the task. This is not always easily possible. It is up to the
        developer to define and use clearing functions in such a way that her requirements are met.
        """
        raise NotImplementedError("Function 'clear' not implemented for this task.")

    @classmethod
    def get_name(cls):
        """Return the task name (= class name)
        """
        return cls.__name__

    def get_param_string(self):
        """Get all parameters of a task as one string

        Returns:
            str: task parameter string
        """
        return get_task_param_string(self)


class ORMTask(TaskBase):
    """A task with one target of type :class:`ozelot.models.target.ORMTarget` representing the task and its parameters.

    The default target type can be overwritten by setting the attribute :attr:`target_class`.
    """

    # node styling options for graph diagram
    diagram_style = ["bold"]

    # the target type to use
    target_class = ORMTarget

    def __init__(self, *args, **kwargs):
        """Initialize class attributes
        """
        super(ORMTask, self).__init__(*args, **kwargs)

        # DB session and client
        self._session = None
        self._client = None

    def clear(self):
        raise NotImplementedError("Function 'clear' not implemented for this task.")

    def output(self):
        """Task output: one target of type :class:`ozelot.models.target.ORMTarget` (or whatever else is defined
        as :attr:`target_class`).
        """
        yield self.target_class.from_task(self)

    def mark_complete(self):
        """Mark a task as complete - create all its target markers (by default there is just one).
        """
        for target in self.output():
            target.create()

    def mark_incomplete(self):
        """Mark a task as complete - create all its target markers (by default there is just one).

        Ignore if the task is not complete yet, i.e. a target does not exist.
        """
        for target in self.output():
            if target.exists():
                target.remove()

    @property
    def client(self):
        """Get the DB client associated with the task (open a new one if necessary)

        Returns:
            ozelot.client.Client: DB client
        """
        if self._client is None:
            self._client = client.get_client()

        return self._client

    @property
    def session(self):
        """Get the DB session associated with the task (open a new one if necessary)

        Returns:
            sqlalchemy.orm.session.Session: DB session
        """
        if self._session is None:
            self._session = self.client.create_session()

        return self._session

    def close_session(self, commit=True):
        """Commit and close the DB session associated with this task (no error is raised if None is open)

        Args:
            commit (bool): commit session before closing (default=True)
        """
        if self._session is not None:
            if commit:
                self._session.commit()
            self._session.close()
            self._session = None

    def done(self, commit=True):
        """Task cleanup: mark task as complete, commit and close the session (if open)

        Args:
            commit (bool): commit session before closing (default=True)
        """
        self.close_session(commit)
        self.mark_complete()


class ORMTestTask(ORMTask):
    """A testing task base class
    """

    # node styling options for graph diagram
    diagram_style = ["diagonals"]

    def clear(self):
        """Clearing the task means marking it incomplete (task stores no data)
        """
        self.mark_incomplete()

    def done(self, commit=False):
        """Override :func:`done` of :class:`ORMTask`: don't commit by default

        Args:
            commit (bool): commit session before closing (default=False)
        """
        super(ORMTestTask, self).done(commit)


# helper function: increment a stats counter
def _increment_stats(stats, label):
    count = stats.setdefault(label, 0)
    stats[label] = count + 1


def check_completion(task, mark_incomplete=False, clear=False, return_stats=False):
    """Recursively check if a task and all its requirements are complete

    Args:
        task (derived from luigi.Task): Task to check completion for; check everything 'downstream'
            from that task.

        mark_incomplete (bool): If ``True`` set any task as incomplete for which a requirement
            is found to be incomplete (checked recursively).
            This works only for tasks derived from :class:`ORMTask`.

        clear (bool): If ``True``, call the :func:`clear()` method of any task for which a requirement
            is found to be incomplete (checked recursively). This implies ``mark_incomplete = True``.
            This works only for tasks derived from :class:`ORMTask`.

        return_stats (bool): If ``True``, return task checking statistics in addition to completion status

    Returns:
        bool: ``True`` if the task, all its requirements and (recursively) all their requirements
            are complete, ``False`` otherwise.
    """
    # run recursive task checking, get stats
    to_clear = dict()
    is_complete, stats = _check_completion(task,
                                           mark_incomplete=mark_incomplete,
                                           clear=clear,
                                           stats={},
                                           visited=dict(),
                                           to_clear=to_clear)

    # task clearing needs to happen top-down: because of foreign key constraints, a task can
    # only be cleared once all tasks that require it have been cleared
    while to_clear:
        # find all tasks that we can currently clear - tasks not required by other tasks;
        # iterate over list of keys to be able to modify dict while iterating
        found_clearable_task = False
        for task_id in list(to_clear.keys()):
            v = to_clear[task_id]
            if not v['required_by']:
                # this is a task that can be cleared - no other task requires it
                found_clearable_task = True

                task = v['task']
                if isinstance(task, ORMTask):
                    task.mark_incomplete()
                    task.clear()
                    _increment_stats(stats, 'Cleared')
                    config.logger.info("Cleared task: " + task_id)
                else:
                    config.logger.info('Cannot clear task, not an ORMTask: ' + task_id)

                # remove the task from the list of tasks that need clearing, remove references
                # in the required_by lists of all other tasks; this is not an efficient implementation,
                # O(n^2), could be made O(n) using lookup tables of the task graph
                del to_clear[task_id]
                for w in to_clear.values():
                    w['required_by'].discard(task_id)

        if not found_clearable_task:
            raise RuntimeError("Error in recursive task clearing, no clearable task found")

    config.logger.info("Task completion checking, summary:\n" + str(stats))

    if return_stats:
        return is_complete, stats
    else:
        return is_complete


def _check_completion(task, mark_incomplete, clear, stats, visited, to_clear):
    """Core recursion function for :func:`check_completion`, see there for more documentation

    Args:
        task (luigi.Task): task instance

        mark_incomplete (bool): see :func:`check_completion`

        clear (bool): see :func:`check_completion`

        stats (dict): task checking statistics (counts of incomplete, to be cleared, ...)

        visited (dict): cache for visited tasks: key = task name + parameter string, value = completion status

        to_clear (dict): dict of dicts of tasks to be cleared, key = task id,
            value =  {task: task instance object, required_by: set of task IDs that this task is required by}

    """
    # is this task (recursively) complete?
    task_complete = task.complete()
    is_complete = task_complete

    # task identification: task name plus parameters
    task_id = get_task_name(task) + ' ' + get_task_param_string(task)

    # check any requirements
    for req in task.requires():

        # task/parameter ID string to identify this task instance
        req_id = get_task_name(req) + ' ' + get_task_param_string(req)

        # skip recursion on already visited tasks, get completion status from cache
        if req_id in visited:
            req_complete = visited[req_id]
        else:
            req_complete, _ = _check_completion(task=req,
                                                mark_incomplete=mark_incomplete,
                                                clear=clear,
                                                stats=stats,
                                                visited=visited,
                                                to_clear=to_clear)
            visited[req_id] = req_complete

        # add any incomplete requirements to the list of tasks to clear, noting the current task as parent (required by)
        if clear and not req_complete:
            clear_entry = to_clear.setdefault(req_id, dict(task=req, required_by=set()))
            clear_entry['required_by'].add(task_id)

        is_complete &= req_complete

    if not is_complete:
        if task_complete:
            config.logger.info("Task complete but requirements incomplete: " + task_id)
        else:
            config.logger.info("Task incomplete: " + task_id)
        _increment_stats(stats, 'Incomplete tasks')

        if mark_incomplete:
            if isinstance(task, ORMTask):
                task.mark_incomplete()
                _increment_stats(stats, 'Marked incomplete')
                config.logger.info("Marked task incomplete: " + task_id)
            else:
                config.logger.info('Cannot mark task incomplete, not an ORMTask: ' + task_id)

    else:
        _increment_stats(stats, 'Complete tasks')
        config.logger.debug("Task complete: " + task_id)

    # if we want to clear and the current task is not in the dict of tasks to clear,
    # it is the root task, add it with no parent (required by) tasks
    if clear and not is_complete and task_id not in to_clear:
        to_clear[task_id] = dict(task=task, required_by=set())

    return is_complete, stats


class ORMObjectCreatorMixin(object):
    """Functionality for tasks that are responsible for creating a certain class of ORM objects

    If a class (or multiple classes) of ORM objects are created only by a single task, clearing
    of the task output is simplified, as is bulk storing of task output to the database.
    """

    #: *list(ORM object classes)*: ORM object classes that are created by this task
    object_classes = list()

    def clear(
            self  # type: ORMTask
    ):
        """Delete all objects created by this task.

        Iterate over `self.object_classes` and delete all objects of the listed classes.
        """
        # mark this task as incomplete
        self.mark_incomplete()

        # delete objects
        for object_class in self.object_classes:
            self.session.query(object_class).delete()

        self.close_session()


class ORMWrapperTask(ORMTask):
    """A wrapper task that bundles other tasks and does not run its own code

    The task's :func:`run` method marks the task as complete and does nothing else.
    Since its requirements are checked before marking the task complete, it serves as
    a wrapper around other tasks, for logical grouping.

    The tasks :func:`complete` method only returns complete if the task itself
    is marked complete and all requirements are complete. This behaviour is useful
    when adding tasks to a wrapper task during development -- adding a new tasks
    triggers re-running the wrapper task.
    """

    # node styling options for graph diagram
    diagram_style = ["dotted"]

    def clear(self):
        """Mark task as incomplete, do nothing else.
        """
        self.mark_incomplete()

    def run(self):
        """Mark task as complete, do nothing else.
        """
        self.mark_complete()

    def complete(self):
        """Task is complete if completion marker is set and all requirements are complete
        """
        is_complete = super(ORMWrapperTask, self).complete()
        for req in self.requires():
            is_complete &= req.complete()

        return is_complete


class InputFileTask(TaskBase):
    """A base class for tasks representing input files.

    This task has an empty run method (does nothing) and specifies the input file as task output.
    Existence of the input file signals task completion; this is checked by , :mod:`luigi`.

    It is useful to supply a `load` method in the derived classes, which loads the respective data file
    to a usable format (e.g., to a :class:`pandas.DataFrame`).
    """

    # the path to the input file, to be defined in the derived class
    input_file = None

    # node styling options for graph diagram
    diagram_style = ["rounded"]

    def output(self):
        """Task output: one :class:`luigi.LocalTarget` for the input file
        """
        if self.input_file is None:
            raise RuntimeError("Please specify 'input_file' for the task class")

        return [luigi.LocalTarget(self.input_file)]

    def clear(self):
        """Input files should never be cleared - log the call but do nothing else.
        """
        config.logger.info("Called 'clear' for input file task {:s} - doing nothing.".format(self.get_name()))

    def run(self):
        """Run method - does nothing
        """
        pass

    def load(self):
        """A method to load the input file to a usable format, to be defined in the derived class.
        """
        raise NotImplementedError()
