#!/bin/bash
#
# Copyright 2018, The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

usage() {
    cat <<EOF
Usage: run_app_with_prefetch --package <name> [OPTIONS]...

    -p, --package <name>        package of the app to test
    -a, --activity <name>       activity to use
    -h, --help                  usage information (this)
    -v, --verbose               enable extra verbose printing
    -i, --input <file>          trace file protobuf (default 'TraceFile.pb')
    -r, --readahead <mode>      cold, warm, fadvise, mlock (default 'warm')
    -w, --when <when>           aot or jit (default 'aot')
    -c, --count <count>         how many times to run (default 1)
    -s, --sleep <sec>           how long to sleep after readahead
    -t, --timeout <sec>         how many seconds to timeout in between each app run (default 10)
    -o, --output <file.csv>     what file to write the performance results into as csv (default stdout)
EOF
}

DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
source "$DIR/lib/common"

needs_trace_file="n"
input_file=""
package=""
mode='warm'
count=2
sleep_time=2
timeout=10
output="" # stdout by default
when="aot"
parse_arguments() {
  while [[ $# -gt 0 ]]; do
    case "$1" in
      -h|--help)
        usage
        exit 0
        ;;
      -p|--package)
        package="$2"
        shift
        ;;
      -a|--activity)
        activity="$2"
        shift
        ;;
      -i|--input)
        input_file="$2"
        shift
        ;;
      -v|--verbose)
        export verbose="y"
        ;;
      -r|--readahead)
        mode="$2"
        shift
        ;;
      -c|--count)
        count="$2"
        ((count+=1))
        shift
        ;;
      -s|--sleep)
        sleep_time="$2"
        shift
        ;;
      -t|--timeout)
        timeout="$2"
        shift
        ;;
      -o|--output)
        output="$2"
        shift
        ;;
      -w|--when)
        when="$2"
        shift
        ;;
      --compiler-filter)
        compiler_filter="$2"
        shift
        ;;
      *)
        echo "Invalid argument: $1" >&2
        exit 1
    esac
    shift
  done
}

echo_to_output_file() {
  if [[ "x$output" != x ]]; then
    echo "$@" >> $output
  fi
  # Always echo to stdout as well.
  echo "$@"
}

find_package_path() {
  local pkg="$1"

  res="$(adb shell find "/data/app/$pkg"-'*' -maxdepth 0 2> /dev/null)"
  if [[ -z $res ]]; then
    res="$(adb shell find "/system/app/$pkg"-'*' -maxdepth 0 2> /dev/null)"
  fi
  echo "$res"
}

