#!/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.
#


"""Mapreduce execution context.

Mapreduce context provides handler code with information about
current mapreduce execution and organizes utility data flow
from handlers such as counters, log messages, mutation pools.
"""


__all__ = ["get",
           "Pool",
           "Context",
           "COUNTER_MAPPER_CALLS",
           "COUNTER_MAPPER_WALLTIME_MS",
           "DATASTORE_DEADLINE",
           "MAX_ENTITY_COUNT",
          ]

import heapq
import logging
import threading

try:
  from google.appengine.ext import ndb
except ImportError:
  ndb = None
from google.appengine.api import datastore
from google.appengine.ext import db
from google.appengine.runtime import apiproxy_errors


MAX_ENTITY_COUNT = 20


DATASTORE_DEADLINE = 15


COUNTER_MAPPER_CALLS = "mapper-calls"


COUNTER_MAPPER_WALLTIME_MS = "mapper-walltime-ms"


def _normalize_entity(value):
  """Return an entity from an entity or model instance."""
  if ndb is not None and isinstance(value, ndb.Model):
    return None
  if getattr(value, "_populate_internal_entity", None):
    return value._populate_internal_entity()
  return value


def _normalize_key(value):
  """Return a key from an entity, model instance, key, or key string."""
  if ndb is not None and isinstance(value, (ndb.Model, ndb.Key)):
    return None
  if getattr(value, "key", None):
    return value.key()
  elif isinstance(value, basestring):
    return datastore.Key(value)
  else:
    return value


class _ItemList(object):
  """A buffer that holds arbitrary items and auto flushes them when full.

  Callers of this class provides the logic on how to flush.
  This class takes care of the common logic of when to flush and when to retry.

  Properties:
    items: list of objects.
    length: length of item list.
    size: aggregate item size in bytes.
  """

  DEFAULT_RETRIES = 3
  _LARGEST_ITEMS_TO_LOG = 5

  def __init__(self,
               max_entity_count,
               flush_function,
               timeout_retries=DEFAULT_RETRIES,
               repr_function=None):
    """Constructor.

    Args:
      max_entity_count: maximum number of entities before flushing it to db.
      flush_function: a function that can flush the items. The function is
        called with a list of items as the first argument, a dict of options
        as second argument. Currently options can contain {"deadline": int}.
        see self.flush on how the function is called.
      timeout_retries: how many times to retry upon timeouts.
      repr_function: a function that turns an item into meaningful
        representation. For debugging large items.
    """
    self.items = []
    self.__max_entity_count = int(max_entity_count)
    self.__flush_function = flush_function
    self.__repr_function = repr_function
    self.__timeout_retries = int(timeout_retries)

  def __str__(self):
    return "ItemList of with %s items" % len(self.items)

  def append(self, item):
    """Add new item to the list.

    If needed, append will first flush existing items and clear existing items.

    Args:
      item: an item to add to the list.
    """
    if self.should_flush():
      self.flush()
    self.items.append(item)

  def flush(self):
    """Force a flush."""
    if not self.items:
      return

    retry = 0
    options = {"deadline": DATASTORE_DEADLINE}
    while retry <= self.__timeout_retries:
      try:
        self.__flush_function(self.items, options)
        self.clear()
        break
      except db.Timeout, e:
        logging.warning(e)
        logging.warning("Flushing '%s' timed out. Will retry for the %s time.",
                        self, retry)
        retry += 1
        options["deadline"] *= 2
      except apiproxy_errors.RequestTooLargeError:
        self._log_largest_items()
        raise
    else:
      raise

  def _log_largest_items(self):
    if not self.__repr_function:
      logging.error("Got RequestTooLargeError but can't interpret items in "
                    "_ItemList %s.", self)
      return

    sizes = [len(self.__repr_function(i)) for i in self.items]
    largest = heapq.nlargest(self._LARGEST_ITEMS_TO_LOG,
                             zip(sizes, self.items),
                             lambda t: t[0])

    self._largest = [(s, self.__repr_function(i)) for s, i in largest]
    logging.error("Got RequestTooLargeError. Largest items: %r", self._largest)

  def clear(self):
    """Clear item list."""
    self.items = []

  def should_flush(self):
    """Whether to flush before append the next entity.

    Returns:
      True to flush. False other.
    """
    return len(self.items) >= self.__max_entity_count


