# Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

from __future__ import print_function

import os
import time
import shlex

from utils import command_executer
import test_flag

TEST_THAT_PATH = '/usr/bin/test_that'
CHROME_MOUNT_DIR = '/tmp/chrome_root'


def GetProfilerArgs(profiler_args):
  # Remove "--" from in front of profiler args.
  args_list = shlex.split(profiler_args)
  new_list = []
  for arg in args_list:
    if arg[0:2] == '--':
      arg = arg[2:]
    new_list.append(arg)
  args_list = new_list

  # Remove "perf_options=" from middle of profiler args.
  new_list = []
  for arg in args_list:
    idx = arg.find('perf_options=')
    if idx != -1:
      prefix = arg[0:idx]
      suffix = arg[idx + len('perf_options=') + 1:-1]
      new_arg = prefix + "'" + suffix + "'"
      new_list.append(new_arg)
    else:
      new_list.append(arg)
  args_list = new_list

  return ' '.join(args_list)


class SuiteRunner(object):
  """This defines the interface from crosperf to test script."""

  def __init__(self,
               logger_to_use=None,
               log_level='verbose',
               cmd_exec=None,
               cmd_term=None):
    self._logger = logger_to_use
    self.log_level = log_level
    self._ce = cmd_exec or command_executer.GetCommandExecuter(
        self._logger,
        log_level=self.log_level)
    self._ct = cmd_term or command_executer.CommandTerminator()

  def Run(self, machine, label, benchmark, test_args, profiler_args):
    for i in range(0, benchmark.retries + 1):
      self.PinGovernorExecutionFrequencies(machine, label.chromeos_root)
      if benchmark.suite == 'telemetry':
        ret_tup = self.Telemetry_Run(machine, label, benchmark, profiler_args)
      elif benchmark.suite == 'telemetry_Crosperf':
        ret_tup = self.Telemetry_Crosperf_Run(machine, label, benchmark,
                                              test_args, profiler_args)
      else:
        ret_tup = self.Test_That_Run(machine, label, benchmark, test_args,
                                     profiler_args)
      if ret_tup[0] != 0:
        self._logger.LogOutput('benchmark %s failed. Retries left: %s' %
                               (benchmark.name, benchmark.retries - i))
      elif i > 0:
        self._logger.LogOutput('benchmark %s succeded after %s retries' %
                               (benchmark.name, i))
        break
      else:
        self._logger.LogOutput('benchmark %s succeded on first try' %
                               benchmark.name)
        break
    return ret_tup

  def GetHighestStaticFrequency(self, machine_name, chromeos_root):
    """Gets the highest static frequency for the specified machine."""
    get_avail_freqs = ('cd /sys/devices/system/cpu/cpu0/cpufreq/; '
                       'if [[ -e scaling_available_frequencies ]]; then '
                       '  cat scaling_available_frequencies; '
                       'else '
                       '  cat scaling_max_freq ; '
                       'fi')
    ret, freqs_str, _ = self._ce.CrosRunCommandWOutput(
        get_avail_freqs,
        machine=machine_name,
        chromeos_root=chromeos_root)
    self._logger.LogFatalIf(ret, 'Could not get available frequencies '
                            'from machine: %s' % machine_name)
    freqs = freqs_str.split()
    # We need to make sure that the frequencies are sorted in decreasing
    # order
    freqs.sort(key=int, reverse=True)

    ## When there is no scaling_available_frequencies file,
    ## we have only 1 choice.
    if len(freqs) == 1:
      return freqs[0]
    # The dynamic frequency ends with a "1000". So, ignore it if found.
    if freqs[0].endswith('1000'):
      return freqs[1]
    else:
      return freqs[0]

  def PinGovernorExecutionFrequencies(self, machine_name, chromeos_root):
    """Set min and max frequencies to max static frequency."""
    highest_freq = self.GetHighestStaticFrequency(machine_name, chromeos_root)
    BASH_FOR = 'for f in {list}; do {body}; done'
    CPUFREQ_DIRS = '/sys/devices/system/cpu/cpu*/cpufreq/'
    change_max_freq = BASH_FOR.format(list=CPUFREQ_DIRS + 'scaling_max_freq',
                                      body='echo %s > $f' % highest_freq)
    change_min_freq = BASH_FOR.format(list=CPUFREQ_DIRS + 'scaling_min_freq',
                                      body='echo %s > $f' % highest_freq)
    change_perf_gov = BASH_FOR.format(list=CPUFREQ_DIRS + 'scaling_governor',
                                      body='echo performance > $f')
    if self.log_level == 'average':
      self._logger.LogOutput('Pinning governor execution frequencies for %s' %
                             machine_name)
    ret = self._ce.CrosRunCommand(' && '.join((
        'set -e ', change_max_freq, change_min_freq, change_perf_gov)),
                                  machine=machine_name,
                                  chromeos_root=chromeos_root)
    self._logger.LogFatalIf(ret, 'Could not pin frequencies on machine: %s' %
                            machine_name)

  def RebootMachine(self, machine_name, chromeos_root):
    command = 'reboot && exit'
    self._ce.CrosRunCommand(command,
                            machine=machine_name,
                            chromeos_root=chromeos_root)
    time.sleep(60)
    # Whenever we reboot the machine, we need to restore the governor settings.
    self.PinGovernorExecutionFrequencies(machine_name, chromeos_root)

  def Test_That_Run(self, machine, label, benchmark, test_args, profiler_args):
    """Run the test_that test.."""
    options = ''
    if label.board:
      options += ' --board=%s' % label.board
    if test_args:
      options += ' %s' % test_args
    if profiler_args:
      self._logger.LogFatal('test_that does not support profiler.')
    command = 'rm -rf /usr/local/autotest/results/*'
    self._ce.CrosRunCommand(command,
                            machine=machine,
                            chromeos_root=label.chromeos_root)

    # We do this because some tests leave the machine in weird states.
    # Rebooting between iterations has proven to help with this.
    self.RebootMachine(machine, label.chromeos_root)

    command = (
        ('%s --autotest_dir ~/trunk/src/third_party/autotest/files --fast '
         '%s %s %s') % (TEST_THAT_PATH, options, machine, benchmark.test_name))
    if self.log_level != 'verbose':
      self._logger.LogOutput('Running test.')
      self._logger.LogOutput('CMD: %s' % command)
    # Use --no-ns-pid so that cros_sdk does not create a different
    # process namespace and we can kill process created easily by
    # their process group.
    return self._ce.ChrootRunCommandWOutput(label.chromeos_root,
                                            command,
                                            command_terminator=self._ct,
                                            cros_sdk_options='--no-ns-pid')

  def RemoveTelemetryTempFile(self, machine, chromeos_root):
    filename = 'telemetry@%s' % machine
    fullname = os.path.join(chromeos_root, 'chroot', 'tmp', filename)
    if os.path.exists(fullname):
      os.remove(fullname)

  def Telemetry_Crosperf_Run(self, machine, label, benchmark, test_args,
                             profiler_args):
    if not os.path.isdir(label.chrome_src):
      self._logger.LogFatal('Cannot find chrome src dir to'
                            ' run telemetry: %s' % label.chrome_src)

    # Check for and remove temporary file that may have been left by
    # previous telemetry runs (and which might prevent this run from
    # working).
    self.RemoveTelemetryTempFile(machine, label.chromeos_root)

    # For telemetry runs, we can use the autotest copy from the source
    # location. No need to have one under /build/<board>.
    autotest_dir_arg = '--autotest_dir ~/trunk/src/third_party/autotest/files'

    profiler_args = GetProfilerArgs(profiler_args)
    fast_arg = ''
    if not profiler_args:
      # --fast works unless we are doing profiling (autotest limitation).
      # --fast avoids unnecessary copies of syslogs.
      fast_arg = '--fast'
    args_string = ''
    if test_args:
      # Strip double quotes off args (so we can wrap them in single
      # quotes, to pass through to Telemetry).
      if test_args[0] == '"' and test_args[-1] == '"':
        test_args = test_args[1:-1]
      args_string = "test_args='%s'" % test_args

    cmd = ('{} {} {} --board={} --args="{} run_local={} test={} '
           '{}" {} telemetry_Crosperf'.format(
               TEST_THAT_PATH, autotest_dir_arg, fast_arg, label.board,
               args_string, benchmark.run_local, benchmark.test_name,
               profiler_args, machine))

    # Use --no-ns-pid so that cros_sdk does not create a different
    # process namespace and we can kill process created easily by their
    # process group.
    chrome_root_options = ('--no-ns-pid '
                           '--chrome_root={} --chrome_root_mount={} '
                           "FEATURES=\"-usersandbox\" "
                           'CHROME_ROOT={}'.format(label.chrome_src,
                                                   CHROME_MOUNT_DIR,
                                                   CHROME_MOUNT_DIR))
    if self.log_level != 'verbose':
      self._logger.LogOutput('Running test.')
      self._logger.LogOutput('CMD: %s' % cmd)
    return self._ce.ChrootRunCommandWOutput(
        label.chromeos_root,
        cmd,
        command_terminator=self._ct,
        cros_sdk_options=chrome_root_options)

  def Telemetry_Run(self, machine, label, benchmark, profiler_args):
    telemetry_run_path = ''
    if not os.path.isdir(label.chrome_src):
      self._logger.LogFatal('Cannot find chrome src dir to' ' run telemetry.')
    else:
      telemetry_run_path = os.path.join(label.chrome_src, 'src/tools/perf')
      if not os.path.exists(telemetry_run_path):
        self._logger.LogFatal('Cannot find %s directory.' % telemetry_run_path)

    if profiler_args:
      self._logger.LogFatal('Telemetry does not support the perf profiler.')

    # Check for and remove temporary file that may have been left by
    # previous telemetry runs (and which might prevent this run from
    # working).
    if not test_flag.GetTestMode():
      self.RemoveTelemetryTempFile(machine, label.chromeos_root)

    rsa_key = os.path.join(
        label.chromeos_root,
        'src/scripts/mod_for_test_scripts/ssh_keys/testing_rsa')

    cmd = ('cd {0} && '
           './run_measurement '
           '--browser=cros-chrome '
           '--output-format=csv '
           '--remote={1} '
           '--identity {2} '
           '{3} {4}'.format(telemetry_run_path, machine, rsa_key,
                            benchmark.test_name, benchmark.test_args))
    if self.log_level != 'verbose':
      self._logger.LogOutput('Running test.')
      self._logger.LogOutput('CMD: %s' % cmd)
    return self._ce.RunCommandWOutput(cmd, print_to_console=False)

  def Terminate(self):
    self._ct.Terminate()


class MockSuiteRunner(object):

  def __init__(self):
    self._true = True

  def Run(self, *_args):
    if self._true:
      return [0, '', '']
    else:
      return [0, '', '']