#!/usr/bin/env bash
#
# interactive_engine command tool
set -eo pipefail
# color
readonly RED="\033[0;31m"
readonly YELLOW="\033[1;33m"
readonly GREEN="\033[0;32m"
readonly NC="\033[0m" # No Color

GS_LOG=/var/log/graphscope

err() {
  echo -e "${RED}[$(date +'%Y-%m-%dT%H:%M:%S%z')]: [ERROR] $*${NC}" >&2
}

warning() {
  echo -e "${YELLOW}[$(date +'%Y-%m-%dT%H:%M:%S%z')]: [WARNING] $*${NC}" >&1
}

log() {
  echo -e "[$(date +'%Y-%m-%dT%H:%M:%S%z')]: $*" >&1
}

succ() {
  echo -e "${GREEN}[$(date +'%Y-%m-%dT%H:%M:%S%z')]: $*${NC}" >&1
}

##########################
# Output usage information.
# Globals:
#   None
# Arguments:
#   None
##########################
usage() {
cat <<END
  A script to launch interactive engine.

  Usage: giectl [options] [command] [parameters]

  Options:

    -h, --help           output help information

  Commands:

    create_gremlin_instance_on_local    create gremlin instance locally
    create_gremlin_instance_on_k8s      create gremlin instance on k8s
    close_gremlin_instance_on_local     close gremlin instance locally
    close_gremlin_instance_on_k8s       close instance on k8s
    start_frontend                      start frontend of instance
    start_executor                      start executor of instance
END
}

##########################
# Start the frontend of interactive engine instance.
# Globals:
#   GRAPHSCOPE_HOME
# Arguments:
#   GRAPHSCOPE_RUNTIME: runtime workspace
#   object_id: id of vineyard object
#   schema_path: path of graph schema file
#   pegasus_hosts: hosts and port of executor
#   executor_count: number of executor, equal to engine count
#   frontend_port: frontend port
##########################
start_frontend() {
  declare -r GRAPHSCOPE_RUNTIME=$1
  declare -r object_id=$2
  declare -r schema_path=$3
  declare -r pegasus_hosts=$4
  declare -r frontend_port=$5
  declare -r threads_per_worker=${THREADS_PER_WORKER:-2}

  # create related directories
  declare -r log_dir=${GS_LOG}/${object_id}
  declare -r config_dir=${GRAPHSCOPE_RUNTIME}/config/${object_id}
  declare -r pid_dir=${GRAPHSCOPE_RUNTIME}/pid/${object_id}
  mkdir -p ${log_dir} ${config_dir} ${pid_dir}

  declare java_opt="-server
            -verbose:gc
            -Xloggc:${log_dir}/frontend.gc.log
            -XX:+PrintGCDetails
            -XX:+PrintGCDateStamps
            -XX:+PrintHeapAtGC
            -XX:+PrintTenuringDistribution
            -Djava.awt.headless=true
            -Dsun.net.client.defaultConnectTimeout=10000
            -Dsun.net.client.defaultReadTimeout=30000
            -XX:+DisableExplicitGC
            -XX:-OmitStackTraceInFastThrow
            -XX:+UseG1GC
            -XX:InitiatingHeapOccupancyPercent=75
            -XX:+IgnoreUnrecognizedVMOptions
            -Dfile.encoding=UTF-8
            -Dsun.jnu.encoding=UTF-8
            -Dlogfilename=${log_dir}/frontend.log
            -Dlogbasedir=${log_dir}/frontend
            -Dlog4j.configurationFile=file:${GRAPHSCOPE_HOME}/conf/log4j2.xml
            -Djna.library.path=${GRAPHSCOPE_HOME}/lib:"

  # set frontend config file
  sed -e "s@GRAPH_NAME@${object_id}@g" \
      -e "s@GRAPH_SCHEMA@${schema_path}@g" \
      -e "s@PEGASUS_HOSTS@${pegasus_hosts}@g" \
      -e "s@FRONTEND_SERVICE_PORT@${frontend_port}@g" \
      -e "s@THREADS_PER_WORKER@${threads_per_worker}@g" \
      ${GRAPHSCOPE_HOME}/conf/frontend.vineyard.properties > ${config_dir}/frontend.vineyard.properties

  # frontend service hold a handle client of coordinator
  java ${java_opt} \
    -cp "${GRAPHSCOPE_HOME}/lib/*:" \
    com.alibaba.graphscope.frontend.Frontend \
    ${config_dir}/frontend.vineyard.properties \
    >> ${log_dir}/frontend.log 2>&1 &
  echo $! >> ${pid_dir}/frontend.pid
}

