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


"""Utilities to aid in testing mapreduces."""


import base64
import collections
import logging
import traceback
import os
import re

from google.appengine.ext.mapreduce import main
from google.appengine.ext.mapreduce import model
from google.appengine.ext.webapp import mock_webapp


def decode_task_payload(task):
  """Decodes POST task payload.

  This can only decode POST payload for a normal task. For huge task,
  use model.HugeTask.decode_payload.

  Args:
    task: a dict representing a taskqueue task as documented in taskqueue_stub.

  Returns:
    parameter_name -> parameter_value dict. If multiple parameter values are
    present, then parameter_value will be a list.
  """
  if not task:
    return {}

  body = base64.b64decode(task["body"])

  return model.HugeTask._decode_payload(body)


def execute_task(task, retries=0, handlers_map=None):
  """Execute mapper's executor task.

  This will try to determine the correct mapper handler for the task, will set
  up all mock environment necessary for task execution, and execute the task
  itself.

  This function can be used for functional-style testing of functionality
  depending on mapper framework.

  Args:
    task: a taskqueue task.
    retries: the current retry of this task.
    handlers_map: a dict from url regex to handler.

  Returns:
    the handler instance used for this task.

  Raises:
    Exception: whatever the task raises.
  """

  if not handlers_map:
    handlers_map = main.create_handlers_map()

  url = task["url"]
  handler = None

  for (re_str, handler_class) in handlers_map:
    re_str = "^" + re_str + "($|\\?)"
    if re.match(re_str, url):
      break
  else:
    raise Exception("Can't determine handler for %s" % task)

  request = mock_webapp.MockRequest()
  request.set_url(url)


  version = "mr-test-support-version.1"
  module = "mr-test-support-module"
  default_version_hostname = "mr-test-support.appspot.com"
  host = "%s.%s.%s" % (version.split(".")[0],
                       module,
                       default_version_hostname)

  if "CURRENT_VERSION_ID" not in os.environ:
    request.environ["CURRENT_VERSION_ID"] = version
  if "DEFAULT_VERSION_HOSTNAME" not in os.environ:
    request.environ["DEFAULT_VERSION_HOSTNAME"] = (
        default_version_hostname)
  if "CURRENT_MODULE_ID" not in os.environ:
    request.environ["CURRENT_MODULE_ID"] = module
  if "HTTP_HOST" not in os.environ:
    request.environ["HTTP_HOST"] = host


  for k, v in task.get("headers", []):
    request.headers[k] = v
    environ_key = "HTTP_" + k.replace("-", "_").upper()
    request.environ[environ_key] = v
  request.headers["X-AppEngine-TaskExecutionCount"] = retries
  request.environ["HTTP_X_APPENGINE_TASKNAME"] = (
      task.get("name", "default_task_name"))
  request.environ["HTTP_X_APPENGINE_QUEUENAME"] = (
      task.get("queue_name", "default"))
  request.environ["PATH_INFO"] = request.path

  if task["method"] == "POST":

    request.body = base64.b64decode(task["body"])
    for k, v in decode_task_payload(task).iteritems():
      request.set(k, v)

  response = mock_webapp.MockResponse()
  try:


    handler = handler_class(request, response)
  except TypeError:

    handler = handler_class()
    handler.initialize(request, response)

  saved_os_environ = os.environ
  try:
    os.environ = dict(os.environ)
    os.environ.update(request.environ)
    if task["method"] == "POST":
      handler.post()
    elif task["method"] == "GET":
      handler.get()
    else:
      raise Exception("Unsupported method: %s" % task.method)
  finally:
    os.environ = saved_os_environ

  if handler.response.status != 200:
    raise Exception("Handler failure: %s (%s). \nTask: %s\nHandler: %s" %
                    (handler.response.status,
                     handler.response.status_message,
                     task,
                     handler))
  return handler


def execute_all_tasks(taskqueue, queue="default", handlers_map=None):
  """Run and remove all tasks in the taskqueue.

  Args:
    taskqueue: An instance of taskqueue stub.
    queue: Queue name to run all tasks from.
    hanlders_map: see main.create_handlers_map.

  Returns:
    task_run_counts: a dict from handler class to the number of tasks
      it handled.
  """
  tasks = taskqueue.GetTasks(queue)
  taskqueue.FlushQueue(queue)
  task_run_counts = collections.defaultdict(lambda: 0)
  for task in tasks:
    retries = 0
    while True:
      try:
        handler = execute_task(task, retries, handlers_map=handlers_map)
        task_run_counts[handler.__class__] += 1
        break

      except:
        retries += 1

        if retries > 100:
          logging.debug("Task %s failed for too many times. Giving up.",
                        task["name"])
          raise
        logging.debug(
            "Task %s is being retried for the %s time",
            task["name"],
            retries)

  return task_run_counts


def execute_until_empty(taskqueue, queue="default", handlers_map=None):
  """Execute taskqueue tasks until it becomes empty.

  Args:
    taskqueue: An instance of taskqueue stub.
    queue: Queue name to run all tasks from.
    hanlders_map: see main.create_handlers_map.

  Returns:
    task_run_counts: a dict from handler class to the number of tasks
      it handled.
  """
  task_run_counts = collections.defaultdict(lambda: 0)
  while taskqueue.GetTasks(queue):
    new_counts = execute_all_tasks(taskqueue, queue, handlers_map)
    for handler_cls in new_counts:
      task_run_counts[handler_cls] += new_counts[handler_cls]
  return task_run_counts