#!/usr/bin/env bash

# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.

source "$(dirname "$0")/uniffle-function.sh"
UNIFFLE_SHELL_EXECNAME="uniffle"

function uniffle_usage
{
  uniffle_add_option "--daemon (start|status|stop)" "operate on a daemon"
  uniffle_add_subcommand "client-cli" client "prints uniffle-cli args information"
  uniffle_add_subcommand "admin-cli" admin "prints uniffle-admin args information"
  uniffle_add_subcommand "shuffle-server" daemon "run the shuffle-server"
  uniffle_add_subcommand "coordinator" daemon "run the coordinator"
  uniffle_generate_usage "${UNIFFLE_SHELL_EXECNAME}" true
}

uniffle_parse_args "$@"
shift "${UNIFFLE_PARSE_COUNTER}"

function uniffle_cmd_case
{
  subcmd=$1
  shift

  case ${subcmd} in
    client-cli)
      UNIFFLE_CLASSNAME=org.apache.uniffle.cli.UniffleCLI
    ;;
    admin-cli)
      UNIFFLE_CLASSNAME=org.apache.uniffle.cli.UniffleAdminCLI
    ;;
    shuffle-server)
      MAX_DIRECT_MEMORY_OPTS=""
      if [ -n "${MAX_DIRECT_MEMORY_SIZE:-}" ]; then
        MAX_DIRECT_MEMORY_OPTS="-XX:MaxDirectMemorySize=$MAX_DIRECT_MEMORY_SIZE"
      fi

      JVM_ARGS=" -server \
          -Xmx${XMX_SIZE} \
          -Xms${XMX_SIZE} \
          ${MAX_DIRECT_MEMORY_OPTS} \
          -XX:+UseG1GC \
          -XX:MaxGCPauseMillis=200 \
          -XX:ParallelGCThreads=20 \
          -XX:ConcGCThreads=5 \
          -XX:InitiatingHeapOccupancyPercent=20 \
          -XX:G1HeapRegionSize=32m \
          -XX:+UnlockExperimentalVMOptions \
          -XX:G1NewSizePercent=10 \
          -XX:+PrintGC \
          -XX:+PrintAdaptiveSizePolicy \
          -XX:+PrintGCDateStamps \
          -XX:+PrintGCTimeStamps \
          -XX:+PrintGCDetails \
          -Xloggc:${RSS_LOG_DIR}/gc-%t.log"

      JAVA11_EXTRA_ARGS=" -XX:+IgnoreUnrecognizedVMOptions \
          -Xlog:gc:tags,time,uptime,level"
      UNIFFLE_OPTS="$UNIFFLE_OPTS $JVM_ARGS $JAVA11_EXTRA_ARGS"

      UNIFFLE_SUBCMD_SUPPORT_DAEMONIZATION="true"
      RSS_CONF_FILE="${RSS_CONF_DIR}/server.conf"
      uniffle_add_classpath "${RSS_HOME}/jars/server/*"
      UNIFFLE_CLASSNAME=org.apache.uniffle.server.ShuffleServer
    ;;
    coordinator)
      JVM_ARGS=" -server \
          -Xmx${XMX_SIZE:-8g} \
          -Xms${XMX_SIZE:-8g} \
          -XX:+UseG1GC \
          -XX:MaxGCPauseMillis=200 \
          -XX:ParallelGCThreads=20 \
          -XX:ConcGCThreads=5 \
          -XX:InitiatingHeapOccupancyPercent=45 \
          -XX:+PrintGC \
          -XX:+PrintAdaptiveSizePolicy \
          -XX:+PrintGCDateStamps \
          -XX:+PrintGCTimeStamps \
          -XX:+PrintGCDetails \
          -Xloggc:${RSS_LOG_DIR}/gc-%t.log"

      JAVA11_EXTRA_ARGS=" -XX:+IgnoreUnrecognizedVMOptions \
          -Xlog:gc:tags,time,uptime,level"
      UNIFFLE_OPTS="$UNIFFLE_OPTS $JVM_ARGS $JAVA11_EXTRA_ARGS"

      UNIFFLE_SUBCMD_SUPPORT_DAEMONIZATION="true"
      RSS_CONF_FILE="${RSS_CONF_DIR}/coordinator.conf"
      uniffle_add_classpath "${RSS_HOME}/jars/coordinator/*"
      UNIFFLE_CLASSNAME=org.apache.uniffle.coordinator.CoordinatorServer
    ;;
    *)
      UNIFFLE_CLASSNAME="${subcmd}"
      if ! uniffle_validate_classname "${UNIFFLE_CLASSNAME}"; then
        uniffle_exit_with_usage 1
      fi
    ;;
  esac
}

if [[ $# = 0 ]]; then
  uniffle_exit_with_usage 1
fi

UNIFFLE_SUBCMD=$1
shift

case $UNIFFLE_SUBCMD in
--help|-help|-h)
  uniffle_exit_with_usage 0
  exit
  ;;
esac

UNIFFLE_SUBCMD_ARGS=("$@")

source "$(dirname "$0")/utils.sh"
UNIFFLE_SHELL_SCRIPT_DEBUG=false
load_rss_env
uniffle_java_setup
uniffle_basic_init
uniffle_finalize_uniffle_opts

CLASSPATH=""

JAR_DIR="${RSS_HOME}/jars"
for file in $(ls ${JAR_DIR}/cli/*.jar 2>/dev/null); do
  CLASSPATH=$CLASSPATH:$file
done

HADOOP_DEPENDENCY="$("$HADOOP_HOME/bin/hadoop" classpath --glob)"
CLASSPATH=$CLASSPATH:$HADOOP_CONF_DIR:$HADOOP_DEPENDENCY

set +u
uniffle_cmd_case "${UNIFFLE_SUBCMD}" "${UNIFFLE_SUBCMD_ARGS[@]}"
uniffle_generic_java_subcmd_handler
set -u