#!/usr/bin/env python
#
# Copyright 2007 Google Inc.
#
# 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.
#




"""Pipelines for mapreduce library."""

from __future__ import with_statement



import google

from appengine_pipeline.src import pipeline
from appengine_pipeline.src.pipeline import common as pipeline_common
from google.appengine.api import files
from google.appengine.ext.mapreduce import base_handler
from google.appengine.ext.mapreduce import mapper_pipeline
from google.appengine.ext.mapreduce import shuffler




MapperPipeline = mapper_pipeline.MapperPipeline


class MapPipeline(base_handler.PipelineBase):
  """Runs the map stage of MapReduce.

  Iterates over input reader and outputs data into key/value format
  for shuffler consumption.

  Args:
    job_name: mapreduce job name as string.
    mapper_spec: specification of map handler function as string.
    input_reader_spec: input reader specification as string.
    params: mapper and input reader parameters as dict.
    shards: number of shards to start as int.

  Returns:
    list of filenames list sharded by hash code.
  """

  def run(self,
          job_name,
          mapper_spec,
          input_reader_spec,
          params,
          shards=None):
    yield MapperPipeline(
        job_name + "-map",
        mapper_spec,
        input_reader_spec,
        output_writer_spec=
            shuffler.__name__ + "._KeyValueBlobstoreOutputWriter",
        params=params,
        shards=shards)


class ReducePipeline(base_handler.PipelineBase):
  """Runs the reduce stage of MapReduce.

  Merge-reads input files and runs reducer function on them.

  Args:
    job_name: mapreduce job name as string.
    reader_spec: specification of reduce function.
    output_writer_spec: specification of output write to use with reduce
      function.
    parmas: mapper parameters to use as dict.
    filenames: list of filenames to reduce.

  Returns:
    filenames from output writer.
  """

  def run(self,
          job_name,
          reducer_spec,
          output_writer_spec,
          params,
          filenames):
    new_params = dict(params or {})
    new_params.update({
        "files": filenames
        })
    yield mapper_pipeline.MapperPipeline(
        job_name + "-reduce",
        reducer_spec,
        shuffler.__name__ + "._MergingReader",
        output_writer_spec,
        new_params)


class ShufflePipeline(base_handler.PipelineBase):
  """A pipeline to sort multiple key-value files.

  Args:
    filenames: list of file names to sort. Files have to be of records format
      defined by Files API and contain serialized file_service_pb.KeyValue
      protocol messages.

  Returns:
    The list of filenames as string. Resulting files have the same format as
    input and are sorted by key.
  """
  def run(self, shards):
    result = []


    shuffled_shards = [[] for _ in shards[0]]
    for shard in shards:
      for i, filename in enumerate(shard):
        shuffled_shards[i].append(filename)

    for filenames in shuffled_shards:
      sorted_files = yield shuffler.SortPipeline(filenames)
      result.append(sorted_files)
    yield pipeline_common.Append(*result)


class CleanupPipeline(base_handler.PipelineBase):
  """A pipeline to do a cleanup for mapreduce jobs.

  Args:
    temp_files: list of lists of temporary files generated by mapreduce
      job to delete.
  """

  def run(self, temp_files):
    for shard in temp_files:
      for filename in shard:

        for _ in range(10):
          try:
            files.delete(filename)
            break
          except:
            pass


class MapreducePipeline(base_handler.PipelineBase):
  """Pipeline to execute MapReduce jobs.

  Args:
    job_name: job name as string.
    mapper_spec: specification of mapper to use.
    reader_spec: specification of reducer to use.
    input_reader_spec: specification of input reader to read data from.
    output_writer_spec: specification of output writer to save reduce output to.
    mapper_params: parameters to use for mapper phase.
    reducer_params: parameters to use for reduce phase.
    shards: number of shards to use as int.

  Returns:
    filenames from output writer.
  """

  def run(self,
          job_name,
          mapper_spec,
          reducer_spec,
          input_reader_spec,
          output_writer_spec=None,
          mapper_params=None,
          reducer_params=None,
          shards=None):
    map_pipeline = yield MapPipeline(job_name,
                                     mapper_spec,
                                     input_reader_spec,
                                     params=mapper_params,
                                     shards=shards)
    shuffler_pipeline = yield ShufflePipeline(map_pipeline)
    reducer_pipeline = yield ReducePipeline(job_name,
                                            reducer_spec,
                                            output_writer_spec,
                                            reducer_params,
                                            shuffler_pipeline)
    with pipeline.After(reducer_pipeline):
      all_temp_files = yield pipeline_common.Extend(
          map_pipeline, shuffler_pipeline)
      yield CleanupPipeline(all_temp_files)
    yield pipeline_common.Return(reducer_pipeline)
