#!/usr/bin/python
#
# Copyright 2011 Google Inc. All Rights Reserved.

"""Script to build chrome with FDO and compare performance against no FDO."""

import getpass
import optparse
import os
import sys

import image_chromeos
import setup_chromeos
from utils import command_executer
from utils import misc
from utils import logger


class Patcher(object):
  def __init__(self, dir_to_patch, patch_file):
    self._dir_to_patch = dir_to_patch
    self._patch_file = patch_file
    self._base_patch_command = "patch -p0 %%s < %s" % patch_file
    self._ce = command_executer.GetCommandExecuter()

  def _RunPatchCommand(self, args):
    patch_command = self._base_patch_command % args
    command = ("cd %s && %s" % (self._dir_to_patch,
                                patch_command))
    return self._ce.RunCommand(command)

  def _ApplyPatch(self, args):
    full_args = "%s --dry-run" % args
    ret = self._RunPatchCommand(full_args)
    if ret:
      raise Exception("Patch dry run failed!")
    ret = self._RunPatchCommand(args)
    if ret:
      raise Exception("Patch application failed!")

  def __enter__(self):
    self._ApplyPatch("")

  def __exit__(self, type, value, traceback):
    self._ApplyPatch("-R")


class FDOComparator(object):
  def __init__(self, board, remotes, ebuild_version, plus_pgo, minus_pgo,
               update_pgo, chromeos_root):
    self._board = board
    self._remotes = remotes
    self._ebuild_version = ebuild_version
    self._remote = remotes.split(",")[0]
    self._chromeos_root = chromeos_root
    self._profile_dir = "profile_dir"
    self._profile_path = os.path.join(self._chromeos_root,
                                      "src",
                                      "scripts",
                                      os.path.basename(self._profile_dir))
    self._plus_pgo = plus_pgo
    self._minus_pgo = minus_pgo
    self._update_pgo = update_pgo

    self._ce = command_executer.GetCommandExecuter()
    self._l = logger.GetLogger()

  def _CheckoutChromeOS(self):
    if not os.path.exists(self._chromeos_root):
      setup_chromeos_args = [setup_chromeos.__file__,
                             "--dir=%s" % self._chromeos_root,
                             "--minilayout"]
      setup_chromeos.Main(setup_chromeos_args)

  def _BuildChromeOSUsingBinaries(self):
    image_dir = misc.GetImageDir(self._chromeos_root, self._board)
    command = "equery-%s l chromeos" % self._board
    ret = self._ce.ChrootRunCommand(self._chromeos_root, command)
    if ret:
      command = misc.GetSetupBoardCommand(self._board,
                                           usepkg=True)
      ret = self._ce.ChrootRunCommand(self._chromeos_root,
                                      command)
      if ret:
        raise Exception("Couldn't run setup_board!")
      command = misc.GetBuildPackagesCommand(self._board,
                                              True)
      ret = self._ce.ChrootRunCommand(self._chromeos_root,
                                    command)
      if ret:
        raise Exception("Couldn't run build_packages!")

  def _ReportMismatches(self, build_log):
    mismatch_signature = "-Wcoverage-mismatch"
    mismatches = build_log.count(mismatch_signature)
    self._l.LogOutput("Total mismatches: %s" % mismatches)
    stale_files = set([])
    for line in build_log.splitlines():
      if mismatch_signature in line:
        filename = line.split(":")[0]
        stale_files.add(filename)
    self._l.LogOutput("Total stale files: %s" % len(stale_files))

  def _BuildChromeAndImage(self, ebuild_version="", env_dict={}, cflags="",
                           cxxflags="", ldflags="", label="",
                           build_image_args=""):
    env_string = misc.GetEnvStringFromDict(env_dict)
    if not label:
      label = " ".join([env_string,
                        cflags,
                        cxxflags,
                        ldflags,
                        ebuild_version])
      label = label.strip()
      label = misc.GetFilenameFromString(label)
    if not misc.DoesLabelExist(self._chromeos_root, self._board, label):
      build_chrome_browser_args = ["--clean",
                                   "--chromeos_root=%s" % self._chromeos_root,
                                   "--board=%s" % self._board,
                                   "--env=%r" % env_string,
                                   "--cflags=%r" % cflags,
                                   "--cxxflags=%r" % cxxflags,
                                   "--ldflags=%r" % ldflags,
                                   "--ebuild_version=%s" % ebuild_version,
                                   "--build_image_args=%s" % build_image_args]

      build_chrome_browser = os.path.join(os.path.dirname(__file__),
                                          "..",
                                          "build_chrome_browser.py")
      command = "python %s %s" % (build_chrome_browser,
                                  " ".join(build_chrome_browser_args))
      ret, out, err = self._ce.RunCommand(command,
                                          return_output=True)
      if "-fprofile-use" in cxxflags:
        self._ReportMismatches(out)

      if ret:
        raise Exception("Couldn't build chrome browser!")
      misc.LabelLatestImage(self._chromeos_root, self._board, label)
    return label

  def _TestLabels(self, labels):
    experiment_file = "pgo_experiment.txt"
    experiment_header = """
    board: %s
    remote: %s
    """ % (self._board, self._remotes)
    experiment_tests = """
    benchmark: desktopui_PyAutoPerfTests {
      iterations: 1
    }
    """
    with open(experiment_file, "w") as f:
      print >>f, experiment_header
      print >>f, experiment_tests
      for label in labels:
        # TODO(asharif): Fix crosperf so it accepts labels with symbols
        crosperf_label = label
        crosperf_label = crosperf_label.replace("-", "minus")
        crosperf_label = crosperf_label.replace("+", "plus")
        experiment_image = """
        %s {
          chromeos_image: %s
        }
        """ % (crosperf_label,
               os.path.join(misc.GetImageDir(self._chromeos_root, self._board),
                            label,
                            "chromiumos_test_image.bin"))
        print >>f, experiment_image
    crosperf = os.path.join(os.path.dirname(__file__),
                            "..",
                            "crosperf",
                            "crosperf")
    command = "%s %s" % (crosperf, experiment_file)
    ret = self._ce.RunCommand(command)
    if ret:
      raise Exception("Couldn't run crosperf!")

  def _ImageRemote(self, label):
    image_path = os.path.join(misc.GetImageDir(self._chromeos_root,
                                                self._board),
                              label,
                              "chromiumos_test_image.bin")
    image_chromeos_args = [image_chromeos.__file__,
                           "--chromeos_root=%s" % self._chromeos_root,
                           "--image=%s" % image_path,
                           "--remote=%s" % self._remote,
                           "--board=%s" % self._board]
    image_chromeos.Main(image_chromeos_args)

  def _ProfileRemote(self):
    profile_cycler = os.path.join(os.path.dirname(__file__),
                                  "profile_cycler.py")
    profile_cycler_args = ["--chromeos_root=%s" % self._chromeos_root,
                           "--cycler=all",
                           "--board=%s" % self._board,
                           "--profile_dir=%s" % self._profile_path,
                           "--remote=%s" % self._remote]
    command = "python %s %s" % (profile_cycler, " ".join(profile_cycler_args))
    ret = self._ce.RunCommand(command)
    if ret:
      raise Exception("Couldn't profile cycler!")

  def _BuildGenerateImage(self):
    # TODO(asharif): add cflags as well.
    labels_list = ["fprofile-generate", self._ebuild_version]
    label = "_".join(labels_list)
    generate_label = self._BuildChromeAndImage(
        env_dict={"USE": "chrome_internal -pgo pgo_generate"},
        label=label,
        ebuild_version=self._ebuild_version,
        build_image_args="--rootfs_boost_size=400")
    return generate_label

  def _BuildUseImage(self):
    ctarget = misc.GetCtargetFromBoard(self._board, self._chromeos_root)
    chroot_profile_dir = os.path.join("/home/%s/trunk" % getpass.getuser(),
                                      "src",
                                      "scripts",
                                      self._profile_dir,
                                      ctarget)
    cflags = ("-fprofile-use "
              "-fprofile-correction "
              "-Wno-error "
              "-fdump-tree-optimized-blocks-lineno "
              "-fdump-ipa-profile-blocks-lineno "
              "-fno-vpt "
              "-fprofile-dir=%s" %
              chroot_profile_dir)
    labels_list = ["updated_pgo", self._ebuild_version]
    label = "_".join(labels_list)
    pgo_use_label = self._BuildChromeAndImage(
        env_dict={"USE": "chrome_internal -pgo"},
        cflags=cflags,
        cxxflags=cflags,
        ldflags=cflags,
        label=label,
        ebuild_version=self._ebuild_version)
    return pgo_use_label

  def DoAll(self):
    self._CheckoutChromeOS()
    self._BuildChromeOSUsingBinaries()
    labels = []

    if self._minus_pgo:
      minus_pgo = self._BuildChromeAndImage(env_dict={"USE": "chrome_internal -pgo"},
                                            ebuild_version=self._ebuild_version)
      labels.append(minus_pgo)
    if self._plus_pgo:
      plus_pgo = self._BuildChromeAndImage(env_dict={"USE": "chrome_internal pgo"},
                                           ebuild_version=self._ebuild_version)
      labels.append(plus_pgo)

    if self._update_pgo:
      if not os.path.exists(self._profile_path):
        # Build Chrome with -fprofile-generate
        generate_label = self._BuildGenerateImage()
        # Image to the remote box.
        self._ImageRemote(generate_label)
        # Profile it using all page cyclers.
        self._ProfileRemote()

      # Use the profile directory to rebuild it.
      updated_pgo_label = self._BuildUseImage()
      labels.append(updated_pgo_label)

    # Run crosperf on all images now.
    self._TestLabels(labels)
    return 0