##########################
# Start the executor of interactive engine instance.
# Globals:
#   GRAPHSCOPE_HOME
# Arguments:
#   GRAPHSCOPE_RUNTIME: runtime workspace
#   object_id: id of vineyard object
#   server_id: global id of executor server
#   server_size
#   rpc_port
#   network_servers
##########################
start_executor() {
  declare -r GRAPHSCOPE_RUNTIME=$1
  declare -r object_id=$2
  declare -r server_id=$3
  declare -r server_size=$4
  declare -r rpc_port=$5
  declare -r network_servers=$6
  declare -r threads_per_worker=${THREADS_PER_WORKER:-2}

  declare -r log_dir=${GS_LOG}/${object_id}
  declare -r config_dir=${GRAPHSCOPE_RUNTIME}/config/${object_id}
  declare -r pid_dir=${GRAPHSCOPE_RUNTIME}/pid/${object_id}
  mkdir -p ${log_dir} ${config_dir} ${pid_dir}
  # log4rs needs LOG_DIRS env
  export LOG_DIRS=${log_dir}
  export LD_LIBRARY_PATH=${GRAPHSCOPE_HOME}/lib:${LD_LIBRARY_PATH}
  export DYLD_LIBRARY_PATH=${GRAPHSCOPE_HOME}/lib:${DYLD_LIBRARY_PATH}

  # set executor config file
  sed -e "s@GRAPH_NAME@${object_id}@g" \
      -e "s@VINEYARD_OBJECT_ID@${object_id}@g" \
      -e "s@RPC_PORT@${rpc_port}@g" \
      -e "s@SERVER_ID@${server_id}@g" \
      -e "s@SERVER_SIZE@${server_size}@g" \
      -e "s@NETWORK_SERVERS@${network_servers}@g" \
      -e "s@THREADS_PER_WORKER@${threads_per_worker}@g" \
      ${GRAPHSCOPE_HOME}/conf/executor.vineyard.properties > ${config_dir}/executor.vineyard.properties

  cp ${GRAPHSCOPE_HOME}/conf/log4rs.yml  ${config_dir}/log4rs.yml

  declare executor_binary="gaia_executor"

  # launch executor
  declare flag="gaia_${object_id}executor"  # Flag is used to kill process
  RUST_BACKTRACE=full ${GRAPHSCOPE_HOME}/bin/gaia_executor ${config_dir}/log4rs.yml ${config_dir}/executor.vineyard.properties ${flag} \
    >> ${log_dir}/executor.log 2>&1 &
  echo $! >> ${pid_dir}/executor.pid
}

##########################
# create interactive instance on local.
# Globals:
#   GRAPHSCOPE_HOME
# Arguments:
#   GRAPHSCOPE_RUNTIME: runtime workspace
#   object_id: id of vineyard object
#   schema_path: path of graph schema file
#   server_id: global server id of executor
#   executor_port
#   executor_rpc_port
#   frontend_port
#   vineyard_ipc_socket
##########################
create_gremlin_instance_on_local() {
  declare -r GRAPHSCOPE_RUNTIME=$1
  declare -r object_id=$2
  declare -r schema_path=$3
  declare -r server_id=$4
  declare -r executor_port=$5
  declare -r executor_rpc_port=$6
  declare -r frontend_port=$7
  export VINEYARD_IPC_SOCKET=$8

  declare -r cluster_type="local"
  declare -r executor_count="1"  # local mode only start one executor

  if [[ ! -d "${GS_LOG}" || ! -w "${GS_LOG}" ]]; then
    # /var/log/graphscope is not existed/writable, switch to ${HOME}/.local/log/graphscope
    GS_LOG=${HOME}/.local/log/graphscope
  fi
  # init Graphscope log location
  readonly GS_LOG
  mkdir -p ${GS_LOG}

  declare -r log_dir=${GS_LOG}/${object_id}
  # make a "current" link
  unlink ${GS_LOG}/current || true
  ln -s ${log_dir} ${GS_LOG}/current

  server_size="1"
  # Frontend use executor rpc port
  pegasus_hosts="127.0.0.1:${executor_rpc_port}"

  start_frontend ${GRAPHSCOPE_RUNTIME} ${object_id} ${schema_path} ${pegasus_hosts} \
                 ${frontend_port}

  log "FRONTEND_ENDPOINT:127.0.0.1:${frontend_port}"
  # executor use executor inner port
  network_servers="127.0.0.1:${executor_port}"
  start_executor ${GRAPHSCOPE_RUNTIME} ${object_id} ${server_id} ${server_size} ${executor_rpc_port} \
                 ${network_servers}
}

