#!/usr/bin/env python3
# Copyright 2019 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""A script to bisect field trials to pin point a culprit for certain behavior.

Chrome runs with many experiments and variations (field trials) that are
randomly selected based on a configuration from a server. They lead to
different code paths and different Chrome behaviors. When a bug is caused by
one of the experiments or variations, it is useful to be able to bisect into
the set and pin-point which one is responsible.

Go to chrome://version/?show-variations-cmd. At the bottom, a few commandline
switches define the current experiments and variations Chrome runs with.

Sample use:

vpython3 bisect_variations.py --input-file="variations_cmd.txt"
    --output-dir=".\out" --browser=canary --url="https://www.youtube.com/"

"variations_cmd.txt" is the command line switches data saved from
chrome://version/?show-variations-cmd.

Sample use for Android:
vpython3 bisect_variations.py --input-file="variations_cmd.txt"
    --output-dir=".\bisect-out" --url="https://www.youtube.com/"
    --browser-path="out/Android/bin/chrome_apk"

Run with --help to get a complete list of options this script runs with.
"""

import logging
import optparse
import os
import shutil
import subprocess
import sys
import tempfile

import split_variations_cmd

_CHROME_PATH_WIN = {
    # The following three paths are relative to %ProgramFiles%
    "stable": r"Google\Chrome\Application\chrome.exe",
    "beta": r"Google\Chrome Beta\Application\chrome.exe",
    "dev": r"Google\Chrome Dev\Application\chrome.exe",
    # The following two paths are relative to %LOCALAPPDATA%
    "canary": r"Google\Chrome SxS\Application\chrome.exe",
    "chromium": r"Chromium\Application\chrome.exe",
}

_CHROME_PATH_MAC = {
    "stable":
    r"/Applications/Google Chrome.app/Contents/MacOS/Google Chrome",
    "beta":
    r"/Applications/Google Chrome Beta.app/Contents/MacOS/Google Chrome Beta",
    "dev":
    r"/Applications/Google Chrome Dev.app/Contents/MacOS/Google Chrome Dev",
    "canary": (r"/Applications/Google Chrome Canary.app/Contents/MacOS/"
               r"Google Chrome Canary"),
}

_CHROME_PATH_LINUX = {
  "stable": r"/usr/bin/google-chrome",
  "beta": r"/usr/bin/google-chrome-beta",
  "dev": r"/usr/bin/google-chrome-unstable",
  "chromium": r"/usr/bin/chromium",
}

# Maximum command length is 32767. Constant below is reduced to leave space
# for executable and chrome arguments.
_MAX_ARGS_LENGTH_WIN = 32000


def _GetSupportedBrowserTypes():
  """Returns the supported browser types on this platform."""
  if sys.platform.startswith('win'):
    return _CHROME_PATH_WIN.keys()
  if sys.platform == 'darwin':
    return _CHROME_PATH_MAC.keys();
  if sys.platform.startswith('linux'):
    return _CHROME_PATH_LINUX.keys();
  raise NotImplementedError('Unsupported platform')


def _LocateBrowser_Win(browser_type):
  """Locates browser executable path based on input browser type.

  Args:
      browser_type: 'stable', 'beta', 'dev', 'canary', or 'chromium'.

  Returns:
      Browser executable path.
  """
  if browser_type in ['stable', 'beta', 'dev']:
    return os.path.join(os.getenv('ProgramFiles'),
                        _CHROME_PATH_WIN[browser_type])
  else:
    assert browser_type in ['canary', 'chromium']
    return os.path.join(os.getenv('LOCALAPPDATA'),
                        _CHROME_PATH_WIN[browser_type])


def _LocateBrowser_Mac(browser_type):
  """Locates browser executable path based on input browser type.

  Args:
      browser_type: A supported browser type on Mac.

  Returns:
      Browser executable path.
  """
  return _CHROME_PATH_MAC[browser_type]


def _LocateBrowser_Linux(browser_type):
  """Locates browser executable path based on input browser type.

  Args:
      browser_type: A supported browser type on Linux.

  Returns:
      Browser executable path.
  """
  return _CHROME_PATH_LINUX[browser_type]


def _LocateBrowser(browser_type):
  """Locates browser executable path based on input browser type.

  Args:
      browser_type: A supported browser types on this platform.

  Returns:
      Browser executable path.
  """
  supported_browser_types = _GetSupportedBrowserTypes()
  if browser_type not in supported_browser_types:
    raise ValueError('Invalid browser type. Supported values are: %s.' %
                         ', '.join(supported_browser_types))
  if sys.platform.startswith('win'):
    return _LocateBrowser_Win(browser_type)
  elif sys.platform == 'darwin':
    return _LocateBrowser_Mac(browser_type)
  elif sys.platform.startswith('linux'):
    return _LocateBrowser_Linux(browser_type)
  else:
    raise NotImplementedError('Unsupported platform')


def _LoadVariations(filename):
  """Reads variations commandline switches from a file.

  Args:
      filename: A file that contains variations commandline switches.

  Returns:
      A list of commandline switches.
  """
  with open(filename, 'r') as f:
    data = f.read().replace('\n', ' ')
  switches = split_variations_cmd.ParseCommandLineSwitchesString(data)
  return ['--%s=%s' % (switch_name, switch_value) for
          switch_name, switch_value in switches.items()]


def _BuildBrowserArgs(user_data_dir, extra_browser_args, variations_args, is_apk):
  """Builds commandline switches browser runs with.

  Args:
      user_data_dir: A path that is used as user data dir.
      extra_browser_args: A list of extra commandline switches browser runs
          with.
      variations_args: A list of commandline switches that defines the
          variations cmd browser runs with.
      is_apk: Whether we're running an APK.

  Returns:
      A list of commandline switches.
  """
  # Make sure each run is fresh, but avoid first run setup steps.
  browser_args = [
      '--no-first-run',
      '--no-default-browser-check',
      '--user-data-dir=%s' % user_data_dir,
      '--disable-field-trial-config',
  ]
  browser_args.extend(extra_browser_args)
  browser_args.extend(variations_args)

  if is_apk:
      return [
        "--args={}".format(" ".join(browser_args))
      ]

  return browser_args


def _RunVariations(browser_path, url, extra_browser_args, variations_args, is_apk):
  """Launches browser with given variations.

  Args:
      browser_path: Browser executable file.
      url: The webpage URL browser goes to after it launches.
      extra_browser_args: A list of extra commandline switches browser runs
          with.
      variations_args: A list of commandline switches that defines the
          variations cmd browser runs with.
      is_apk: Whether we're running an APK.

  Returns:
      A set of (returncode, stdout, stderr) from browser subprocess.
  """
  command = [os.path.abspath(browser_path)]
  if is_apk:
    command.append("run")
  if url:
    command.append(url)
  tempdir = tempfile.mkdtemp(prefix='bisect_variations_tmp')
  command.extend(_BuildBrowserArgs(user_data_dir=tempdir,
                                   extra_browser_args=extra_browser_args,
                                   variations_args=variations_args,
                                   is_apk=is_apk))
  logging.debug(' '.join(command))

  subproc = subprocess.Popen(
      command, bufsize=-1, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  try:
    stdout, stderr = subproc.communicate(timeout=None if not is_apk else 15)
  except subprocess.TimeoutExpired:
    # APK wrapper scripts do not exit (even if the browser's closed on-device),
    # so hold on for a bit before prompting. After the prompt is dismissed,
    # the next iteration of the APK run script will close the browser and
    # reopen it with new command line args.
    return 0, "", ""
  if not is_apk:
    shutil.rmtree(tempdir, True)
  return (subproc.returncode, stdout, stderr)


def _AskCanReproduce(exit_status, stdout, stderr):
  """Asks whether running Chrome with given variations reproduces the issue.

  Args:
      exit_status: Chrome subprocess return code.
      stdout: Chrome subprocess stdout.
      stderr: Chrome subprocess stderr.

  Returns:
      One of ['y', 'n', 'r']:
        'y': yes
        'n': no
        'r': retry
  """
  # Loop until we get a response that we can parse.
  while True:
    response = input('Can we reproduce with given variations file '
                     '[(y)es/(n)o/(r)etry/(s)tdout/(q)uit]: ').lower()
    if response in ('y', 'n', 'r'):
      return response
    if response == 'q':
      sys.exit()
    if response == 's':
      logging.info(stdout)
      logging.info(stderr)


def Bisect(browser_type, url, browser_path, extra_browser_args, variations_file, output_dir):
  """Bisect variations interactively.

  Args:
      browser_type: One of the supported browser type on this platform. See
          --help for the list.
      url: The webpage URL browser launches with.
      browser_path: Location of the compiled output browser.
      extra_browser_args: A list of commandline switches browser runs with.
      variations_file: A file contains variations commandline switches that
          need to be bisected.
      output_dir: A folder where intermediate bisecting data are stored.
  """
  if not browser_path:
    browser_path = _LocateBrowser(browser_type)
  if sys.platform.startswith('win'):
    runs = _EnsureCommandLineLength(variations_file, output_dir)
  else:
    runs = [variations_file]

  # All Android wrapper scripts end with _apk
  is_apk = browser_path.endswith("_apk")

  # Verify that the issue not be reproduced without variations.
  while True:
    exit_status, stdout, stderr = _RunVariations(
        browser_path=browser_path, url=url,
        extra_browser_args=extra_browser_args,
        variations_args=[], is_apk=is_apk)
    answer = _AskCanReproduce(exit_status, stdout, stderr)
    if answer == 'y':
      raise Exception(
          'The issue was reproduced without any variation flags set. Consider'
          ' using tools/bisect-builds.py instead.\n'
          'You might want to try the following command (substitute M100 for a'
          ' good revision):\n'
          'python3 tools/bisect-builds.py -g M100 --verify-range --'
          ' --disable-field-trial-config\n'
          'See https://www.chromium.org/developers/bisect-builds-py/ for more'
          ' details.'
      )
    elif answer == 'n':
      # We are expected to not reproduce the issue without variation flags.
      break
    else:
      assert answer == 'r'

  while runs:
    run = runs[0]
    print('Run Chrome with variations file', run)
    variations_args = _LoadVariations(run)
    exit_status, stdout, stderr = _RunVariations(
        browser_path=browser_path, url=url,
        extra_browser_args=extra_browser_args,
        variations_args=variations_args, is_apk=is_apk)

    answer = _AskCanReproduce(exit_status, stdout, stderr)
    if answer == 'y':
      runs = split_variations_cmd.SplitVariationsCmdFromFile(run, output_dir)
      if len(runs) == 1:
        # Can divide no further.
        print('Bisecting succeeded:', ' '.join(variations_args))
        return
    elif answer == 'n':
      if len(runs) == 1:
        raise ValueError('Bisecting failed: should reproduce but did not: %s' %
                         ' '.join(variations_args))
      runs = runs[1:]
    else:
      assert answer == 'r'


def _EnsureCommandLineLength(filename, output_dir):
  """Splits command-line to ensure it isn't too long for Windows.

  Args:
      filename: A file that contains variations commandline switches.
      output_dir: A folder where intermediate bisecting data are stored.
  Returns:
       List of files containing variations from the input file, split
       such that no file has a command line too long for Windows.
  """
  files_to_process = [filename]

  result = []
  while len(files_to_process) > 0:
    new_files = []
    for f in files_to_process:
      variations_args = ' '.join(_LoadVariations(f))
      if len(variations_args) <= _MAX_ARGS_LENGTH_WIN:
        result.append(f)
      else:
        split = split_variations_cmd.SplitVariationsCmdFromFile(f, output_dir)
        if len(split) == 1:
          raise ValueError('Can not split long argument list %s' %
                           variations_args)
        new_files.extend(split)
    files_to_process = new_files
  return result


def main():
  parser = optparse.OptionParser()
  parser.add_option("--browser",
                    help="select which browser to run. Options include: %s."
                    " By default, stable is selected." %
                        ", ".join(_GetSupportedBrowserTypes()))
  parser.add_option("-v", "--verbose", action="store_true", default=False,
                    help="print out debug information.")
  parser.add_option("--extra-browser-args",
                    help="specify extra command line switches for the browser "
                    "that are separated by spaces (quoted).")
  parser.add_option("--url",
                    help="specify the webpage URL the browser launches with. "
                    "This is optional.")
  parser.add_option("--input-file",
                    help="specify the filename that contains variations cmd "
                    "to bisect. This has to be specified.")
  parser.add_option("--output-dir",
                    help="specify a folder where output files are saved. "
                    "If not specified, it is the folder of the input file.")
  parser.add_option("--browser-path", help="specify location of the browser "
                    "executable or run script. Overrides the default location " \
                    "from --browser")
  options, _ = parser.parse_args()
  if options.verbose:
    logging.basicConfig(level=logging.DEBUG)
  if options.input_file is None:
    raise ValueError('Missing input through --input-file.')
  output_dir = options.output_dir
  if output_dir is None:
    output_dir, _ = os.path.split(options.input_file)
  if not os.path.exists(output_dir):
    os.makedirs(output_dir)
  browser_type = options.browser
  if browser_type is None:
    browser_type = 'stable'
  extra_browser_args = []
  if options.extra_browser_args is not None:
    extra_browser_args = options.extra_browser_args.split()
  Bisect(browser_type=browser_type, url=options.url,
         browser_path=options.browser_path,
         extra_browser_args=extra_browser_args,
         variations_file=options.input_file, output_dir=output_dir)
  return 0


if __name__ == '__main__':
  sys.exit(main())