#!/bin/bash
# Copyright (C) 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.
set -e
SCRIPT_DIR="$(realpath "$(dirname "${BASH_SOURCE[0]}")")"

if [ "$TMPDIR" == "" ]; then
  TMPDIR=/tmp
fi

function get_gn_value() {
  local out=$1
  local key=$2
  "$SCRIPT_DIR/gn" args "$out" --list=$key --short | awk '{print $3}' \
    | tr -d '"'
}

function is_monolithic() {
  local out=$1
  value=$(get_gn_value "$out" "monolithic_binaries")
  test "$value" == "true"
}

function is_android() {
  local out=$1
  value=$(get_gn_value "$out" "target_os")
  test "$value" == "android"
}

function is_ssh_target() {
  [[ -n "$SSH_TARGET" ]]
}

function is_mac() {
  # shellcheck disable=2251
  ! test -d /proc
  return $?
}

function tmux_ensure_bash() {
  if [[ $SHELL == *"fish" ]]; then
    tmux send-keys "bash" Enter
  fi
}

function reset_tracing() {
  if is_android "$OUT"; then
    adb shell 'test -d /sys/kernel/tracing && echo 0 > /sys/kernel/tracing/tracing_on'
  elif ! is_mac; then
    # shellcheck disable=SC2016
    local script='
    if [ ! -w /sys/kernel/tracing ]; then
      echo "tracefs not accessible, try sudo chown -R $USER /sys/kernel/tracing"
      sudo chown -R "$USER" /sys/kernel/tracing
    fi

    echo 0 > /sys/kernel/tracing/tracing_on
    '

    if is_ssh_target; then
      # shellcheck disable=SC2029
      ssh -t "$SSH_TARGET" "sh -c '$script'"
    else
      sh -c "$script"
    fi
  fi
}

function adb_supports_push_sync() {
  adb --help 2>&1 | grep 'push.*\[--sync\]' >/dev/null 2>&1
}

function push() {
  if is_android "$OUT"; then
    local maybe_sync=''
    if adb_supports_push_sync; then
      maybe_sync='--sync'
    fi
    echo adb push $maybe_sync "$1" "$DIR"
    adb push $maybe_sync "$1" "$DIR"
  elif is_ssh_target; then
    echo scp "$1" "$SSH_TARGET:$DIR"
    scp "$1" "$SSH_TARGET:$DIR"
  else
    echo cp "$1" "$DIR"
    cp "$1" "$DIR"
  fi
}

function pull() {
  if is_android "$OUT"; then
    echo adb pull "$DIR/$1" "$2"
    adb pull "$DIR/$1" "$2"
  elif is_ssh_target; then
    echo scp "$SSH_TARGET:$DIR/$1" "$2"
    scp "$SSH_TARGET:$DIR/$1" "$2"
  else
    echo mv "$DIR/$1" "$2"
    mv "$DIR/$1" "$2"
  fi
}

BACKGROUND=0
SKIP_CONVERTERS=0
TMUX_LAYOUT="even-vertical"
CPU_MASK=""

while getopts "bl:nt:c:C:z:s:" o; do
  case "$o" in
    b) BACKGROUND=1 ;;
    l) TMUX_LAYOUT=${OPTARG} ;;
    n) SKIP_CONVERTERS=1 ;;
    t) SSH_TARGET=${OPTARG} ;;
    c) CONFIG=${OPTARG} ;;
    C) OUT=${OPTARG} ;;
    z) CPU_MASK=${OPTARG} ;;
    s) SCRIPT=${OPTARG} ;;
    *)
      echo >&2 "Invalid option $o"
      exit
      ;;
  esac
done

# Allow out to be passed as argument
shift $((OPTIND - 1))
OUT="${OUT:-$1}"

# Provide useful usage information
if [ -z "$OUT" ]; then
  echo "Usage: $0 [OPTION]... [OUTPUT]"
  echo ""
  echo "Options:"
  echo "  -b          run in the background"
  echo "  -l LAYOUT   tmux pane layout"
  echo "  -n          skip post-trace convertors"
  echo "  -t TARGET   SSH device target"
  echo "  -c CONFIG   trace configuration file"
  echo "  -C OUTPUT   output directory"
  echo "  -z MASK     constrain binaries to given cpu mask (taskset syntax)"
  echo "  -s SCRIPT   a script to put into a tmux pane"
  echo ""
  echo "Environment variables:"
  echo "  SSH_TARGET  SSH device target"
  echo "  CONFIG      trace configuration file"
  echo "  OUT         output directory"
  exit 1
