#!/usr/bin/env vpython3
# Copyright 2024 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Runs benchmarks and generates an orderfile, similar to generate_profile.py.

Example:
Build trichrome_chrome_64_32_bundle and install it on device.

Run this script with:
$ tools/cygprofile/generate_orderfile.py -C out/orderfile-arm64 \
    --android-browser android-trichrome-chrome-64-32-bundle \
    --target-arch arm64

The orderfiles should be located in out/orderfile-arm64/orderfiles.
"""

import argparse
import logging
import pathlib
import sys

import android_profile_tool
import orderfile_shared

_SRC_PATH = pathlib.Path(__file__).resolve().parents[2]
sys.path.append(str(_SRC_PATH / 'third_party/catapult/devil'))
from devil.android import device_utils


def _GetOrderfilesDir(options) -> pathlib.Path:
  if options.isolated_script_test_output:
    orderfiles_dir = options.isolated_script_test_output.parent / 'orderfiles'
  else:
    orderfiles_dir = options.out_dir / 'orderfiles'
  orderfiles_dir.mkdir(exist_ok=True)
  return orderfiles_dir


def _GetOrderfileFilename(options):
  """Gets the path to the architecture-specific orderfile."""
  arch = options.arch
  orderfiles_dir = _GetOrderfilesDir(options)
  return str(orderfiles_dir / f'orderfile.{arch}.out')


def _GetUnpatchedOrderfileFilename(options):
  """Gets the path to the architecture-specific unpatched orderfile."""
  arch = options.arch
  orderfiles_dir = _GetOrderfilesDir(options)
  return str(orderfiles_dir / f'unpatched_orderfile.{arch}')


def CreateArgumentParser():
  """Creates and returns the argument parser."""
  parser = argparse.ArgumentParser()
  orderfile_shared.AddCommonArguments(parser)

  # Essential arguments for profiling and processing:
  parser.add_argument(
    '--android-browser',
    required=True,
    help='Browser string to pass to run_benchmark.',
  )
  parser.add_argument(
    '-C',
    '--out-dir',
    type=pathlib.Path,
    required=True,
    help='Path to the output directory (e.g. out/Release).',
  )
  # The following two are bot-specific args.
  parser.add_argument(
    '--isolated-script-test-output',
    type=pathlib.Path,
    help='Output.json file that the script can write to.',
  )
  parser.add_argument(
    '--isolated-script-test-perf-output',
    help='Deprecated and ignored, but bots pass it.',
  )

  return parser


def GenerateOrderfile(options, device):
  """Generates an orderfile for a given device."""
  host_profile_root = options.out_dir / 'profile_data'
  profiler = android_profile_tool.AndroidProfileTool(
    str(host_profile_root),
    device,
    debug=options.streamline_for_debugging,
    verbosity=options.verbosity,
  )

  lib_chrome_so = orderfile_shared.GetLibchromeSoPath(
    options.out_dir, options.arch, options.profile_webview
  )
  try:
    if options.profile_webview:
      if options.arch == 'arm64':
        webview_target = 'system_webview_64_32_apk'
        apk_name = 'SystemWebView6432.apk'
      else:
        webview_target = 'system_webview_apk'
        apk_name = 'SystemWebView.apk'
      webview_installer_path = str(options.out_dir / 'bin' / webview_target)
      apk_or_browser = str(options.out_dir / 'apks' / apk_name)
    else:
      apk_or_browser = options.android_browser
      webview_installer_path = None
    files = orderfile_shared.CollectProfiles(
      profiler,
      options.profile_webview,
      options.arch,
      apk_or_browser,
      str(options.out_dir),
      webview_installer_path,
    )
    ordered_symbols, _ = orderfile_shared.ProcessProfiles(files, lib_chrome_so)
    with open(_GetUnpatchedOrderfileFilename(options), 'w') as orderfile:
      orderfile.write('\n'.join(ordered_symbols))
  finally:
    if not options.save_profile_data:
      profiler.Cleanup()
    logging.getLogger().setLevel(logging.INFO)

  orderfile_shared.AddDummyFunctions(
    _GetUnpatchedOrderfileFilename(options), _GetOrderfileFilename(options)
  )


def main():
  parser = CreateArgumentParser()
  options = parser.parse_args()
  if options.verbosity >= 2:
    level = logging.DEBUG
  elif options.verbosity == 1:
    level = logging.INFO
  else:
    level = logging.WARNING
  logging.basicConfig(
    level=level, format='%(levelname).1s %(relativeCreated)6d %(message)s'
  )

  logging.info('Generate Profile Data')

  # Ensure that the output directory is an absolute path.
  options.out_dir = options.out_dir.resolve(strict=True)
  logging.info('Using options.out_dir=%s', options.out_dir)

  devices = device_utils.DeviceUtils.HealthyDevices()
  assert devices, 'Expected at least one connected device'
  device = devices[0]

  logging.getLogger().setLevel(logging.DEBUG)
  GenerateOrderfile(options, device)


if __name__ == '__main__':
  main()