#!/usr/bin/env python3
# Copyright (c) 2022, 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.

import argparse
import os
import subprocess
import sys
import threading
import time

from enum import Enum

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import profile_utils
import utils

DEVNULL=subprocess.DEVNULL

class ProcessReader(threading.Thread):

  def __init__(self, process):
    threading.Thread.__init__(self)
    self.lines = []
    self.process = process

  def run(self):
    for line in self.process.stdout:
      line = line.decode('utf-8').strip()
      self.lines.append(line)

  def stop(self):
    self.process.kill()

class ScreenState(Enum):
  OFF_LOCKED = 1,
  OFF_UNLOCKED = 2
  ON_LOCKED = 3
  ON_UNLOCKED = 4

  def is_off(self):
    return self == ScreenState.OFF_LOCKED or self == ScreenState.OFF_UNLOCKED

  def is_on(self):
    return self == ScreenState.ON_LOCKED or self == ScreenState.ON_UNLOCKED

  def is_on_and_locked(self):
    return self == ScreenState.ON_LOCKED

  def is_on_and_unlocked(self):
    return self == ScreenState.ON_UNLOCKED

def broadcast(action, component, device_id=None):
  print('Sending broadcast %s' % action)
  cmd = create_adb_cmd('shell am broadcast -a %s %s' % (action, component), device_id)
  return subprocess.check_output(cmd).decode('utf-8').strip().splitlines()

def build_apks_from_bundle(bundle, output, overwrite=False):
  print('Building %s' % bundle)
  cmd = [
      'java', '-jar', utils.BUNDLETOOL_JAR,
      'build-apks',
      '--bundle=%s' % bundle,
      '--output=%s' % output]
  if overwrite:
    cmd.append('--overwrite')
  subprocess.check_call(cmd, stdout=DEVNULL, stderr=DEVNULL)

def capture_screen(target, device_id=None):
  print('Taking screenshot to %s' % target)
  tmp = '/sdcard/screencap.png'
  cmd = create_adb_cmd('shell screencap -p %s' % tmp, device_id)
  subprocess.check_call(cmd, stdout=DEVNULL, stderr=DEVNULL)
  pull(tmp, target, device_id)

def create_adb_cmd(arguments, device_id=None):
  assert isinstance(arguments, list) or isinstance(arguments, str)
  cmd = ['adb']
  if device_id is not None:
    cmd.append('-s')
    cmd.append(device_id)
  cmd.extend(arguments if isinstance(arguments, list) else arguments.split(' '))
  return cmd

def capture_app_profile_data(app_id, device_id=None):
  ps_cmd = create_adb_cmd('shell ps -o NAME', device_id)
  stdout = subprocess.check_output(ps_cmd).decode('utf-8').strip()
  killed_any = False
  for process_name in stdout.splitlines():
    if process_name.startswith(app_id):
      print('Flushing profile for process %s' % process_name)
      killall_cmd = create_adb_cmd(
          'shell killall -s SIGUSR1 %s' % process_name, device_id)
      killall_result = subprocess.run(killall_cmd, capture_output=True)
      stdout = killall_result.stdout.decode('utf-8')
      stderr = killall_result.stderr.decode('utf-8')
      if killall_result.returncode == 0:
        killed_any = True
      else:
        print('Error: stdout: %s, stderr: %s' % (stdout, stderr))
      time.sleep(5)
  assert killed_any, 'Expected to find at least one process'

def check_app_has_profile_data(app_id, device_id=None):
  profile_path = get_profile_path(app_id)
  cmd = create_adb_cmd(
      'shell du /data/misc/profiles/cur/0/%s/primary.prof' % app_id,
      device_id)
  stdout = subprocess.check_output(cmd).decode('utf-8').strip()
  size_str = stdout[:stdout.index('\t')]
  assert size_str.isdigit()
  size = int(size_str)
  if size == 4:
    raise ValueError('Expected size of profile at %s to be > 4K' % profile_path)

def clear_logcat(device_id=None):
  cmd = create_adb_cmd('logcat -c', device_id)
  subprocess.check_call(cmd, stdout=DEVNULL, stderr=DEVNULL)

def clear_profile_data(app_id, device_id=None):
  cmd = create_adb_cmd(
      'shell cmd package compile --reset %s' % app_id, device_id)
  subprocess.check_call(cmd, stdout=DEVNULL, stderr=DEVNULL)

def drop_caches(device_id=None):
  cmd = create_adb_cmd(
      ['shell', 'echo 3 > /proc/sys/vm/drop_caches'], device_id)
  subprocess.check_call(cmd, stdout=DEVNULL, stderr=DEVNULL)

def ensure_screen_on(device_id=None):
  if get_screen_state(device_id).is_off():
    toggle_screen(device_id)
  assert get_screen_state(device_id).is_on()