fi

# Warn about invalid output directories
if [ ! -f "$OUT/args.gn" ]; then
  echo >&2 "OUT=$OUT doesn't look like an output directory."
  echo >&2 "Please specify a directory by doing:"
  echo >&2 "  export OUT=out/xxx $0"
  exit 1
fi

# Check SSH target is valid
if is_ssh_target && ! ssh -q "$SSH_TARGET" exit; then
  echo >&2 "SSH_TARGET=$SSH_TARGET doesn't look like a valid SSH target."
  echo >&2 "Please specify a SSH cross-compilation target by doing:"
  echo >&2 "  export SSH_TARGET=<user>@<host> $0"
  exit 1
fi

if ! builtin type -P tmux &>/dev/null; then
  echo >&2 "tmux not found"
  exit 1
fi

# You can set the config to one of the files under test/configs e.g.
# CONFIG=ftrace.cfg or to :test. Defaults to :test.
CONFIG="${CONFIG:-:test}"

if is_android "$OUT"; then
  DIR=/data/local/tmp
elif is_ssh_target; then
  DIR=$(ssh "$SSH_TARGET" mktemp -d $TMPDIR/perfetto.XXXXXX)
elif is_mac; then
  DIR=$(mktemp -d $TMPDIR/perfetto.XXXXXX)
else
  DIR=$(mktemp -p $TMPDIR -d perfetto.XXXXXX)
fi

if is_android "$OUT"; then
  TRACECONV="gcc_like_host/traceconv"
else
  TRACECONV="traceconv"
fi


# (re)build the binaries
BUILD_TARGETS=(traced traced_probes perfetto test/configs)
if [[ SKIP_CONVERTERS -eq 0 ]]; then
  BUILD_TARGETS+=($TRACECONV)
fi

"$SCRIPT_DIR/ninja" -C "$OUT" ${BUILD_TARGETS[*]}

push "$OUT/traced"
push "$OUT/traced_probes"
push "$OUT/perfetto"
reset_tracing

if is_android "$OUT"; then
  PREFIX="PERFETTO_CONSUMER_SOCK_NAME=@perfetto_test_consumer PERFETTO_PRODUCER_SOCK_NAME=@perfetto_test_producer"
else
  PREFIX=""
fi

if ! is_monolithic "$OUT"; then
  PREFIX="$PREFIX LD_LIBRARY_PATH=$DIR"
  push "$OUT/libperfetto.so"
fi

CONFIG_DEVICE_PATH="$CONFIG"
CMD_OPTS=""
# Shorthand for using serialized test configs.
if [[ "$CONFIG" == *.protobuf ]]; then
  CONFIG_DEVICE_PATH="$CONFIG"
  CONFIG_PATH=$OUT/$CONFIG
  if [[ ! -f $CONFIG_PATH ]]; then
    echo >&2 "Config \"$CONFIG_PATH\" not known."
    exit 1
  fi
  push "$CONFIG_PATH"
elif [[ "$CONFIG" != ":test" ]]; then
  CONFIG_DEVICE_PATH="$(basename "$CONFIG")"
  CONFIG_PATH=$CONFIG
  # If path isn't valid, assume it's a name of a test config.
  if [[ ! -f $CONFIG_PATH ]]; then
    CONFIG_PATH=test/configs/$CONFIG
    if [[ ! -f $CONFIG_PATH ]]; then
      echo >&2 "Config \"$CONFIG\" not known."
      exit 1
    fi
  fi
  CMD_OPTS="--txt $CMD_OPTS"
  push "$CONFIG_PATH"
fi

if [[ -f "$SCRIPT" ]]; then
  push "$SCRIPT"
fi

POSTFIX=""

if [[ -n "$CPU_MASK" ]]; then
  PREFIX="$PREFIX taskset $CPU_MASK"
fi

