# Copyright (C) 2008 The Android Open Source Project
#
# 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.

import errno
import functools
import io
import multiprocessing
import re
import os
import signal
import sys
import subprocess

from color import Coloring
from command import DEFAULT_LOCAL_JOBS, Command, MirrorSafeCommand, WORKER_BATCH_SIZE
from error import ManifestInvalidRevisionError

_CAN_COLOR = [
    'branch',
    'diff',
    'grep',
    'log',
]


class ForallColoring(Coloring):
  def __init__(self, config):
    Coloring.__init__(self, config, 'forall')
    self.project = self.printer('project', attr='bold')


class Forall(Command, MirrorSafeCommand):
  COMMON = False
  helpSummary = "Run a shell command in each project"
  helpUsage = """
%prog [<project>...] -c <command> [<arg>...]
%prog -r str1 [str2] ... -c <command> [<arg>...]
"""
  helpDescription = """
Executes the same shell command in each project.

The -r option allows running the command only on projects matching
regex or wildcard expression.

By default, projects are processed non-interactively in parallel.  If you want
to run interactive commands, make sure to pass --interactive to force --jobs 1.
While the processing order of projects is not guaranteed, the order of project
output is stable.

# Output Formatting

The -p option causes '%prog' to bind pipes to the command's stdin,
stdout and stderr streams, and pipe all output into a continuous
stream that is displayed in a single pager session.  Project headings
are inserted before the output of each command is displayed.  If the
command produces no output in a project, no heading is displayed.

The formatting convention used by -p is very suitable for some
types of searching, e.g. `repo forall -p -c git log -SFoo` will
print all commits that add or remove references to Foo.

The -v option causes '%prog' to display stderr messages if a
command produces output only on stderr.  Normally the -p option
causes command output to be suppressed until the command produces
at least one byte of output on stdout.

# Environment

pwd is the project's working directory.  If the current client is
a mirror client, then pwd is the Git repository.

REPO_PROJECT is set to the unique name of the project.

REPO_PATH is the path relative the the root of the client.

REPO_OUTERPATH is the path of the sub manifest's root relative to the root of
the client.

REPO_INNERPATH is the path relative to the root of the sub manifest.

REPO_REMOTE is the name of the remote system from the manifest.

REPO_LREV is the name of the revision from the manifest, translated
to a local tracking branch.  If you need to pass the manifest
revision to a locally executed git command, use REPO_LREV.

REPO_RREV is the name of the revision from the manifest, exactly
as written in the manifest.

REPO_COUNT is the total number of projects being iterated.

REPO_I is the current (1-based) iteration count. Can be used in
conjunction with REPO_COUNT to add a simple progress indicator to your
command.

REPO__* are any extra environment variables, specified by the
"annotation" element under any project element.  This can be useful
for differentiating trees based on user-specific criteria, or simply
annotating tree details.

shell positional arguments ($1, $2, .., $#) are set to any arguments
following <command>.

Example: to list projects:

  %prog -c 'echo $REPO_PROJECT'

Notice that $REPO_PROJECT is quoted to ensure it is expanded in
the context of running <command> instead of in the calling shell.

Unless -p is used, stdin, stdout, stderr are inherited from the
terminal and are not redirected.

If -e is used, when a command exits unsuccessfully, '%prog' will abort
without iterating through the remaining projects.
"""
  PARALLEL_JOBS = DEFAULT_LOCAL_JOBS

  @staticmethod
  def _cmd_option(option, _opt_str, _value, parser):
    setattr(parser.values, option.dest, list(parser.rargs))
    while parser.rargs:
      del parser.rargs[0]

  def _Options(self, p):
    p.add_option('-r', '--regex',
                 dest='regex', action='store_true',
                 help='execute the command only on projects matching regex or wildcard expression')
    p.add_option('-i', '--inverse-regex',
                 dest='inverse_regex', action='store_true',
                 help='execute the command only on projects not matching regex or '
                      'wildcard expression')
    p.add_option('-g', '--groups',
                 dest='groups',
                 help='execute the command only on projects matching the specified groups')
    p.add_option('-c', '--command',
                 help='command (and arguments) to execute',
                 dest='command',
                 action='callback',
                 callback=self._cmd_option)
    p.add_option('-e', '--abort-on-errors',
                 dest='abort_on_errors', action='store_true',
                 help='abort if a command exits unsuccessfully')
    p.add_option('--ignore-missing', action='store_true',
                 help='silently skip & do not exit non-zero due missing '
                      'checkouts')

    g = p.get_option_group('--quiet')
    g.add_option('-p',
                 dest='project_header', action='store_true',
                 help='show project headers before output')
    p.add_option('--interactive',
                 action='store_true',
                 help='force interactive usage')

  def WantPager(self, opt):
    return opt.project_header and opt.jobs == 1

  def ValidateOptions(self, opt, args):
    if not opt.command:
      self.Usage()

  def Execute(self, opt, args):
    cmd = [opt.command[0]]
    all_trees = not opt.this_manifest_only

    shell = True
    if re.compile(r'^[a-z0-9A-Z_/\.-]+$').match(cmd[0]):
      shell = False

    if shell:
      cmd.append(cmd[0])
    cmd.extend(opt.command[1:])

    # Historically, forall operated interactively, and in serial.  If the user
    # has selected 1 job, then default to interacive mode.
    if opt.jobs == 1:
      opt.interactive = True

    if opt.project_header \
            and not shell \
            and cmd[0] == 'git':
      # If this is a direct git command that can enable colorized
      # output and the user prefers coloring, add --color into the
      # command line because we are going to wrap the command into
      # a pipe and git won't know coloring should activate.
      #
      for cn in cmd[1:]:
        if not cn.startswith('-'):
          break
      else:
        cn = None
      if cn and cn in _CAN_COLOR:
        class ColorCmd(Coloring):
          def __init__(self, config, cmd):
            Coloring.__init__(self, config, cmd)
        if ColorCmd(self.manifest.manifestProject.config, cn).is_on:
          cmd.insert(cmd.index(cn) + 1, '--color')

    mirror = self.manifest.IsMirror
    rc = 0

    smart_sync_manifest_name = "smart_sync_override.xml"
    smart_sync_manifest_path = os.path.join(
        self.manifest.manifestProject.worktree, smart_sync_manifest_name)

    if os.path.isfile(smart_sync_manifest_path):
      self.manifest.Override(smart_sync_manifest_path)

    if opt.regex:
      projects = self.FindProjects(args, all_manifests=all_trees)
    elif opt.inverse_regex:
      projects = self.FindProjects(args, inverse=True, all_manifests=all_trees)
    else:
      projects = self.GetProjects(args, groups=opt.groups, all_manifests=all_trees)

    os.environ['REPO_COUNT'] = str(len(projects))

    try:
      config = self.manifest.manifestProject.config
      with multiprocessing.Pool(opt.jobs, InitWorker) as pool:
        results_it = pool.imap(
            functools.partial(DoWorkWrapper, mirror, opt, cmd, shell, config),
            enumerate(projects),
            chunksize=WORKER_BATCH_SIZE)
        first = True
        for (r, output) in results_it:
          if output:
            if first:
              first = False
            elif opt.project_header:
              print()
            # To simplify the DoWorkWrapper, take care of automatic newlines.
            end = '\n'
            if output[-1] == '\n':
              end = ''
            print(output, end=end)
          rc = rc or r
          if r != 0 and opt.abort_on_errors:
            raise Exception('Aborting due to previous error')
    except (KeyboardInterrupt, WorkerKeyboardInterrupt):
      # Catch KeyboardInterrupt raised inside and outside of workers
      rc = rc or errno.EINTR
    except Exception as e:
      # Catch any other exceptions raised
      print('forall: unhandled error, terminating the pool: %s: %s' %
            (type(e).__name__, e),
            file=sys.stderr)
      rc = rc or getattr(e, 'errno', 1)
    if rc != 0:
      sys.exit(rc)