def ensure_screen_off(device_id=None):
  if get_screen_state(device_id).is_on():
    toggle_screen(device_id)
  assert get_screen_state(device_id).is_off()

def force_compilation(app_id, device_id=None):
  print('Applying AOT (full)')
  cmd = create_adb_cmd(
      'shell cmd package compile -m speed -f %s' % app_id, device_id)
  subprocess.check_call(cmd, stdout=DEVNULL, stderr=DEVNULL)

def force_profile_compilation(app_id, device_id=None):
  print('Applying AOT (profile)')
  cmd = create_adb_cmd(
      'shell cmd package compile -m speed-profile -f %s' % app_id, device_id)
  subprocess.check_call(cmd, stdout=DEVNULL, stderr=DEVNULL)

def get_apk_path(app_id, device_id=None):
  cmd = create_adb_cmd('shell pm path %s' % app_id, device_id)
  stdout = subprocess.check_output(cmd).decode('utf-8').strip()
  if not stdout.startswith('package:'):
    raise ValueError(
        'Expected stdout to start with "package:", was: %s' % stdout)
  apk_path = stdout[len('package:'):]
  if not apk_path.endswith('.apk'):
    raise ValueError(
        'Expected stdout to end with ".apk", was: %s' % stdout)
  return apk_path

def get_component_name(app_id, activity):
  if activity.startswith(app_id):
    return '%s/.%s' % (app_id, activity[len(app_id)+1:])
  else:
    return '%s/%s' % (app_id, activity)

def get_meminfo(app_id, device_id=None):
  cmd = create_adb_cmd('shell dumpsys meminfo -s %s' % app_id, device_id)
  stdout = subprocess.check_output(cmd).decode('utf-8').strip()
  for line in stdout.splitlines():
    if 'TOTAL PSS: ' in line:
      elements = [s for s in line.replace('TOTAL ', 'TOTAL_').split()]
      assert elements[0] == 'TOTAL_PSS:', elements[0]
      assert elements[1].isdigit()
      assert elements[2] == 'TOTAL_RSS:'
      assert elements[3].isdigit()
      return { 'total_pss': int(elements[1]), 'total_rss': int(elements[3]) }
  raise ValueError('Unexpected stdout: %s' % stdout)

def get_profile_data(app_id, device_id=None):
  with utils.TempDir() as temp:
    source = get_profile_path(app_id)
    target = os.path.join(temp, 'primary.prof')
    pull(source, target, device_id)
    with open(target, 'rb') as f:
      return f.read()

def get_profile_path(app_id):
  return '/data/misc/profiles/cur/0/%s/primary.prof' % app_id

def get_minor_major_page_faults(app_id, device_id=None):
  pid = get_pid(app_id, device_id)
  cmd = create_adb_cmd('shell ps -p %i -o MINFL,MAJFL' % pid, device_id)
  stdout = subprocess.check_output(cmd).decode('utf-8')
  lines_it = iter(stdout.splitlines())
  first_line = next(lines_it)
  assert first_line == ' MINFL  MAJFL'
  second_line = next(lines_it)
  minfl, majfl = second_line.split()
  assert minfl.isdigit()
  assert majfl.isdigit()
  return (int(minfl), int(majfl))

def get_pid(app_id, device_id=None):
  cmd = create_adb_cmd('shell pidof %s' % app_id, device_id)
  stdout = subprocess.check_output(cmd).decode('utf-8').strip()
  assert stdout.isdigit()
  pid = int(stdout)
  return pid

def get_screen_state(device_id=None):
  cmd = create_adb_cmd('shell dumpsys nfc', device_id)
  stdout = subprocess.check_output(cmd).decode('utf-8').strip()
  screen_state_value = None
  for line in stdout.splitlines():
    if line.startswith('mScreenState='):
      value_start_index = len('mScreenState=')
      screen_state_value=line[value_start_index:]
  if screen_state_value is None:
    raise ValueError('Expected to find mScreenState in: adb shell dumpsys nfc')
  if not hasattr(ScreenState, screen_state_value):
    raise ValueError(
        'Expected mScreenState to be a value of ScreenState, was: %s'
            % screen_state_value)
  return ScreenState[screen_state_value]

def get_classes_and_methods_from_app_profile(app_id, device_id=None):
  apk_path = get_apk_path(app_id, device_id)
  profile_path = get_profile_path(app_id)
  cmd = create_adb_cmd(
    'shell profman --dump-classes-and-methods'
    ' --profile-file=%s --apk=%s --dex-location=%s'
        % (profile_path, apk_path, apk_path), device_id)
  stdout = subprocess.check_output(cmd).decode('utf-8').strip()
  lines = stdout.splitlines()
  return profile_utils.parse_art_profile(lines)

