# -*- coding: utf-8 -*-
"""The task manager."""

from __future__ import unicode_literals

import collections
import heapq
import threading
import time

from dfvfs.lib import definitions as dfvfs_definitions

from plaso.containers import tasks
from plaso.lib import definitions
from plaso.engine import processing_status
from plaso.engine import profilers
from plaso.multi_processing import logger


class _PendingMergeTaskHeap(object):
  """Heap to manage pending merge tasks."""

  def __init__(self):
    """Initializes a pending merge task heap."""
    super(_PendingMergeTaskHeap, self).__init__()
    self._heap = []
    self._task_identifiers = set()

  def __contains__(self, task_identifier):
    """Checks for an task identifier being present in the heap.

    Args:
      task_identifier (str): task identifier to check for.

    Returns:
      bool: True if the task with the given identifier is present in the heap.
    """
    return task_identifier in self._task_identifiers

  def __len__(self):
    """Determines the number of tasks on the heap.

    Returns:
      int: number of tasks on the heap.
    """
    return len(self._heap)

  def PeekTask(self):
    """Retrieves the first task from the heap without removing it.

    Returns:
      Task: task or None if the heap is empty.
    """
    try:
      _, task = self._heap[0]

    except IndexError:
      return None

    return task

  def PopTask(self):
    """Retrieves and removes the first task from the heap.

    Returns:
      Task: the task or None if the heap is empty.
    """
    try:
      _, task = heapq.heappop(self._heap)

    except IndexError:
      return None
    self._task_identifiers.remove(task.identifier)
    return task

  def PushTask(self, task):
    """Pushes a task onto the heap.

    Args:
      task (Task): task.

    Raises:
      ValueError: if the size of the storage file is not set in the task.
    """
    storage_file_size = getattr(task, 'storage_file_size', None)
    if not storage_file_size:
      raise ValueError('Task storage file size not set.')

    if task.file_entry_type == dfvfs_definitions.FILE_ENTRY_TYPE_DIRECTORY:
      weight = 1
    else:
      weight = storage_file_size

    task.merge_priority = weight

    heap_values = (weight, task)
    heapq.heappush(self._heap, heap_values)
    self._task_identifiers.add(task.identifier)