class WorkerKeyboardInterrupt(Exception):
  """ Keyboard interrupt exception for worker processes. """


def InitWorker():
  signal.signal(signal.SIGINT, signal.SIG_IGN)


def DoWorkWrapper(mirror, opt, cmd, shell, config, args):
  """ A wrapper around the DoWork() method.

  Catch the KeyboardInterrupt exceptions here and re-raise them as a different,
  ``Exception``-based exception to stop it flooding the console with stacktraces
  and making the parent hang indefinitely.

  """
  cnt, project = args
  try:
    return DoWork(project, mirror, opt, cmd, shell, cnt, config)
  except KeyboardInterrupt:
    print('%s: Worker interrupted' % project.name)
    raise WorkerKeyboardInterrupt()


def DoWork(project, mirror, opt, cmd, shell, cnt, config):
  env = os.environ.copy()

  def setenv(name, val):
    if val is None:
      val = ''
    env[name] = val

  setenv('REPO_PROJECT', project.name)
  setenv('REPO_OUTERPATH', project.manifest.path_prefix)
  setenv('REPO_INNERPATH', project.relpath)
  setenv('REPO_PATH', project.RelPath(local=opt.this_manifest_only))
  setenv('REPO_REMOTE', project.remote.name)
  try:
    # If we aren't in a fully synced state and we don't have the ref the manifest
    # wants, then this will fail.  Ignore it for the purposes of this code.
    lrev = '' if mirror else project.GetRevisionId()
  except ManifestInvalidRevisionError:
    lrev = ''
  setenv('REPO_LREV', lrev)
  setenv('REPO_RREV', project.revisionExpr)
  setenv('REPO_UPSTREAM', project.upstream)
  setenv('REPO_DEST_BRANCH', project.dest_branch)
  setenv('REPO_I', str(cnt + 1))
  for annotation in project.annotations:
    setenv("REPO__%s" % (annotation.name), annotation.value)

  if mirror:
    setenv('GIT_DIR', project.gitdir)
    cwd = project.gitdir
  else:
    cwd = project.worktree

  if not os.path.exists(cwd):
    # Allow the user to silently ignore missing checkouts so they can run on
    # partial checkouts (good for infra recovery tools).
    if opt.ignore_missing:
      return (0, '')

    output = ''
    if ((opt.project_header and opt.verbose)
            or not opt.project_header):
      output = 'skipping %s/' % project.RelPath(local=opt.this_manifest_only)
    return (1, output)

  if opt.verbose:
    stderr = subprocess.STDOUT
  else:
    stderr = subprocess.DEVNULL

  stdin = None if opt.interactive else subprocess.DEVNULL

  result = subprocess.run(
      cmd, cwd=cwd, shell=shell, env=env, check=False,
      encoding='utf-8', errors='replace',
      stdin=stdin, stdout=subprocess.PIPE, stderr=stderr)

  output = result.stdout
  if opt.project_header:
    if output:
      buf = io.StringIO()
      out = ForallColoring(config)
      out.redirect(buf)
      if mirror:
        project_header_path = project.name
      else:
        project_header_path = project.RelPath(local=opt.this_manifest_only)
      out.project('project %s/' % project_header_path)
      out.nl()
      buf.write(output)
      output = buf.getvalue()
  return (result.returncode, output)