#!/usr/bin/python
#
# Copyright 2010 Google Inc. All Rights Reserved.
"""Script to run ChromeOS benchmarks

Inputs:
    chromeos_root
    board
    [chromeos/cpu/<benchname>|chromeos/browser/[pagecycler|sunspider]|chromeos/startup]
    hostname/IP of Chromeos machine

    chromeos/cpu/<benchname>
       - Read run script rules from bench.mk perflab-bin, copy benchmark to
       host, run
       and return results.

    chromeos/startup
       - Re-image host with image in perflab-bin
       - Call run_tests to run startup test, gather results.
       - Restore host back to what it was.

    chromeos/browser/*
       - Call build_chromebrowser to build image with new browser
       - Copy image to perflab-bin

"""

__author__ = 'bjanakiraman@google.com (Bhaskar Janakiraman)'

import optparse
import os
import re
import sys

import image_chromeos
import run_tests
from utils import command_executer
from utils import logger

KNOWN_BENCHMARKS = [
    'chromeos/startup', 'chromeos/browser/pagecycler',
    'chromeos/browser/sunspider', 'chromeos/browser/v8bench',
    'chromeos/cpu/bikjmp'
]

name_map = {
    'pagecycler': 'Page',
    'sunspider': 'SunSpider',
    'v8bench': 'V8Bench',
    'startup': 'BootPerfServer'
}

# Run command template

# Common initializations
cmd_executer = command_executer.GetCommandExecuter()


def Usage(parser, message):
  print 'ERROR: ' + message
  parser.print_help()
  sys.exit(0)


def RunBrowserBenchmark(chromeos_root, board, bench, workdir, machine):
  """Run browser benchmarks.

  Args:
    chromeos_root: ChromeOS src dir
    board: Board being tested
    bench: Name of benchmark (chromeos/browser/*)
    workdir: Directory containing benchmark directory
    machine: name of chromeos machine
  """
  benchname = re.split('/', bench)[2]
  benchdir = '%s/%s' % (workdir, benchname)
  benchname = name_map[benchname]
  retval = run_tests.RunRemoteTests(chromeos_root, machine, board, benchname)
  return retval


def RunStartupBenchmark(chromeos_root, board, bench, workdir, machine):
  """Run browser benchmarks.

  Args:
    chromeos_root: ChromeOS src dir
    board: Board being tested
    bench: Name of benchmark (chromeos/browser/*)
    workdir: Directory containing benchmark directory
    machine: name of chromeos machine
  """
  benchname = 'startup'
  benchdir = '%s/%s' % (workdir, benchname)
  benchname = name_map[benchname]
  retval = run_tests.RunRemoteTests(chromeos_root, machine, board, benchname)
  return retval


def RunCpuBenchmark(chromeos_root, bench, workdir, machine):
  """Run CPU benchmark.

  Args:
    bench: Name of benchmark
    workdir: directory containing benchmark directory
    machine: name of chromeos machine

  Returns:
    status: 0 on success
  """

  benchname = re.split('/', bench)[2]
  benchdir = '%s/%s' % (workdir, benchname)

  # Delete any existing run directories on machine.
  # Since this has exclusive access to the machine,
  # we do not worry about duplicates.
  args = 'rm -rf /tmp/%s' % benchname
  retval = cmd_executer.CrosRunCommand(args,
                                       chromeos_root=chromeos_root,
                                       machine=machine)
  if retval:
    return retval

  # Copy benchmark directory.
  retval = cmd_executer.CopyFiles(benchdir,
                                  '/tmp/' + benchname,
                                  chromeos_root=chromeos_root,
                                  dest_machine=machine,
                                  dest_cros=True)
  if retval:
    return retval

  # Parse bench.mk to extract run flags.

  benchmk_file = open('%s/bench.mk' % benchdir, 'r')
  for line in benchmk_file:
    line.rstrip()
    if re.match('^run_cmd', line):
      line = re.sub('^run_cmd.*\${PERFLAB_PATH}', './out', line)
      line = re.sub('\${PERFLAB_INPUT}', './data', line)
      run_cmd = line
      break

  # Execute on remote machine
  # Capture output and process it.
  sshargs = "'cd /tmp/%s;" % benchname
  sshargs += "time -p %s'" % run_cmd
  cmd_executer.CrosRunCommand(sshargs,
                              chromeos_root=chromeos_root,
                              machine=machine)

  return retval


def Main(argv):
  """Build ChromeOS."""
  # Common initializations

  parser = optparse.OptionParser()
  parser.add_option('-c',
                    '--chromeos_root',
                    dest='chromeos_root',
                    help='Target directory for ChromeOS installation.')
  parser.add_option('-m',
                    '--machine',
                    dest='machine',
                    help='The chromeos host machine.')
  parser.add_option('--workdir',
                    dest='workdir',
                    default='./perflab-bin',
                    help='Work directory for perflab outputs.')
  parser.add_option('--board',
                    dest='board',
                    help='ChromeOS target board, e.g. x86-generic')

  (options, args) = parser.parse_args(argv[1:])

  # validate args
  for arg in args:
    if arg not in KNOWN_BENCHMARKS:
      logger.GetLogger().LogFatal('Bad benchmark %s specified' % arg)

  if options.chromeos_root is None:
    Usage(parser, '--chromeos_root must be set')

  if options.board is None:
    Usage(parser, '--board must be set')

  if options.machine is None:
    Usage(parser, '--machine must be set')

  found_err = 0
  retval = 0
  for arg in args:
    # CPU benchmarks
    comps = re.split('/', arg)
    if re.match('chromeos/cpu', arg):
      benchname = comps[2]
      print 'RUNNING %s' % benchname
      retval = RunCpuBenchmark(options.chromeos_root, arg, options.workdir,
                               options.machine)
      if not found_err:
        found_err = retval
    elif re.match('chromeos/startup', arg):
      benchname = comps[1]
      image_args = [
          os.path.dirname(os.path.abspath(__file__)) + '/image_chromeos.py',
          '--chromeos_root=' + options.chromeos_root,
          '--remote=' + options.machine, '--image=' + options.workdir + '/' +
          benchname + '/chromiumos_image.bin'
      ]
      logger.GetLogger().LogOutput('Reimaging machine %s' % options.machine)
      image_chromeos.Main(image_args)

      logger.GetLogger().LogOutput('Running %s' % arg)
      retval = RunStartupBenchmark(options.chromeos_root, options.board, arg,
                                   options.workdir, options.machine)
      if not found_err:
        found_err = retval
    elif re.match('chromeos/browser', arg):
      benchname = comps[2]
      image_args = [
          os.path.dirname(os.path.abspath(__file__)) + '/image_chromeos.py',
          '--chromeos_root=' + options.chromeos_root,
          '--remote=' + options.machine, '--image=' + options.workdir + '/' +
          benchname + '/chromiumos_image.bin'
      ]
      logger.GetLogger().LogOutput('Reimaging machine %s' % options.machine)
      image_chromeos.Main(image_args)

      logger.GetLogger().LogOutput('Running %s' % arg)
      retval = RunBrowserBenchmark(options.chromeos_root, options.board, arg,
                                   options.workdir, options.machine)
      if not found_err:
        found_err = retval

  return found_err


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