def get_screen_off_timeout(device_id=None):
  cmd = create_adb_cmd(
      'shell settings get system screen_off_timeout', device_id)
  stdout = subprocess.check_output(cmd).decode('utf-8').strip()
  assert stdout.isdigit()
  screen_off_timeout = int(stdout)
  return screen_off_timeout

def grant(app_id, permission, device_id=None):
  cmd = create_adb_cmd('shell pm grant %s %s' % (app_id, permission), device_id)
  subprocess.check_call(cmd)

def install(apk, device_id=None):
  print('Installing %s' % apk)
  cmd = create_adb_cmd('install %s' % apk, device_id)
  stdout = subprocess.check_output(cmd).decode('utf-8')
  assert 'Success' in stdout

def install_apks(apks, device_id=None, max_attempts=3):
  print('Installing %s' % apks)
  cmd = [
      'java', '-jar', utils.BUNDLETOOL_JAR,
      'install-apks',
      '--apks=%s' % apks]
  if device_id is not None:
    cmd.append('--device-id=%s' % device_id)
  for i in range(max_attempts):
    process_result = subprocess.run(cmd, capture_output=True)
    stdout = process_result.stdout.decode('utf-8')
    stderr = process_result.stderr.decode('utf-8')
    if process_result.returncode == 0:
      return
    print('Failed to install %s' % apks)
    print('Stdout: %s' % stdout)
    print('Stderr: %s' % stderr)
    print('Retrying...')
  raise Exception('Unable to install apks in %s attempts' % max_attempts)

def install_bundle(bundle, device_id=None):
  print('Installing %s' % bundle)
  with utils.TempDir() as temp:
    apks = os.path.join(temp, 'Bundle.apks')
    build_apks_from_bundle(bundle, apks)
    install_apks(apks, device_id)

def install_profile_using_adb(app_id, host_profile_path, device_id=None):
  device_profile_path = get_profile_path(app_id)
  cmd = create_adb_cmd('push %s %s' % (host_profile_path, device_profile_path))
  subprocess.check_call(cmd)
  stop_app(app_id, device_id)
  force_profile_compilation(app_id, device_id)

def install_profile_using_profileinstaller(app_id, device_id=None):
  # This assumes that the profileinstaller library has been added to the app,
  # https://developer.android.com/jetpack/androidx/releases/profileinstaller.
  action = 'androidx.profileinstaller.action.INSTALL_PROFILE'
  component = '%s/androidx.profileinstaller.ProfileInstallReceiver' % app_id
  stdout = broadcast(action, component, device_id)
  assert len(stdout) == 2
  assert stdout[0] == ('Broadcasting: Intent { act=%s flg=0x400000 cmp=%s }' % (action, component))
  assert stdout[1] == 'Broadcast completed: result=1', stdout[1]
  stop_app(app_id, device_id)
  force_profile_compilation(app_id, device_id)

def issue_key_event(key_event, device_id=None, sleep_in_seconds=1):
  cmd = create_adb_cmd('shell input keyevent %s' % key_event, device_id)
  stdout = subprocess.check_output(cmd).decode('utf-8').strip()
  assert len(stdout) == 0
  time.sleep(sleep_in_seconds)

def launch_activity(
    app_id,
    activity,
    device_id=None,
    intent_data_uri=None,
    wait_for_activity_to_launch=False):
  args = ['shell', 'am', 'start', '-n', '%s/%s' % (app_id, activity)]
  if intent_data_uri:
    args.extend(['-d', intent_data_uri])
  if wait_for_activity_to_launch:
    args.append('-W')
  cmd = create_adb_cmd(args, device_id)
  stdout = subprocess.check_output(cmd).decode('utf-8').strip()
  assert stdout.startswith('Starting: Intent {')
  expected_component = 'cmp=%s' % get_component_name(app_id, activity)
  assert expected_component in stdout, \
      'was %s, expected %s' % (stdout, expected_component)
  lines = stdout.splitlines()
  result = {}
  for line in lines:
    if line.startswith('TotalTime: '):
      total_time_str = line.removeprefix('TotalTime: ')
      assert total_time_str.isdigit()
      result['total_time'] = int(total_time_str)
  assert not wait_for_activity_to_launch or 'total_time' in result, lines
  return result

def navigate_to_home_screen(device_id=None):
  cmd = create_adb_cmd('shell input keyevent KEYCODE_HOME', device_id)
  subprocess.check_call(cmd, stdout=DEVNULL, stderr=DEVNULL)