# Main entry point
if [[ $# -eq 0 ]]; then
  usage
  exit 1
else
  parse_arguments "$@"

  # if we do not have have package exit early with an error
  [[ "$package" == "" ]] && echo "--package not specified" 1>&2 && exit 1

  if [[ $mode != "cold" && $mode != "warm" ]]; then
    needs_trace_file="y"
    if [[ -z "$input_file" ]] || ! [[ -f $input_file ]]; then
      echo "--input not specified" 1>&2
      exit 1
    fi
  fi

  if [[ "$activity" == "" ]]; then
    activity="$(get_activity_name "$package")"
    if [[ "$activity" == "" ]]; then
      echo "Activity name could not be found, invalid package name?" 1>&2
      exit 1
    else
      verbose_print "Activity name inferred: " "$activity"
    fi
  fi
fi

adb root > /dev/null

if [[ ($when == jit) || ($when == aot) ]] && [[ "$(adb shell getenforce)" != "Permissive" ]]; then
  echo "Disable selinux permissions and restart framework."
  adb shell setenforce 0
  adb shell stop
  adb shell start
  adb wait-for-device
fi

# TODO: set performance governor etc, preferrably only once
# before every single app run.

# Kill everything before running.
remote_pkill "$package"
sleep 1

timings_array=()

package_path="$(find_package_path "$package")"
if [[ $? -ne 0 ]]; then
  echo "Failed to detect package path for '$package'" >&2
  exit 1
fi
verbose_print "Package was in path '$package_path'"

keep_application_trace_file=n
application_trace_file_path="$package_path/TraceFile.pb"
trace_file_directory="$package_path"
if [[ $needs_trace_file == y ]]; then
  # system server always passes down the package path in a hardcoded spot.
  if [[ $when == "jit" ]]; then
    verbose_print adb push "$input_file" "$application_trace_file_path"
    adb push "$input_file" "$application_trace_file_path"
    keep_application_trace_file="y"
  else
    # otherwise use a temporary directory to get normal non-jit behavior.
    trace_file_directory="/data/local/tmp/prefetch/$package"
    adb shell mkdir -p "$trace_file_directory"
    verbose_print  adb push "$input_file" "$trace_file_directory/TraceFile.pb"
    adb push "$input_file" "$trace_file_directory/TraceFile.pb"
  fi
fi

# Everything other than JIT: remove the trace file,
# otherwise system server activity hints will kick in
# and the new just-in-time app pre-warmup will happen.
if [[ $keep_application_trace_file == "n" ]]; then
  adb shell "[[ -f '$application_trace_file_path' ]] && rm '$application_trace_file_path'"
fi

# Perform AOT readahead/pinning/etc when an application is about to be launched.
# For JIT readahead, we allow the system to handle it itself (this is a no-op).
#
# For warm, cold, etc modes which don't need readahead this is always a no-op.
perform_aot() {
  local the_when="$1" # user: aot, jit
  local the_mode="$2" # warm, cold, fadvise, mlock, etc.

  if [[ $the_when != "aot" ]]; then
    # TODO: just in time implementation.. should probably use system server.
    return 0
  fi

  # any non-warm/non-cold modes should use the iorap-activity-hint wrapper script.
  if [[ $the_mode != 'warm' && $the_mode != 'cold' ]]; then

    # TODO: add activity_hint_sender.exp
    verbose_print "starting with package=$package package_path=$trace_file_directory"
    coproc hint_sender_fd { $ANDROID_BUILD_TOP/system/iorap/src/sh/activity_hint_sender.exp "$package" "$trace_file_directory" "$the_mode"; }
    hint_sender_pid=$!
    verbose_print "Activity hint sender began"

    notification_success="n"
    while read -r -u "${hint_sender_fd[0]}" hint_sender_output; do
      verbose_print "$hint_sender_output"
      if [[ "$hint_sender_output" == "Press any key to send completed event..."* ]]; then
        verbose_print "WE DID SEE NOTIFICATION SUCCESS."
        notification_success='y'
        # Give it some time to actually perform the readaheads.
        sleep $sleep_time
        break
      fi
    done

    if [[ $notification_success == 'n' ]]; then
      echo "[FATAL] Activity hint notification failed." 1>&2
      exit 1
    fi
  fi
}

perform_aot_cleanup() {
  local the_when="$1" # user: aot, jit
  local the_mode="$2" # warm, cold, fadvise, mlock, etc.

  if [[ $the_when != "aot" ]]; then
    # TODO: just in time implementation.. should probably use system server.
    return 0
  fi

  # any non-warm/non-cold modes should use the iorap-activity-hint wrapper script.
  if [[ $the_mode != 'warm' && $the_mode != 'cold' ]]; then
    # Clean up the hint sender by telling it that the launch was completed,
    # and to shutdown the watcher.
    echo "Done\n" >&"${hint_sender_fd[1]}"

    while read -r -u "${hint_sender_fd[0]}" hint_sender_output; do
      verbose_print "$hint_sender_output"
    done

    wait $hint_sender_pid
  fi
}

configure_compiler_filter() {
  local the_compiler_filter="$1"
  local the_package="$2"
  local the_activity="$3"

  if [[ -z $the_compiler_filter ]]; then
    verbose_print "No --compiler-filter specified, don't need to force it."
    return 0
  fi

  local current_compiler_filter_info="$("$DIR"/query_compiler_filter.py --package "$the_package")"
  local res=$?
  if [[ $res -ne 0 ]]; then
    return $res
  fi

  local current_compiler_filter
  local current_reason
  local current_isa
  read current_compiler_filter current_reason current_isa <<< "$current_compiler_filter_info"

  verbose_print "Compiler Filter="$current_compiler_filter "Reason="$current_reason "Isa="$current_isa

  # Don't trust reasons that aren't 'unknown' because that means we didn't manually force the compilation filter.
  # (e.g. if any automatic system-triggered compilations are not unknown).
  if [[ $current_reason != "unknown" ]] || [[ $current_compiler_filter != $the_compiler_filter ]]; then
    verbose_print "$DIR"/force_compiler_filter --compiler-filter "$the_compiler_filter" --package "$the_package" --activity "$the_activity"
    "$DIR"/force_compiler_filter --compiler-filter "$the_compiler_filter" --package "$the_package" --activity "$the_activity"
    res=$?
  else
    verbose_print "Queried compiler-filter matched requested compiler-filter, skip forcing."
    res=0
  fi

  return $res
}

# Ensure the APK is currently compiled with whatever we passed in via --compiler-filter.
# No-op if this option was not passed in.
configure_compiler_filter "$compiler_filter" "$package" "$activity" || exit 1

# TODO: This loop logic could probably be moved into app_startup_runner.py
for ((i=0;i<count;++i)) do
  verbose_print "=========================================="
  verbose_print "====         ITERATION $i             ===="
  verbose_print "=========================================="
  if [[ $mode != "warm" ]]; then
    verbose_print "Drop caches for non-warm start."
    # Drop all caches to get cold starts.
    adb shell "echo 3 > /proc/sys/vm/drop_caches"
  fi

  perform_aot "$when" "$mode"

  verbose_print "Running with timeout $timeout"

  # TODO: multiple metrics output.
  total_time="$(timeout $timeout $DIR/launch_application "$package" "$activity")"

  if [[ $? -ne 0 ]]; then
    echo "WARNING: Skip bad result, try iteration again." >&2
    ((i=i-1))
    continue
  fi

  perform_aot_cleanup "$when" "$mode"

  echo "Iteration $i. Total time was: $total_time"

  timings_array+=($total_time)
done

# drop the first result which is usually garbage.
timings_array=("${timings_array[@]:1}")


# Print out interactive/debugging timings and averages.
# Other scripts should use the --output flag and parse the CSV.
for tim in "${timings_array[@]}"; do
  echo_to_output_file -ne "$tim,"
done
echo_to_output_file ""

average_string=$(echo "${timings_array[@]}" | awk '{s+=$0}END{print "Average:",s/NR}' RS=" ")
echo -ne ${average_string}.
if [[ x$output != x ]]; then
  echo " Saved results to '$output'"
fi

# Temporary hack around multiple activities being launched with different package paths (for same app):
# Clean up all left-over TraceFile.pb
adb shell 'for i in $(find /data/app -name TraceFile.pb); do rm \$i; done'

# Kill the process to ensure AM isn't keeping it around.
remote_pkill "$package"

exit 0