#!/bin/sh

verbose=0

die () {
  echo
  echo "$*"
  echo
  exit 1
} >&2

error() {
  echo
  echo "ERROR: $*"
} >&2

checkJava() {
  set +eu
  if [ -n "$JAVA_HOME" ]; then
    JAVA_CMD=$JAVA_HOME/bin/java
    if [ ! -x "$JAVA_CMD" ]; then
      die "ERROR: JAVA_HOME is set to an invalid directory: $JAVA_HOME

  Please set the JAVA_HOME variable in your environment to match the
  location of your Java installation."
    fi
  else
    JAVA_CMD=java
    which java > /dev/null 2>&1 || die "ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH.

  Please set the JAVA_HOME variable in your environment to match the
  location of your Java installation."
  fi
  set -eu
}

set -eu

usage() {
  echo
  echo "Usage: $0 [options] <pid>          run collecting for target process"
  echo "    or $0 new-config               create a configuration from template"
  echo
  echo "Options:"
  echo "  -d duration                      run collecting for <duration> seconds"
  echo "  -c configuration_path            config file"
  echo "  -a architecture                  system architecture (x86_64, aarch64, etc.)"
  echo
  exit 1
} >&2

jattach() {
  set +e
  if [ $verbose -eq 1 ]; then
    echo "$JATTACH" "$PID" load instrument false "$COLLECTOR=$1"
  fi
  "$JATTACH" "$PID" load instrument false "$COLLECTOR=$1" >/dev/null
  RET=$?

  # Check if jattach failed
  if [ $RET -ne 0 ]; then
    error "Failed to inject collector into $PID, return code: $RET"
    exit $RET
  fi

  set -e
}

check_if_terminated() {
  if ! kill -0 "$PID" 2>/dev/null; then
    exit 0
  fi
}

CONTAINER_DIR=""
CONFIGURATION_FILE="configuration.toml"
CONFIGURATION_PATH=
SCRIPT_BIN="$0"
while [ -h "$SCRIPT_BIN" ]; do
  SCRIPT_BIN="$(readlink "$SCRIPT_BIN")"
done
SCRIPT_DIR="$(
  cd "$(dirname "$SCRIPT_BIN")" >/dev/null 2>&1
  pwd -P
)"

# 默认架构，可以通过 -a 参数覆盖
ARCH=""
JATTACH=$SCRIPT_DIR/jattach
COLLECTOR=$SCRIPT_DIR/continuous-profile-collector-agent.jar
PID=""

newConfig() {
  set +e
  checkJava
  "$JAVA_CMD" -cp "$COLLECTOR" "com.alibaba.cpc.ConfigurationGenerator" "$SCRIPT_DIR"
  RET=$?
  if [ $RET -ne 0 ]; then
    die "Failed to generate configuration"
  fi
  set -e
}

if [ $# -eq 1 ]; then
  if [ "$1" = "new-config" ]; then
    newConfig
    exit 0
  fi
fi

DURATION=10
while [ $# -gt 0 ]; do
  case $1 in
  -h | "-?")
    usage
    ;;
  -p)
    CONTAINER_DIR="$2"
    shift
    ;;
  -d)
    DURATION="$2"
    shift
    ;;
  -a)
    ARCH="$2"
    shift
    ;;
  -b | --start)
    DURATION="0"
    ;;
  -e | --stop)
    DURATION="-1"
    ;;
  -c)
    CONFIGURATION_FILE="$2"
    shift
    ;;
  -v)
    verbose=1
    ;;
  [0-9]*)
    PID="$1"
    ;;
  -*)
    error "Unrecognized option: $1"
    usage
    ;;
  esac
  shift
done

# 根据架构选择 jattach 文件（在参数解析后立即处理）
if [ -n "$ARCH" ]; then
  # 标准化架构名称
  case "$ARCH" in
    x86_64|amd64)
      ARCH="x86_64"
      ;;
    aarch64|arm64)
      ARCH="aarch64"
      ;;
  esac
  JATTACH=$SCRIPT_DIR/jattach-$ARCH
else
  # 如果没有指定架构，使用通用 jattach
  JATTACH=$SCRIPT_DIR/jattach
fi

genJattach() {
  set +e
  checkJava
  # 使用源目录的 jar 文件来生成 jattach，而不是容器目录中的
  local source_collector=$SCRIPT_DIR/continuous-profile-collector-agent.jar
  if [ ! -f "$source_collector" ]; then
    error "Source collector jar not found: $source_collector"
    die "Cannot generate jattach without collector jar"
  fi
  "$JAVA_CMD" -cp "$source_collector" "com.alibaba.cpc.JattachExtractor" "$SCRIPT_DIR"
  RET=$?
  if [ $RET -ne 0 ]; then
    die "Failed to generate jattach"
  fi

  # 如果生成的是通用 jattach，但我们需要架构特定的，则重命名
  if [ -f "$SCRIPT_DIR/jattach" ] && [ "$JATTACH" != "$SCRIPT_DIR/jattach" ]; then
    mv "$SCRIPT_DIR/jattach" "$JATTACH"
  fi
  
  chmod +x $JATTACH
  set -e
}

# 检查 jattach 文件是否存在，如果不存在则生成（此时 JATTACH 已经根据架构设置）
if [ ! -f "$JATTACH" ]; then
  # 如果架构特定的文件不存在，先尝试使用通用 jattach
  if [ "$JATTACH" != "$SCRIPT_DIR/jattach" ] && [ -f "$SCRIPT_DIR/jattach" ]; then
    JATTACH=$SCRIPT_DIR/jattach
  else
    # 如果都不存在，尝试生成
    genJattach
    # 生成后再次检查
    if [ ! -f "$JATTACH" ]; then
      error "jattach not found: $JATTACH"
      error "Failed to generate jattach"
      exit 1
    fi
  fi
fi

if [ ! "$CONTAINER_DIR" = "" ]; then
  CONFIGURATION_PATH="$CONTAINER_DIR/$CONFIGURATION_FILE"
  COLLECTOR=$CONTAINER_DIR/continuous-profile-collector-agent.jar
else
  CONFIGURATION_PATH=$SCRIPT_DIR/$CONFIGURATION_FILE
  if [ ! -f "$CONFIGURATION_PATH" ]; then
    error "config file not found: $CONFIGURATION_PATH"
    exit 1
  fi
fi

if [ "$PID" = "" ]; then
  error "No process specified"
  usage
fi

if [ $DURATION -gt 0 ]; then
  jattach "configuration_path=$CONFIGURATION_PATH"
  echo Collecting for "$DURATION" seconds >&2
  set +e
  trap 'DURATION=0' INT

  while [ "$DURATION" -gt 0 ]; do
    DURATION=$((DURATION - 1))
    check_if_terminated
    sleep 1
  done

  set -e
  trap - INT
  echo Done >&2
  jattach ""
elif [ $DURATION -eq 0 ]; then
  jattach "configuration_path=$CONFIGURATION_PATH"
  echo Profiling started >&2
elif [ $DURATION -lt 0 ]; then
  set -e
  jattach ""
  echo Profiling stopped >&2
fi