##########################
# create interactive engine instance on k8s.
# Globals:
#   None
# Arguments:
#   GRAPHSCOPE_RUNTIME: runtime workspace
#   object_id: id of vineyard object
#   schema_path: path of graph schema file
#   pod_name_list:
#   engine_container: container name of engine
#   executor_port
#   executor_rpc_port
#   frontend_port
#   coordinator_name: name of coordinator deployment object in k8s
##########################
create_gremlin_instance_on_k8s() {
  declare -r GRAPHSCOPE_RUNTIME=$1
  declare -r object_id=$2
  declare -r schema_path=$3

  declare -r pod_hosts=$(echo $4 | awk -F"," '{for(i=1;i<=NF;++i) {print $i" "}}')
  declare -r server_size=$(echo $4 | awk -F"," '{print NF}')

  declare -r engine_container=$5
  declare -r executor_port=$6
  declare -r executor_rpc_port=$7
  declare -r frontend_port=$8
  declare -r coordinator_name=$9 # deployment name of coordinator

  instance_id=${coordinator_name##*-}

  pod_ips=$(kubectl get pod -lapp.kubernetes.io/component=engine,app.kubernetes.io/instance=${instance_id} -o jsonpath='{.items[*].status.podIP}')
  pegasus_hosts=""
  for pod in ${pod_ips}; do
    pegasus_hosts="${pegasus_hosts},${pod}:${executor_rpc_port}"
  done
  pegasus_hosts=${pegasus_hosts:1}

  frontend_name=$(kubectl get pod -lapp.kubernetes.io/component=frontend,app.kubernetes.io/instance=${instance_id} -o jsonpath='{.items[*].metadata.name}')

  launch_frontend_cmd="GRAPHSCOPE_HOME=${GRAPHSCOPE_HOME} \
      ${GRAPHSCOPE_HOME}/bin/giectl start_frontend \
      ${GRAPHSCOPE_RUNTIME} ${object_id} ${schema_path} ${pegasus_hosts} ${frontend_port}"
  kubectl cp ${schema_path} ${frontend_name}:${schema_path}
  kubectl exec ${frontend_name} -- /bin/bash -c "${launch_frontend_cmd}"
  

  network_servers=""
  for pod in ${pod_ips}; do
    network_servers="${network_servers},${pod}:${executor_port}"
  done
  network_servers=${network_servers:1}
  log "Launch interactive engine(executor) in per engine pod."
  _server_id=0
  for pod in $(echo ${pod_hosts})
  do
    launch_executor_cmd="GRAPHSCOPE_HOME=${GRAPHSCOPE_HOME} ${GRAPHSCOPE_HOME}/bin/giectl start_executor ${GRAPHSCOPE_RUNTIME} ${object_id} ${_server_id} ${server_size} ${executor_rpc_port} ${network_servers}"
    # kubectl exec ${pod} -c ${engine_container} -- sudo mkdir -p /var/log/graphscope
    # kubectl exec ${pod} -c ${engine_container} -- sudo chown -R graphscope:graphscope /var/log/graphscope
    kubectl exec ${pod} -c ${engine_container} -- /bin/bash -c "${launch_executor_cmd}"
    (( _server_id+=1 ))
  done

  log "Expose gremlin server."
  # random from range [50001, 51000) for interactive engine
  frontend_external_port=$(( ((RANDOM<<15)|RANDOM) % 1000 + 50000 ))
  frontend_deployment_name=$(kubectl get deployment -lapp.kubernetes.io/component=frontend,app.kubernetes.io/instance=${instance_id} -o jsonpath='{.items[*].metadata.name}')
  if [ "${GREMLIN_EXPOSE}" = "LoadBalancer" ]; then
    kubectl expose deployment ${frontend_deployment_name} --name=gremlin-${object_id} --port=${frontend_external_port} \
      --target-port=${frontend_port} --type=LoadBalancer 1>/dev/null 2>&1
    [ $? -eq 0 ] || exit 1
    wait_period_seconds=0
    while true
    do
      external_ip=$(kubectl describe service gremlin-${object_id} | grep "LoadBalancer Ingress" | awk -F'[ :]+' '{print $3}')
      if [ -n "${external_ip}" ]; then
        break
      fi
      wait_period_seconds=$(($wait_period_seconds+5))
      if [ ${wait_period_seconds} -gt ${timeout_seconds} ];then
        echo "Get external ip of ${GREMLIN_EXPOSE} failed."
        break
      fi
      sleep 5
    done
  else
    # NodePort service type
    # expose gremlin service
    kubectl expose deployment ${frontend_deployment_name} --name=gremlin-${object_id} --port=${frontend_external_port} \
      --target-port=${frontend_port} --type=NodePort 1>/dev/null 2>&1
    [ $? -eq 0 ] || exit 1
    wait_period_seconds=0
    while true
    do
      frontend_external_port=$(kubectl describe services gremlin-${object_id} | grep "NodePort" | grep "TCP" | tr -cd "0-9")
      if [ -n "${frontend_external_port}" ]; then
        break
      fi
      wait_period_seconds=$(($wait_period_seconds+5))
      if [ ${wait_period_seconds} -gt ${timeout_seconds} ];then
        log "Get interactive engine frontend node port failed."
        break
      fi
      sleep 5
    done
    wait_period_seconds=0
    while true
    do
      external_ip=$(kubectl describe pods ${frontend_deployment_name} | grep "Node:" | head -1 | awk -F '[ /]+' '{print $3}')
      if [ -n "${external_ip}" ]; then
        break
      fi
      wait_period_seconds=$(($wait_period_seconds+5))
      if [ ${wait_period_seconds} -gt ${timeout_seconds} ];then
        log "Get interactive engine frontend host ip of ${GREMLIN_EXPOSE} failed."
        break
      fi
      sleep 5
    done
  fi
  # currently support only 1 pod.
  frontend_ip=$(kubectl get pod -lapp.kubernetes.io/component=frontend,app.kubernetes.io/instance=${instance_id} -o jsonpath='{.items[*].status.podIP}')
  log "FRONTEND_ENDPOINT:${frontend_ip}:${frontend_port}"
  log "FRONTEND_EXTERNAL_ENDPOINT:${external_ip}:${frontend_external_port}"
}

##########################
# close interactive engine instance on local.
# Globals:
# Arguments:
#   GRAPHSCOPE_RUNTIME: runtime workspace
#   object_id: id of vineyard object
##########################
close_gremlin_instance_on_local() {
  declare -r GRAPHSCOPE_RUNTIME=$1
  declare -r object_id=$2
  declare -r pid_dir=${GRAPHSCOPE_RUNTIME}/pid/${object_id}

  declare -a components=("frontend" "executor")

  for component in "${components[@]}"; do
    declare str=$(cat ${pid_dir}/${component}.pid)

    # The file may have multiple pids, each in a single line
    # This will read each line into an array
    while read -r pid; do pids+=("$pid"); done <<<"${str}"

    for pid in "${pids[@]}"; do
        kill ${pid} || true
    done
  done
}

##########################
# Close interactive engine instance on k8s.
# Globals:
#   None
# Arguments:
#   GRAPHSCOPE_RUNTIME: runtime workspace
#   object_id: id of vineyard object
#   pod_name_list
#   engine_container
##########################
close_gremlin_instance_on_k8s() {
  declare -r GRAPHSCOPE_RUNTIME=$1
  declare -r object_id=$2
  declare -r pod_hosts=$(echo $3 | awk -F"," '{for(i=1;i<=NF;++i) {print $i" "}}')
  declare -r engine_container=$4
  declare -r instance_id=$5

  declare -r pid_dir=${GRAPHSCOPE_RUNTIME}/pid/${object_id}

  # delete service
  log "delete gremlin service"
  kubectl delete service gremlin-${object_id} || true

  # kill frontend and coordinator process
  log "Close frontend process."
  frontend_name=$(kubectl get pod -lapp.kubernetes.io/component=frontend,app.kubernetes.io/instance=${instance_id} -o jsonpath='{.items[*].metadata.name}')

  kill_frontend_cmd="ps -ef | grep ${object_id} | grep -v grep | awk '{print \$2}' | xargs kill -9"
  kubectl exec ${frontend_name} -- sh -c "${kill_frontend_cmd}"


  # kill executor process on engine container.
  log "Close executor process on engine container."
  for pod in $(echo ${pod_hosts})
  do
    kill_executor_process_cmd="ps -ef | grep gaia_${object_id}executor |
        grep -v grep | awk '{print \$2}' | xargs kill -9"
    kubectl exec ${pod} -c ${engine_container} -- sh -c "${kill_executor_process_cmd}"
  done
}

# parse argv
while test $# -ne 0; do
  arg=$1; shift
  case $arg in
    -h|--help) usage; exit ;;
    create_gremlin_instance_on_local) create_gremlin_instance_on_local "$@"; exit;;
    create_gremlin_instance_on_k8s) create_gremlin_instance_on_k8s "$@"; exit;;
    close_gremlin_instance_on_local) close_gremlin_instance_on_local "$@"; exit;;
    close_gremlin_instance_on_k8s) close_gremlin_instance_on_k8s "$@"; exit;;
    start_frontend) start_frontend "$@"; exit;;
    start_executor) start_executor "$@"; exit;;
    *)
      echo "unrecognized option or command '${arg}'"
      usage; exit;;
  esac
done

set +e
set +o pipefail
