#!/usr/bin/env python3
# Copyright (c) 2018, the R8 project authors. Please see the AUTHORS file
# for details. All rights reserved. Use of this source code is governed by a
# BSD-style license that can be found in the LICENSE file.

# Script for running kotlin based benchmarks

import optparse
import os
import subprocess
import sys
import toolhelper
import utils


BENCHMARK_ROOT = os.path.join(utils.REPO_ROOT, 'third_party', 'benchmarks',
                              'kotlin-benches')

BENCHMARK_PATTERN = '{benchmark}/kotlin/perf/build/libs/perf-1.0-BENCH.jar'
BENCHMARK_MAIN_CLASS = 'com.android.kt.bms.cli.Runner'
ART = os.path.join(utils.TOOLS_DIR, 'linux', 'art', 'bin', 'art')

PROGUARD_CONF = """
# From Android rules
-keepclasseswithmembers public class * {
  public static void main(java.lang.String[]);
}
# Disable obfuscation to only focus on shrinking
-dontobfuscate
# Once we're ready for optimization, we might want to relax access modifiers.
-allowaccessmodification
"""

DEVICE_TEMP='/data/local/temp/bench'


def parse_options():
  result = optparse.OptionParser()
  result.add_option('--api',
                    help='Android api level',
                    default='26',
                    choices=['21', '22', '23', '24', '25', '26'])
  result.add_option('--benchmark',
                    help='The benchmark to run',
                    default='rgx',
                    choices=['rgx', 'deltablue', 'sta', 'empty'])
  result.add_option('--use-device',
                    help='Run the benchmark on an attaced device',
                    default=False, action='store_true')
  return result.parse_args()


def get_jar_for_benchmark(benchmark):
  return os.path.join(BENCHMARK_ROOT,
                      BENCHMARK_PATTERN.format(benchmark=benchmark))

def run_art(dex):
  command = ['bash', ART, '-cp', dex, BENCHMARK_MAIN_CLASS]
  utils.PrintCmd(command)
  benchmark_output = subprocess.check_output(command)
  return get_result(benchmark_output)

def adb(args):
  command = ['adb'] + args
  utils.PrintCmd(command)
  return subprocess.check_output(['adb'] + args)

def get_result(output):
  # There is a lot of debug output, with the actual results being in the line with:
  # RESULTS,KtBench,KtBench,15719
  # structure.
  for result in [s for s in output.splitlines() if s.startswith('RESULTS')]:
    return s.split('RESULTS,KtBench,KtBench,')[1]

def run_art_device(dex):
  adb(['wait-for-device', 'root'])
  device_dst = os.path.join(DEVICE_TEMP, os.path.basename(dex))
  adb(['push', dex, device_dst])
  benchmark_output = adb(['shell', 'dalvikvm', '-cp', device_dst, BENCHMARK_MAIN_CLASS])
  return get_result(benchmark_output)

def Main():
  (options, args) = parse_options()
  with utils.TempDir() as temp:
    dex_path = os.path.join(temp, "classes.jar")
    proguard_conf = os.path.join(temp, 'proguard.conf')
    with open(proguard_conf, 'w') as f:
      f.write(PROGUARD_CONF)
    benchmark_jar = get_jar_for_benchmark(options.benchmark)
    r8_args = [
        '--lib', utils.get_android_jar(26), # Only works with api 26
        '--output', dex_path,
        '--pg-conf', proguard_conf,
        '--min-api', str(options.api),
        benchmark_jar
    ]
    toolhelper.run('r8', r8_args, True)
    if options.use_device:
      result = run_art_device(dex_path)
    else:
      result = run_art(dex_path)
    print('Kotlin_{}(RunTimeRaw): {} ms'.format(options.benchmark, result))

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