def prepare_for_interaction_with_device(device_id=None, device_pin=None):
  # Increase screen off timeout to avoid device screen turns off.
  twenty_four_hours_in_millis = 24 * 60 * 60 * 1000
  previous_screen_off_timeout = get_screen_off_timeout(device_id)
  set_screen_off_timeout(twenty_four_hours_in_millis, device_id)

  # Unlock device.
  unlock(device_id, device_pin)

  teardown_options = {
    'previous_screen_off_timeout': previous_screen_off_timeout
  }
  return teardown_options

def pull(source, target, device_id=None):
  cmd = create_adb_cmd('pull %s %s' % (source, target), device_id)
  subprocess.check_call(cmd, stdout=DEVNULL, stderr=DEVNULL)

def root(device_id=None):
  cmd = create_adb_cmd('root', device_id)
  subprocess.check_call(cmd, stdout=DEVNULL, stderr=DEVNULL)

def set_screen_off_timeout(screen_off_timeout_in_millis, device_id=None):
  cmd = create_adb_cmd(
      'shell settings put system screen_off_timeout %i'
          % screen_off_timeout_in_millis,
      device_id)
  stdout = subprocess.check_output(cmd).decode('utf-8').strip()
  assert len(stdout) == 0

def start_logcat(device_id=None, format=None, filter=None, silent=False):
  args = ['logcat']
  if format:
    args.extend(['--format', format])
  if silent:
    args.append('-s')
  if filter:
    args.append(filter)
  cmd = create_adb_cmd(args, device_id)
  logcat_process = subprocess.Popen(
      cmd, bufsize=1024*1024, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  reader = ProcessReader(logcat_process)
  reader.start()
  return reader

def stop_logcat(logcat_reader):
  logcat_reader.stop()
  logcat_reader.join()
  return logcat_reader.lines

def stop_app(app_id, device_id=None):
  print('Shutting down %s' % app_id)
  cmd = create_adb_cmd('shell am force-stop %s' % app_id, device_id)
  subprocess.check_call(cmd, stdout=DEVNULL, stderr=DEVNULL)

def teardown_after_interaction_with_device(teardown_options, device_id=None):
  # Reset screen off timeout.
  set_screen_off_timeout(
      teardown_options['previous_screen_off_timeout'],
      device_id)

def toggle_screen(device_id=None):
  issue_key_event('KEYCODE_POWER', device_id)

def uninstall(app_id, device_id=None):
  print('Uninstalling %s' % app_id)
  cmd = create_adb_cmd('uninstall %s' % app_id, device_id)
  process_result = subprocess.run(cmd, capture_output=True)
  stdout = process_result.stdout.decode('utf-8')
  stderr = process_result.stderr.decode('utf-8')
  if process_result.returncode == 0:
    assert 'Success' in stdout
  elif stdout.startswith('cmd: Failure calling service package: Broken pipe'):
    assert app_id == 'com.google.android.youtube'
    print('Waiting after broken pipe')
    time.sleep(15)
  else:
    expected_error = (
        'java.lang.IllegalArgumentException: Unknown package: %s' % app_id)
    assert 'Failure [DELETE_FAILED_INTERNAL_ERROR]' in stdout \
        or expected_error in stderr, \
        'stdout: %s, stderr: %s' % (stdout, stderr)

def unlock(device_id=None, device_pin=None):
  ensure_screen_on(device_id)
  screen_state = get_screen_state(device_id)
  assert screen_state.is_on(), 'was %s' % screen_state
  if screen_state.is_on_and_locked():
    if device_pin is not None:
      raise NotImplementedError('Device unlocking with pin not implemented')
    issue_key_event('KEYCODE_MENU', device_id)
    screen_state = get_screen_state(device_id)
  assert screen_state.is_on_and_unlocked(), 'was %s' % screen_state

def parse_options(argv):
  result = argparse.ArgumentParser(description='Run adb utils.')
  result.add_argument('--capture-screen',
                      help='Capture screen to given file')
  result.add_argument('--device-id',
                      help='Device id (e.g., emulator-5554).')
  result.add_argument('--device-pin',
                      help='Device pin code (e.g., 1234)')
  result.add_argument('--ensure-screen-off',
                      help='Ensure screen off',
                      action='store_true',
                      default=False)
  result.add_argument('--get-screen-state',
                      help='Get screen state',
                      action='store_true',
                      default=False)
  result.add_argument('--unlock',
                      help='Unlock device',
                      action='store_true',
                      default=False)
  options, args = result.parse_known_args(argv)
  return options, args

def main(argv):
  (options, args) = parse_options(argv)
  if options.capture_screen:
    capture_screen(options.capture_screen, options.device_id)
  if options.ensure_screen_off:
    ensure_screen_off(options.device_id)
  elif options.get_screen_state:
    print(get_screen_state(options.device_id))
  elif options.unlock:
    unlock(options.device_id, options.device_pin)


if __name__ == '__main__':
  sys.exit(main(sys.argv[1:]))