class TaskManager(object):
  """Manages tasks and tracks their completion and status.

  A task being tracked by the manager must be in exactly one of the
  following states:

  * abandoned: a task assumed to be abandoned because a tasks that has been
      queued or was processing exceeds the maximum inactive time.
  * merging: a task that is being merged by the engine.
  * pending_merge: the task has been processed and is ready to be merged with
      the session storage.
  * processed: a worker has completed processing the task, but it is not ready
      to be merged into the session storage.
  * processing: a worker is processing the task.
  * queued: the task is waiting for a worker to start processing it. It is also
      possible that a worker has already completed the task, but no status
      update was collected from the worker while it processed the task.

  Once the engine reports that a task is completely merged, it is removed
  from the task manager.

  Tasks are considered "pending" when there is more work that needs to be done
  to complete these tasks. Pending applies to tasks that are:
  * not abandoned;
  * abandoned, but need to be retried.

  Abandoned tasks without corresponding retry tasks are considered "failed"
  when the foreman is done processing.
  """

  # Stop pylint from reporting:
  # Context manager 'lock' doesn't implement __enter__ and __exit__.
  # pylint: disable=not-context-manager

  # Consider a task inactive after 5 minutes of no activity.
  _TASK_INACTIVE_TIME = 5.0 * 60.0

  def __init__(self):
    """Initializes a task manager."""
    super(TaskManager, self).__init__()
    self._lock = threading.Lock()

    # This dictionary maps task identifiers to tasks that have been abandoned,
    # as no worker has reported processing the task in the expected interval.
    self._tasks_abandoned = {}

    # The latest processing time observed in a task. This value is set to
    # the current time to not have to handle None as a special case.
    self._latest_task_processing_time = int(
        time.time() * definitions.MICROSECONDS_PER_SECOND)

    # This dictionary maps task identifiers to tasks that are being merged
    # by the foreman.
    self._tasks_merging = {}

    # Use ordered dictionary to preserve the order in which tasks were added.
    # This dictionary maps task identifiers to tasks that are waiting to be
    # processed.
    self._tasks_queued = collections.OrderedDict()

    self._tasks_pending_merge = _PendingMergeTaskHeap()

    # This dictionary maps task identifiers to tasks that are currently
    # being processed by a worker.
    self._tasks_processing = {}

    self._tasks_profiler = None

    # TODO: implement a limit on the number of tasks.
    self._total_number_of_tasks = 0

  def _AbandonInactiveProcessingTasks(self):
    """Marks processing tasks that exceed the inactive time as abandoned.

    This method does not lock the manager and should be called by a method
    holding the manager lock.
    """
    if self._tasks_processing:
      inactive_time = time.time() - self._TASK_INACTIVE_TIME
      inactive_time = int(inactive_time * definitions.MICROSECONDS_PER_SECOND)

      # Abandon all tasks after they're identified so as not to modify the
      # dict while iterating over it.
      tasks_to_abandon = []
      for task_identifier, task in iter(self._tasks_processing.items()):
        if task.last_processing_time < inactive_time:
          logger.debug('Abandoned processing task: {0:s}.'.format(
              task_identifier))

          self.SampleTaskStatus(task, 'abandoned_processing')
          tasks_to_abandon.append((task_identifier, task))

      for task_identifier, task in tasks_to_abandon:
        self._tasks_abandoned[task_identifier] = task
        del self._tasks_processing[task_identifier]

  def _AbandonQueuedTasks(self):
    """Marks queued tasks abandoned.

    This method does not lock the manager and should be called by a method
    holding the manager lock.
    """
    for task_identifier, task in iter(self._tasks_queued.items()):
      logger.debug('Abandoned queued task: {0:s}.'.format(task_identifier))

      self._tasks_abandoned[task_identifier] = task
      del self._tasks_queued[task_identifier]

  def _GetTaskPendingRetry(self):
    """Retrieves an abandoned task that should be retried.

    This method does not lock the manager and should be called by a method
    holding the manager lock.

    Returns:
      Task: a task that was abandoned but should be retried or None if there are
          no abandoned tasks that should be retried.
    """
    for task in self._tasks_abandoned.values():
      if not task.has_retry:
        return task

    return None

  def _HasTasksPendingMerge(self):
    """Determines if there are tasks waiting to be merged.

    This method does not lock the manager and should be called by a method
    holding the manager lock.

    Returns:
      bool: True if there are abandoned tasks that need to be retried.
    """
    return bool(self._tasks_pending_merge)

  def _HasTasksPendingRetry(self):
    """Determines if there are abandoned tasks that still need to be retried.

    This method does not lock the manager and should be called by a method
    holding the manager lock.

    Returns:
      bool: True if there are abandoned tasks that need to be retried.
    """
    return bool(self._GetTaskPendingRetry())

  def _UpdateLatestProcessingTime(self, task):
    """Updates the latest processing time of the task manager from the task.

    This method does not lock the manager and should be called by a method
    holding the manager lock.

    Args:
      task (Task): task to update the processing time of.
    """
    self._latest_task_processing_time = max(
        self._latest_task_processing_time, task.last_processing_time)

  def CheckTaskToMerge(self, task):
    """Checks if the task should be merged.

    Args:
      task (Task): task.

    Returns:
      bool: True if the task should be merged.

    Raises:
      KeyError: if the task was not queued, processing or abandoned.
    """
    with self._lock:
      is_abandoned = task.identifier in self._tasks_abandoned
      is_processing = task.identifier in self._tasks_processing
      is_queued = task.identifier in self._tasks_queued

      if not is_queued and not is_processing and not is_abandoned:
        raise KeyError('Status of task {0:s} is unknown.'.format(
            task.identifier))

      return is_queued or is_processing or is_abandoned and not task.has_retry

  def CreateRetryTask(self):
    """Creates a task that to retry a previously abandoned task.

    Returns:
      Task: a task that was abandoned but should be retried or None if there are
          no abandoned tasks that should be retried.
    """
    with self._lock:
      abandoned_task = self._GetTaskPendingRetry()
      if not abandoned_task:
        return None

      # The abandoned task is kept in _tasks_abandoned so it can be still
      # identified in CheckTaskToMerge and UpdateTaskAsPendingMerge.

      retry_task = abandoned_task.CreateRetryTask()
      logger.debug('Retrying task {0:s} as {1:s}.'.format(
          abandoned_task.identifier, retry_task.identifier))

      self._tasks_queued[retry_task.identifier] = retry_task
      self._total_number_of_tasks += 1

      self.SampleTaskStatus(retry_task, 'created_retry')

      return retry_task

  # TODO: add support for task types.
  def CreateTask(self, session_identifier):
    """Creates a task.

    Args:
      session_identifier (str): the identifier of the session the task is
          part of.

    Returns:
      Task: task attribute container.
    """
    task = tasks.Task(session_identifier)
    logger.debug('Created task: {0:s}.'.format(task.identifier))

    with self._lock:
      self._tasks_queued[task.identifier] = task
      self._total_number_of_tasks += 1

      self.SampleTaskStatus(task, 'created')

    return task

  def CompleteTask(self, task):
    """Completes a task.

    The task is complete and can be removed from the task manager.

    Args:
      task (Task): task.

    Raises:
      KeyError: if the task was not merging.
    """
    with self._lock:
      if task.identifier not in self._tasks_merging:
        raise KeyError('Task {0:s} was not merging.'.format(task.identifier))

      self.SampleTaskStatus(task, 'completed')

      del self._tasks_merging[task.identifier]

      logger.debug('Completed task {0:s}.'.format(task.identifier))

  def GetFailedTasks(self):
    """Retrieves all failed tasks.

    Failed tasks are tasks that were abandoned and have no retry task once
    the foreman is done processing.

    Returns:
      list[Task]: tasks.
    """
    # TODO: add check to determine foreman is done processing.

    with self._lock:
      return [task for task in self._tasks_abandoned.values()
              if not task.has_retry]

  def GetProcessedTaskByIdentifier(self, task_identifier):
    """Retrieves a task that has been processed.

    Args:
      task_identifier (str): unique identifier of the task.

    Returns:
      Task: a task that has been processed.

    Raises:
      KeyError: if the task was not processing, queued or abandoned.
    """
    with self._lock:
      task = self._tasks_processing.get(task_identifier, None)
      if not task:
        task = self._tasks_queued.get(task_identifier, None)
      if not task:
        task = self._tasks_abandoned.get(task_identifier, None)
      if not task:
        raise KeyError('Status of task {0:s} is unknown'.format(
            task_identifier))

    return task

  def GetStatusInformation(self):
    """Retrieves status information about the tasks.

    Returns:
      TasksStatus: tasks status information.
    """
    status = processing_status.TasksStatus()

    with self._lock:
      status.number_of_abandoned_tasks = len(self._tasks_abandoned)
      status.number_of_queued_tasks = len(self._tasks_queued)
      status.number_of_tasks_pending_merge = (
          len(self._tasks_pending_merge) + len(self._tasks_merging))
      status.number_of_tasks_processing = len(self._tasks_processing)
      status.total_number_of_tasks = self._total_number_of_tasks

    return status

  def GetTaskPendingMerge(self, current_task):
    """Retrieves the first task that is pending merge or has a higher priority.

    This function will check if there is a task with a higher merge priority
    than the current_task being merged. If so, that task with the higher
    priority is returned.

    Args:
      current_task (Task): current task being merged or None if no such task.

    Returns:
      Task: the next task to merge or None if there is no task pending merge or
          with a higher priority.
    """
    next_task = self._tasks_pending_merge.PeekTask()
    if not next_task:
      return None

    if current_task and next_task.merge_priority > current_task.merge_priority:
      return None

    with self._lock:
      next_task = self._tasks_pending_merge.PopTask()

    self._tasks_merging[next_task.identifier] = next_task
    return next_task

  def HasPendingTasks(self):
    """Determines if there are tasks running or in need of retrying.

    Returns:
      bool: True if there are tasks that are active, ready to be merged or
          need to be retried.
    """
    with self._lock:
      self._AbandonInactiveProcessingTasks()

      if self._tasks_processing:
        return True

      # There are no tasks being processed, but we might be
      # waiting for some tasks to be merged.
      if self._HasTasksPendingMerge():
        return True

      # There are no tasks processing or pending merge, but there may
      # still be some waiting to be retried, so we check that.
      if self._HasTasksPendingRetry():
        return True

      # It is possible that a worker has processed a task and the foreman has
      # not been informed about it, since there is no feedback from the worker
      # when it pops a task from the queue.

      # If we believe all the workers are idle for longer than the task
      # inactive time (timeout) abandon all queued tasks. This ensures
      # that processing actually stops when the foreman never gets an
      # update from a worker.

      if self._tasks_queued:
        inactive_time = time.time() - self._TASK_INACTIVE_TIME
        inactive_time = int(inactive_time * definitions.MICROSECONDS_PER_SECOND)

        if self._latest_task_processing_time < inactive_time:
          self._AbandonQueuedTasks()

      if self._tasks_queued:
        return True

      if self._tasks_merging:
        return True

    # There are no tasks pending any work.
    return False

  def RemoveTask(self, task):
    """Removes an abandoned task.

    Args:
      task (Task): task.

    Raises:
      KeyError: if the task was not abandoned or the task was abandoned and
          was not retried.
    """
    with self._lock:
      if task.identifier not in self._tasks_abandoned:
        raise KeyError('Task {0:s} was not abandoned.'.format(task.identifier))

      if not task.has_retry:
        raise KeyError(
            'Will not remove a task {0:s} without retry task.'.format(
                task.identifier))

      del self._tasks_abandoned[task.identifier]

      logger.debug('Removed task {0:s}.'.format(task.identifier))

  def SampleTaskStatus(self, task, status):
    """Takes a sample of the status of the task for profiling.

    Args:
      task (Task): a task.
      status (str): status.
    """
    if self._tasks_profiler:
      self._tasks_profiler.Sample(task, status)

  def StartProfiling(self, configuration, identifier):
    """Starts profiling.

    Args:
      configuration (ProfilingConfiguration): profiling configuration.
      identifier (str): identifier of the profiling session used to create
          the sample filename.
    """
    if not configuration:
      return

    if configuration.HaveProfileTasks():
      self._tasks_profiler = profilers.TasksProfiler(identifier, configuration)
      self._tasks_profiler.Start()

  def StopProfiling(self):
    """Stops profiling."""
    if self._tasks_profiler:
      self._tasks_profiler.Stop()
      self._tasks_profiler = None

  def UpdateTaskAsPendingMerge(self, task):
    """Updates the task manager to reflect the task is ready to be merged.

    Args:
      task (Task): task.

    Raises:
      KeyError: if the task was not queued, processing or abandoned, or
          the task was abandoned and has a retry task.
    """
    with self._lock:
      is_abandoned = task.identifier in self._tasks_abandoned
      is_processing = task.identifier in self._tasks_processing
      is_queued = task.identifier in self._tasks_queued

      if not is_queued and not is_processing and not is_abandoned:
        raise KeyError('Status of task {0:s} is unknown.'.format(
            task.identifier))

      if is_abandoned and task.has_retry:
        raise KeyError('Will not merge a task {0:s} with retry task.'.format(
            task.identifier))

      if is_queued:
        logger.debug('Task {0:s} was queued, now merging.'.format(
            task.identifier))
        del self._tasks_queued[task.identifier]

      if is_processing:
        logger.debug('Task {0:s} was processing, now merging.'.format(
            task.identifier))
        del self._tasks_processing[task.identifier]

      if is_abandoned:
        logger.debug('Task {0:s} was abandoned, now merging.'.format(
            task.identifier))
        del self._tasks_abandoned[task.identifier]

      self._tasks_pending_merge.PushTask(task)

      self.SampleTaskStatus(task, 'pending_merge')

      task.UpdateProcessingTime()
      self._UpdateLatestProcessingTime(task)

  def UpdateTaskAsProcessingByIdentifier(self, task_identifier):
    """Updates the task manager to reflect the task is processing.

    Args:
      task_identifier (str): unique identifier of the task.

    Raises:
      KeyError: if the task is not known to the task manager.
    """
    with self._lock:
      task_processing = self._tasks_processing.get(task_identifier, None)
      if task_processing:
        task_processing.UpdateProcessingTime()
        self._UpdateLatestProcessingTime(task_processing)
        return

      task_queued = self._tasks_queued.get(task_identifier, None)
      if task_queued:
        logger.debug('Task {0:s} was queued, now processing.'.format(
            task_identifier))
        self._tasks_processing[task_identifier] = task_queued
        del self._tasks_queued[task_identifier]

        task_queued.UpdateProcessingTime()
        self._UpdateLatestProcessingTime(task_queued)
        return

      task_abandoned = self._tasks_abandoned.get(task_identifier, None)
      if task_abandoned:
        del self._tasks_abandoned[task_identifier]
        self._tasks_processing[task_identifier] = task_abandoned
        logger.debug('Task {0:s} was abandoned, but now processing.'.format(
            task_identifier))

        task_abandoned.UpdateProcessingTime()
        self._UpdateLatestProcessingTime(task_abandoned)
        return

      if task_identifier in self._tasks_pending_merge:
        # No need to update the processing time, as this task is already
        # finished processing and is just waiting for merge.
        return

    # If we get here, we don't know what state the tasks is in, so raise.
    raise KeyError('Status of task {0:s} is unknown.'.format(task_identifier))
