"""Classes to assist in merging attribute containers of tasks."""

from plaso.containers import analysis_results
from plaso.containers import artifacts
from plaso.containers import event_sources
from plaso.containers import events
from plaso.containers import reports
from plaso.containers import warnings


class BaseTaskMergeHelper:
  """Interface of heler for merging task related attribute containers.

  Attributes:
    task_identifier (str): identifier of the task that is merged.
  """

  _CONTAINER_TYPES = ()

  def __init__(self, task_storage_reader, task_identifier):
    """Initialize a helper for merging task related attribute containers.

    Args:
      task_storage_reader (StorageReader): task storage reader.
      task_identifier (str): identifier of the task that is merged.
    """
    super().__init__()
    self._container_identifier_mappings = {}
    self._generator = self._GetAttributeContainers(task_storage_reader)
    self._task_storage_reader = task_storage_reader

    self.fully_merged = False
    self.task_identifier = task_identifier

  def _GetAttributeContainers(self, task_storage_reader):
    """Retrieves attribute containers to merge.

    Args:
      task_storage_reader (StorageReader): task storage reader.

    Yields:
      AttributeContainer: attribute container.
    """
    for container_type in self._CONTAINER_TYPES:
      yield from task_storage_reader.GetAttributeContainers(container_type)

    self.fully_merged = True

  def Close(self):
    """Closes the task storage reader."""
    self._task_storage_reader.Close()
    self._task_storage_reader = None

  def GetAttributeContainer(self):
    """Retrieves an attribute container to merge.

    Returns:
      AttributeContainer: attribute container or None if not available.
    """
    try:
      container = next(self._generator)
    except StopIteration:
      container = None

    return container

  def GetAttributeContainerIdentifier(self, lookup_key):
    """Retrieves an attribute container.

    Args:
      lookup_key (str): lookup key that identifies the attribute container.

    Returns:
      AttributeContainerIdentifier: attribute container identifier that maps
          to the lookup key or None if not available.
    """
    return self._container_identifier_mappings.get(lookup_key)

  def SetAttributeContainerIdentifier(self, lookup_key, identifier):
    """Sets an attribute container.

    Args:
      lookup_key (str): lookup key that identifies the attribute container.
      identifier (AttributeContainerIdentifier): attribute container identifier.
    """
    self._container_identifier_mappings[lookup_key] = identifier


class AnalysisTaskMergeHelper(BaseTaskMergeHelper):
  """Assists in merging attribute containers of an analysis task."""

  # Container types produced by the analysis worker processes that need to be
  # merged. Note that some container types reference other container types and
  # therefore container types that are referenced, must be defined before
  # container types that reference them.

  _CONTAINER_TYPES = (
      events.EventTag.CONTAINER_TYPE,
      reports.AnalysisReport.CONTAINER_TYPE,
      warnings.AnalysisWarning.CONTAINER_TYPE,
      analysis_results.BrowserSearchAnalysisResult.CONTAINER_TYPE,
      analysis_results.ChromeExtensionAnalysisResult.CONTAINER_TYPE)


class ExtractionTaskMergeHelper(BaseTaskMergeHelper):
  """Assists in merging attribute containers of an extraction task."""

  # Container types produced by the extraction worker processes that need to be
  # merged. Note that some container types reference other container types and
  # therefore container types that are referenced, must be defined before
  # container types that reference them.

  _CONTAINER_TYPES = (
      event_sources.EventSource.CONTAINER_TYPE,
      events.EventDataStream.CONTAINER_TYPE,
      # The date-less log helper is needed to generate event from the event
      # data by the timeliner and therefore needs to be merged before event
      # data containers.
      events.DateLessLogHelper.CONTAINER_TYPE,
      events.EventData.CONTAINER_TYPE,
      warnings.ExtractionWarning.CONTAINER_TYPE,
      warnings.RecoveryWarning.CONTAINER_TYPE,
      artifacts.WindowsEventLogMessageFileArtifact.CONTAINER_TYPE,
      artifacts.WindowsEventLogMessageStringArtifact.CONTAINER_TYPE,
      artifacts.WindowsWevtTemplateEvent.CONTAINER_TYPE)
