# Copyright 2013 Google Inc. All Rights Reserved.
# 
# 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.

"""Data pipeline runner."""

import copy
import json
import logging


from src.pipelines import pipelines


class PipelineRunner(object):
  """Generates and runs data pipelines."""

  def Build(self, defn, sink_generator=None):
    """Constructs a pipeline.

    Args:
      defn: The linted pipeline definition as a dictionary of stages.
        Looks like:
        {
          'inputs': [],
          'transforms': [],
          'outputs': [],
          'options': {},
        }
      sink_generator: Required for filling in missing stage output sinks.

    Returns:
      An asynchronous pipeline made up of the defined stages.

    Raises:
      ValueError: If no definition is provided.
    """
    if not defn:
      raise ValueError('No definition provided for pipeline.')

    # Here is where we could select different types but for now there is only
    # the PushPipeline.
    defn = self.Scrub(defn, sink_generator)
    return pipelines.PushPipeline(defn)

  def Scrub(self, defn, sink_generator):
    """Sanitizes the definition.

    Ensures the definition of the pipeline is fully wired together according to
    these rules:
      - Input stages must have at least one sink. If missing from the definition
        and not explicitly declared as None, first try to use the source defined
        on the next stage. If this link is missing, then use the provided
        generator and set its result as the input's sink and the next stage's
        source. No action will be taken if an input stage is missing a source.
      - Transform stages must have both at least one source and sink each. If
        not originally defined, the source will be generated by the preceding
        stage. If the sink is missing, follow the same process as above for
        setting the transform's sink and next stage's source.
      - Output stages must have at least one source. If not originally defined
        and not explicitly declared as None, the source will be generated by the
        preceding stage. No action will be taken if an output stage is missing a
        sink.

    Args:
      defn: The pipeline definition as a list of stages.
      sink_generator: Generates temporary storage links.

    Returns:
      The sanitized definition.

    Raises:
      PipelineError: Invalid definition.
    """
    scrub_defn = copy.deepcopy(defn)
    inputs = scrub_defn.get('inputs', [])
    outputs = scrub_defn.get('outputs', [])
    transforms = scrub_defn.get('transforms', [])

    if not inputs and not outputs:
      raise pipelines.PipelineError('Definition missing input or output stage.')

    filtered_inputs = [s for s in scrub_defn.get('inputs', [])
                       if 'sinks' not in s or s['sinks'] is not None]
    filtered_outputs = [s for s in scrub_defn.get('outputs', [])
                        if 'sources' not in s or s['sources'] is not None]

    if filtered_inputs:
      if transforms:
        # Process all inputs -> first transform
        self._ScrubFanIn(filtered_inputs, transforms[0], sink_generator)
        # Process transform -> transform
        for i in range(len(transforms)-1):
          self._ScrubFanIn([transforms[i]], transforms[i + 1], sink_generator)
        # Process last transform -> all outputs
        if filtered_outputs:
          self._ScrubFanOut(transforms[-1], filtered_outputs, sink_generator)
      elif filtered_outputs:
        if len(filtered_inputs) == 1:
          # Process one input -> all outputs
          self._ScrubFanOut(filtered_inputs[0],
                            filtered_outputs,
                            sink_generator)
        else:
          # Process all inputs -> first output. We don't support N:N right now.
          self._ScrubFanIn(filtered_inputs, filtered_outputs[0], sink_generator)
          if len(filtered_outputs) > 1:
            logging.info('Skipping %d output stages.', len(filtered_outputs)-1)
      else:
        # Still need to scrub the inputs
        for stage in filtered_inputs:
          if 'sinks' not in stage or not stage['sinks']:
            raise pipelines.PipelineError('No sinks found for input stage %s.',
                                          str(stage))
    else:
      if transforms:
        # Process transform -> transform
        for i in range(len(transforms)-1):
          self._ScrubFanIn([transforms[i]], transforms[i + 1], sink_generator)
        # Process last transform -> all outputs
        self._ScrubFanOut(transforms[-1], filtered_outputs, sink_generator)
      else:
        # Still need to scrub the outputs
        for stage in filtered_outputs:
          if 'sources' not in stage or not stage['sources']:
            raise pipelines.PipelineError(
                'No sources found for output stage %s.', str(stage))

    logging.info('scrubbed definition:\n%s',
                 json.dumps(scrub_defn, indent=2, separators=(',', ': ')))
    return scrub_defn

  def _ScrubFanIn(self, input_stages, output_stage, sink_generator):
    """Wires one or more stages' sinks to one stage's sources.

    Args:
      input_stages: A list of stages.
      output_stage: The next stage.
      sink_generator: Generates temporary storage links.

    Raises:
      PipelineError: Invalid definition.
    """
    logging.info('scrubbing stages Fan In:\n%s\n->\n%s\n',
                 json.dumps(input_stages, indent=2, separators=(',', ': ')),
                 json.dumps(output_stage, indent=2, separators=(',', ': ')))
    if 'sources' not in output_stage:
      output_stage['sources'] = []

    # Write any input stage's sink as a source on the output stage,
    # if none exist, generate a temporary link and bind together.
    for stage in input_stages:
      if 'sinks' in stage and stage['sinks']:
        for sink in stage['sinks']:
          if sink not in output_stage['sources']:
            output_stage['sources'].append(sink)
      else:
        link = sink_generator()
        stage['sinks'] = [link]
        output_stage['sources'].append(link)

  def _ScrubFanOut(self, input_stage, output_stages, sink_generator):
    """Wires one stage's sinks to multiple stages' sources.

    Args:
      input_stage: A stages.
      output_stages: The next stages.
      sink_generator: Generates temporary storage links.

    Raises:
      PipelineError: Invalid definition.
    """
    logging.info('scrubbing stages Fan Out:\n%s\n->\n%s\n',
                 json.dumps(input_stage, indent=2, separators=(',', ': ')),
                 json.dumps(output_stages, indent=2, separators=(',', ': ')))
    if 'sinks' not in input_stage:
      input_stage['sinks'] = []

    # Write any output stage's source as a sink on the input stage,
    # if none exist, generate a temporary link and bind together.
    for stage in output_stages:
      if 'sources' in stage and stage['sources']:
        for source in stage['sources']:
          if source not in input_stage['sinks']:
            input_stage['sinks'].append(source)
      else:
        link = sink_generator()
        stage['sources'] = [link]
        input_stage['sinks'].append(link)