class Pool(object):
  """Mutation pool accumulates changes to perform them in patch.

  Any Pool subclass should not be public. Instead, Pool should define an
  operation.Operation class and let user uses that. For example, in a map
  function, user can do:

  def map(foo):
    yield OperationOnMyPool(any_argument)

  Since Operation is a callable object, Mapreduce library will invoke
  any Operation object that is yielded with context.Context instance.
  The operation object can then access MyPool from Context.get_pool.
  """

  def flush(self):
    """Flush all changes."""
    raise NotImplementedError()


class _MutationPool(Pool):
  """Mutation pool accumulates datastore changes to perform them in batch.

  Properties:
    puts: _ItemList of entities to put to datastore.
    deletes: _ItemList of keys to delete from datastore.
    ndb_puts: _ItemList of ndb entities to put to datastore.
    ndb_deletes: _ItemList of ndb keys to delete from datastore.
  """

  def __init__(self,
               max_entity_count=MAX_ENTITY_COUNT,
               mapreduce_spec=None):
    """Constructor.

    Args:
      max_entity_count: maximum number of entities before flushing it to db.
      mapreduce_spec: An optional instance of MapperSpec.
    """
    self.max_entity_count = max_entity_count
    params = mapreduce_spec.params if mapreduce_spec is not None else {}
    self.force_writes = bool(params.get("force_ops_writes", False))
    self.puts = _ItemList(max_entity_count,
                          self._flush_puts,
                          repr_function=self._db_repr)
    self.deletes = _ItemList(max_entity_count,
                             self._flush_deletes)
    self.ndb_puts = _ItemList(max_entity_count,
                              self._flush_ndb_puts,
                              repr_function=self._ndb_repr)
    self.ndb_deletes = _ItemList(max_entity_count,
                                 self._flush_ndb_deletes)

  def put(self, entity):
    """Registers entity to put to datastore.

    Args:
      entity: an entity or model instance to put.
    """
    actual_entity = _normalize_entity(entity)
    if actual_entity is None:
      return self.ndb_put(entity)
    self.puts.append(actual_entity)

  def ndb_put(self, entity):
    """Like put(), but for NDB entities."""
    assert ndb is not None and isinstance(entity, ndb.Model)
    self.ndb_puts.append(entity)

  def delete(self, entity):
    """Registers entity to delete from datastore.

    Args:
      entity: an entity, model instance, or key to delete.
    """
    key = _normalize_key(entity)
    if key is None:
      return self.ndb_delete(entity)
    self.deletes.append(key)

  def ndb_delete(self, entity_or_key):
    """Like delete(), but for NDB entities/keys."""
    if ndb is not None and isinstance(entity_or_key, ndb.Model):
      key = entity_or_key.key
    else:
      key = entity_or_key
    self.ndb_deletes.append(key)

  def flush(self):
    """Flush(apply) all changed to datastore."""
    self.puts.flush()
    self.deletes.flush()
    self.ndb_puts.flush()
    self.ndb_deletes.flush()

  @classmethod
  def _db_repr(cls, entity):
    """Converts entity to a readable repr.

    Args:
      entity: datastore.Entity or datastore_types.Key.

    Returns:
      Proto in str.
    """
    return str(entity._ToPb())

  @classmethod
  def _ndb_repr(cls, entity):
    """Converts entity to a readable repr.

    Args:
      entity: ndb.Model

    Returns:
      Proto in str.
    """
    return str(entity._to_pb())

  def _flush_puts(self, items, options):
    """Flush all puts to datastore."""
    datastore.Put(items, config=self._create_config(options))

  def _flush_deletes(self, items, options):
    """Flush all deletes to datastore."""
    datastore.Delete(items, config=self._create_config(options))

  def _flush_ndb_puts(self, items, options):
    """Flush all NDB puts to datastore."""
    assert ndb is not None
    ndb.put_multi(items, config=self._create_config(options))

  def _flush_ndb_deletes(self, items, options):
    """Flush all deletes to datastore."""
    assert ndb is not None
    ndb.delete_multi(items, config=self._create_config(options))

  def _create_config(self, options):
    """Creates datastore Config.

    Returns:
      A datastore_rpc.Configuration instance.
    """
    return datastore.CreateConfig(deadline=options["deadline"],
                                  force_writes=self.force_writes)