def Main(argv):
  """The main function."""
  # Common initializations
###  command_executer.InitCommandExecuter(True)
  command_executer.InitCommandExecuter()
  parser = optparse.OptionParser()
  parser.add_option("--remote",
                    dest="remote",
                    help="Remote machines to run tests on.")
  parser.add_option("--board",
                    dest="board",
                    default="x86-zgb",
                    help="The target board.")
  parser.add_option("--ebuild_version",
                    dest="ebuild_version",
                    default="",
                    help="The Chrome ebuild version to use.")
  parser.add_option("--plus_pgo",
                    dest="plus_pgo",
                    action="store_true",
                    default=False,
                    help="Build USE=+pgo.")
  parser.add_option("--minus_pgo",
                    dest="minus_pgo",
                    action="store_true",
                    default=False,
                    help="Build USE=-pgo.")
  parser.add_option("--update_pgo",
                    dest="update_pgo",
                    action="store_true",
                    default=False,
                    help="Update pgo and build Chrome with the update.")
  parser.add_option("--chromeos_root",
                    dest="chromeos_root",
                    default=False,
                    help="The chromeos root directory")
  options, _ = parser.parse_args(argv)
  if not options.board:
    print "Please give a board."
    return 1
  if not options.remote:
    print "Please give at least one remote machine."
    return 1
  if not options.chromeos_root:
    print "Please provide the chromeos root directory."
    return 1
  if not any((options.minus_pgo, options.plus_pgo, options.update_pgo)):
    print "Please provide at least one build option."
    return 1
  fc = FDOComparator(options.board,
                     options.remote,
                     options.ebuild_version,
                     options.plus_pgo,
                     options.minus_pgo,
                     options.update_pgo,
                     os.path.expanduser(options.chromeos_root))
  return fc.DoAll()


if __name__ == "__main__":
  retval = Main(sys.argv)
  sys.exit(retval)