if [[ BACKGROUND -eq 1 ]]; then
  PREFIX="$PREFIX nohup"
  POSTFIX=" &> /dev/null &"
fi

if tmux has-session -t demo; then
  tmux kill-session -t demo
fi
tmux -2 new-session -d -s demo

if [ ! -z "$ANDROID_ADB_SERVER_PORT" ]; then
  tmux set-environment -t demo ANDROID_ADB_SERVER_PORT $ANDROID_ADB_SERVER_PORT
fi

if tmux -V | awk '{split($2, ver, "."); if (ver[1] < 2) exit 1 ; else if (ver[1] == 2 && ver[2] < 1) exit 1 }'; then
  tmux set-option -g mouse on
else
  tmux set-option -g mode-mouse on
  tmux set-option -g mouse-resize-pane on
  tmux set-option -g mouse-select-pane on
  tmux set-option -g mouse-select-window on
fi

tmux split-window -v
tmux split-window -v

if [[ -n "$SCRIPT" ]]; then
  tmux split-window -v
fi

tmux select-layout "${TMUX_LAYOUT}"

tmux select-pane -t 0
tmux send-keys "clear" C-m
if is_android "$OUT"; then
  tmux send-keys "adb shell" C-m
fi

tmux select-pane -t 1
tmux send-keys "clear" C-m
if is_android "$OUT"; then
  tmux send-keys "adb shell" C-m
fi

tmux select-pane -t 2
tmux send-keys "clear" C-m
if is_android "$OUT"; then
  tmux send-keys "adb shell" C-m
fi

if [[ -n "$SCRIPT" ]]; then
  tmux select-pane -t 3
  tmux send-keys "clear" C-m
  if is_android "$OUT"; then
    tmux send-keys "adb shell" C-m
  fi
fi

sleep 2

tmux select-pane -t 1
if is_ssh_target; then
  tmux send-keys "ssh $SSH_TARGET" Enter
fi
tmux_ensure_bash
tmux send-keys "PS1='[traced]$ '" Enter
tmux send-keys "cd $DIR" Enter
tmux send-keys "clear" C-m
tmux send-keys "$PREFIX ./traced $POSTFIX" Enter

tmux select-pane -t 0
if is_ssh_target; then
  tmux send-keys "ssh $SSH_TARGET" Enter
fi
tmux_ensure_bash
tmux send-keys "PS1='[traced_probes]$ '" Enter
tmux send-keys "cd $DIR" Enter
tmux send-keys "clear" C-m
tmux send-keys "$PREFIX ./traced_probes $POSTFIX" Enter

tmux select-pane -t 2
if is_ssh_target; then
  tmux send-keys "ssh $SSH_TARGET" Enter
fi
tmux_ensure_bash
tmux send-keys "PS1='[consumer]$ '" Enter
tmux send-keys "cd $DIR" Enter
tmux send-keys "clear" C-m
tmux send-keys "$PREFIX ./perfetto $CMD_OPTS -c $CONFIG_DEVICE_PATH -o trace $POSTFIX"

if [[ -n "$SCRIPT" ]]; then
  tmux select-pane -t 3
  if is_ssh_target; then
    tmux send-keys "ssh $SSH_TARGET" Enter
  fi
  tmux_ensure_bash
  tmux send-keys "PS1='[script]$ '" Enter
  tmux send-keys "cd $DIR" Enter
  tmux send-keys "clear" C-m
  if [[ -f "$SCRIPT" ]]; then
    tmux send-keys "./$(basename "$SCRIPT")"
  else
    tmux send-keys "$SCRIPT"
  fi
fi

# Select consumer pane.
tmux select-pane -t 2

tmux -2 attach-session -t demo
if [[ BACKGROUND -eq 1 ]]; then
  exit 0
fi

reset_tracing

TRACE=$TMPDIR/trace
echo -e "\n\x1b[32mPulling trace into $TRACE.perfetto-trace\x1b[0m"
pull trace "$TRACE.perfetto-trace"

if [[ SKIP_CONVERTERS -eq 0 ]]; then
  echo -e "\n\x1b[32mPulling trace into $TRACE.pbtext\x1b[0m"
  "$OUT/$TRACECONV" text <"$TRACE.perfetto-trace" >"$TRACE.pbtext"
fi