class _Counters(Pool):
  """Regulates access to counters.

  Counters Pool is a str to int map. It is saved as part of ShardState so it
  is flushed when ShardState commits to datastore successfully.
  """

  def __init__(self, shard_state):
    """Constructor.

    Args:
      shard_state: current mapreduce shard state as model.ShardState.
    """
    self._shard_state = shard_state

  def increment(self, counter_name, delta=1):
    """Increment counter value.

    Args:
      counter_name: name of the counter as string.
      delta: increment delta as int.
    """
    self._shard_state.counters_map.increment(counter_name, delta)

  def flush(self):
    """Flush unsaved counter values."""
    pass


class Context(object):
  """MapReduce execution context.

  The main purpose of Context is to facilitate IO. User code, input reader,
  and output writer code can plug in pools (see Pool class) to Context to
  batch operations.

  There is a single Context instance associated with each worker thread.
  It can be accessed via context.get(). handlers.MapperWorkerHandler creates
  this instance before any IO code (input_reader, output_writer, user functions)
  is called.

  Each Pool decides how to batch and when to flush.
  Context and all its pools are flushed by the end of a slice.
  Upon error in slice execution, what is flushed is undefined. (See _Counters
  for an exception).

  Properties:
    mapreduce_spec: current mapreduce specification as model.MapreduceSpec.
  """


  _local = threading.local()

  def __init__(self, mapreduce_spec, shard_state, task_retry_count=0):
    """Constructor.

    Args:
      mapreduce_spec: mapreduce specification as model.MapreduceSpec.
      shard_state: an instance of model.ShardState. This has to be the same
        instance as the one MapperWorkerHandler mutates. All mutations are
        flushed to datastore in the end of the slice.
      task_retry_count: how many times this task has been retried.
    """
    self._shard_state = shard_state
    self.mapreduce_spec = mapreduce_spec


    self.task_retry_count = task_retry_count

    if self.mapreduce_spec:
      self.mapreduce_id = self.mapreduce_spec.mapreduce_id
    else:

      self.mapreduce_id = None
    if shard_state:
      self.shard_id = shard_state.get_shard_id()
    else:

      self.shard_id = None


    self._mutation_pool = _MutationPool(mapreduce_spec=mapreduce_spec)
    self._counters = _Counters(shard_state)


    self.counters = self._counters

    self._pools = {}
    self.register_pool("mutation_pool", self._mutation_pool)
    self.register_pool("counters", self.counters)

  def flush(self):
    """Flush all information recorded in context."""
    for pool in self._pools.values():
      pool.flush()

  def register_pool(self, key, pool):
    """Register an arbitrary pool to be flushed together with this context.

    Args:
      key: pool key as string.
      pool: a pool instance.
    """
    self._pools[key] = pool

  def get_pool(self, key):
    """Obtains an instance of registered pool.

    Args:
      key: pool key as string.

    Returns:
      an instance of the pool registered earlier, or None.
    """
    return self._pools.get(key, None)

  @classmethod
  def _set(cls, context):
    """Set current context instance.

    Args:
      context: new context as Context or None.
    """
    cls._local._context_instance = context


def get():
  """Get current context instance.

  Returns:
    current context as Context.
  """
  if not hasattr(Context._local, "_context_instance") :
    return None
  return Context._local._context_instance