"""Build step class for executing a simple shell command."""

from __future__ import absolute_import
import logging
import time

import six
from helpers import retry_utils
from helpers import signal_utils
from slave import base_step

class ShellStep(base_step.BaseStep):
  """Build step class for executing a shell command."""

  class ShellCommand(object):
    """A simple object wrapping a shell command to execute."""

    def __init__(self, command, validator=None, cwd=None, env=None,
                 num_retries=0, setup_for_retry=None, halt_on_failure=False,
                 retry_if=None):
      """Constructor for a ShellCommand object.

      Args:
        command: The command to be run, as a list of strings (required).
        validator: A function that takes 3 arguments: returncode, stdout, stderr
            from the run command and returns a boolean value indicating if the
            command succeeded (True) or failed (False).
            (Default checks if returncode == 0)
        cwd:     The directory from which this command is to be run. This will
                 override the current directory of the step.
        env:     A dictionary mapping environment variables to their values.
        num_retries: The number of times to retry this step if it fails.
            (Defaults to 0)
      retry_if: A function that takes 3 arguments: returncode, stdout, stderr
          from the run command and returns whether to retry the command (True)
          or to skip retrying (False).
        setup_for_retry: A function to run before attempting to retry this
            command. Can be used to do work that may help the command
            succeed upon retry. Only used if num_retries > 0.
        halt_on_failure: Boolean if we should stop running subsequent steps on
            failure. (Defaults to False)
      """
      assert isinstance(command, list)
      assert not cwd or isinstance(cwd, six.string_types)
      assert not env or isinstance(env, dict)

      self.command = command
      self._validator = validator
      self.cwd = cwd
      self.env = env or {}
      self._num_retries = num_retries
      self._retry_if = retry_if
      self._setup_for_retry = setup_for_retry
      self._halt_on_failure = halt_on_failure

    @property
    def num_retries(self):
      """Number of retries to run."""
      return self._num_retries

    @property
    def halt_on_failure(self):
      """Boolean if we should halt on a failure."""
      return self._halt_on_failure

    def exec_retry_setup(self, returncode, stdout, stderr):
      """Run this before attempting to retry this command.

      If command has failed from a recoverable issue, then this can be used
      to recover so that the next retry attempt of |self.command| will succeed.

      Args:
        returncode: returncode of |self.command| that failed and is about to
            be retried.
        stdout: stdout of |self.command| that failed and is about to be retried.
        stderr: stderr of |self.command| that failed and is about to be retried.
      """
      if self._setup_for_retry:
        self._setup_for_retry(returncode, stdout, stderr)

    def validate(self, returncode, stdout, stderr):
      if self._validator:
        return self._validator(returncode, stdout, stderr)
      return returncode == 0

    def retry_if(self, returncode, stdout, stderr):
      """If command fails validator, check if it should be retried."""
      if self._retry_if:
        return self._retry_if(returncode, stdout, stderr)
      return returncode != 0

  def __init__(self, name=None, command=None, validator=None,
               properties_handler=None, force_success=False, num_retries=0,
               retry_if=None, cwd=None, env=None, **kwargs):
    """Creates a ShellStep instance.

    Args:
      name: name of this step
      command: shell command to run, as a list (argv).
      validator: A function that takes 3 arguments: returncode, stdout, stderr
          from the run command and returns a boolean value indicating if the
          command succeeded (True) or failed (False).
          (Default checks if returncode == 0)
      properties_handler: if provided, method for providing properties from the
                          command output.
      force_success: Force this step to return success, even if the shell
          command returns non-zero. (Defaults to False)
      num_retries: The number of times to retry this step if it fails.
          (Defaults to 0)
      retry_if: A function that takes 3 arguments: returncode, stdout, stderr
          from the run command and returns whether to retry the command (True)
          or to skip retrying (False).
      cwd: The directory from which the ShellCommands will be run if cwd is
          not set on the ShellCommand. This will override the current directory
          for the step also.
      env: A dictionary mapping environment variables to their values to be
          used if no env is set on the ShelLCommands.
      **kwargs: Any additional args to pass to BaseStep.
    """
    base_step.BaseStep.__init__(self, name=name, **kwargs)

    self._command = command
    self._validator = validator
    self._properties_handler = properties_handler
    self._force_success = force_success
    self._num_retries = num_retries
    self._retry_if = retry_if
    self._cwd = cwd
    self._env = env

  def _execute_cmd(self, cmd):
    """Executes a ShellCommand.

    Args:
      cmd: A ShellCommand instance.

    Returns:
      A tuple of:
        Boolean if the command was successful, returncode, stdout, stderr
    """
    directory = cmd.cwd or self._cwd or self.directory or None
    environment = cmd.env or self._env or {}
    retry = True
    attempt = 0
    success = True
    while retry:
      returncode, stdout, stderr = self.exec_subprocess(
          cmd.command, cwd=directory, env=environment)
      success = cmd.validate(returncode, stdout, stderr)
      retry = attempt < cmd.num_retries and not success
      if retry and cmd.retry_if(returncode, stdout, stderr):
        delay = retry_utils.retry_delay(attempt)
        attempt += 1
        logging.info(
            'Retrying command in %s seconds. Attempt %s of %s.',
            delay, attempt, cmd.num_retries)
        time.sleep(delay)
        cmd.exec_retry_setup(returncode, stdout, stderr)
    return success, returncode, stdout, stderr

  def run(self):
    """Executes this step by running all commands found in get_commands().

    Returns:
      True iff the return code of every command was 0.
    """
    retry, success = True, True
    returncode, stdout, stderr = 0, '', ''
    attempt = 0

    while retry:
      try:
        success = True  # Reset success for each retry
        for shell_command in self.get_commands():
          cmd_success, returncode, stdout, stderr = self._execute_cmd(
              shell_command
          )

          success = self._force_success or cmd_success
          if shell_command.halt_on_failure and not success:
            break
      except signal_utils.SigtermException:
        # We want to skip retries in case of a Sigterm because this means we
        # got a signal from outside to stop current execution.
        raise
      except Exception:  # pylint: disable=broad-except
        # Absorb all exceptions until there are no retries left.
        if attempt < self._num_retries:
          logging.exception(('Step exception during [%s] ignored because there '
                             'are still retries(%d) remainining.'),
                            self.get_name(),
                            self._num_retries - attempt)
          success = False
        else:
          logging.exception('Step exception during [%s] with no retries left.',
                            self.get_name(),)
          return False

      finally:
        retry = attempt < self._num_retries and not success

      if retry:
        delay = retry_utils.retry_delay(attempt)
        attempt += 1
        logging.info(
            'Retrying [%s] in %s seconds. Attempt %s of %s.',
            self.get_name(), delay, attempt, self._num_retries)
        time.sleep(delay)

    # Only set properties from the final run, otherwise properties from
    # failed runs and successful runs could be mixed together.
    if self._properties_handler:
      new_properties = self._properties_handler(returncode, stdout, stderr)
      for property_name in new_properties:
        if property_name == 'review':
          self.add_review(new_properties[property_name])
        else:
          self.set_build_property(property_name, new_properties[property_name])

    return success

  def get_commands(self):
    """Returns which commands to be run. Can be overridden by a subclass.

    Returns:
      List of ShellCommand objects to execute.
    """
    assert isinstance(self._command, list)
    return [self.ShellCommand(self._command, validator=self._validator,
                              retry_if=self._retry_if)]