#!/bin/bash
COLORBLACK="\033[30m"
COLORRED="\033[31m"
COLORGREEN="\033[32m"
COLORBLUE="\033[34m"
COLORYELLOW="\033[33m"
COLORPURPLE="\033[35m"
COLORSKYBLUE="\033[36m"
COLORWHITE="\033[37m"
COLOREND="\033[0m"

timeout=$1 # timeout in seconds
# Default timeout is 60s.
if [ -z "$timeout" ]; then
  timeout=60
fi
if [ $timeout -le 30 ]; then
  echo -e "${COLORRED}[START ERROR]$COLOREND timeout should be greater than 30s."
  exit 1
fi

rank_size=$2 # mpi rank size
# Default rank_size is 1.
if [ -z "$rank_size" ]; then
  rank_size=1
fi

if ! command -v lsof >/dev/null 2>&1; then
  echo -e "${COLORRED}[START ERROR]$COLOREND lsof not installed, please install it first."
  exit 1
fi

if ! command -v curl >/dev/null 2>&1; then
  echo -e "${COLORRED}[START ERROR]$COLOREND curl not installed, please install it first."
  exit 1
fi

if ! command -v tail >/dev/null 2>&1; then
  echo -e "${COLORRED}[START ERROR]$COLOREND tail not installed, please install it first."
  exit 1
fi

# Check if the server is running.
PID="./PID"
if [ -e $PID ]; then
  pid=$(cat $PID)
  if [ -n "$(ps -p $pid -o pid=)" ]; then
    echo -e "${COLORPURPLE}[ALREADY RUNNING]$COLOREND"
    exit 0
  fi
fi

rm -f $PID
rm -f nohup.out

if [ ! -e "conf/server.yml" ]; then
  echo -e "${COLORRED}[START ERROR]$COLOREND conf/server.yml not exist."
  exit 1
fi
if [ ! -e "conf/inference.yml" ]; then
  echo -e "${COLORRED}[START ERROR]$COLOREND conf/inference.yml not exist."
  exit 1
fi

parse_yaml() {
    local yaml_file="$1"
    local key_str="$2"
    local result=$(python3 -c "import yaml; conf = yaml.load(open('$yaml_file', 'r'), Loader=yaml.FullLoader); print(conf${key_str})")
    echo "$result"
}

# parse server.yml to get framework and port.
mode=$(parse_yaml conf/server.yml [\'interface\'][\'framework\'])
mode=${mode// /}
host=$(parse_yaml conf/server.yml [\'interface\'][\'host\'])
host=${host// /}
port=$(parse_yaml conf/server.yml [\'interface\'][\'port\'])
port=${port%,*} # remove string after ',' and itself
port=${port// /}
log_dir=$(parse_yaml conf/server.yml [\'log\'][\'log_dir\'])
log_dir=${log_dir// /}

# mkdir log_dir if not exist.
if [ ! -e "$log_dir" ]; then
  mkdir -p "$log_dir"
else
  # check if is file.
  if [ -f "$log_dir" ]; then
    echo -e "${COLORRED}[START ERROR]$COLOREND $log_dir is a file, please remove it first."
    exit 1
  fi
  # clear old logs.
  echo -e "${COLORYELLOW}[WARNING]$COLOREND $log_dir already exist, clearing ${log_dir}/grps_server.log ${log_dir}/grps_usr.log ${log_dir}/grps_monitor.log ..."
  if [ -e "${log_dir}/grps_server.log" ]; then
    rm "${log_dir}"/grps_server.log*
  fi
  if [ -e "${log_dir}/grps_usr.log" ]; then
    rm "${log_dir}"/grps_usr.log*
  fi
  if [ -e "${log_dir}/grps_monitor.log" ]; then
    rm "${log_dir}"/grps_monitor.log*
  fi
fi

# Check if port is already used by lsof.
lsof -i:$port >/dev/null 2>&1
if [ $? -eq 0 ]; then
  echo -e "${COLORRED}[START ERROR]$COLOREND Port $port is already used."
  exit 1
fi

pip3 install -r requirements.txt

# Start server.
# For check if critical log generated by current starting, in order to check if server started failed.
start_time_str=$(date '+%Y%m%d_%H:%M:%S')
echo "Starting server in $mode mode..."
# Add libjemalloc.so to LD_PRELOAD to avoid memory fragmentation.
libjemalloc_path="/usr/local/lib/libjemalloc.so"
if [ -e "$libjemalloc_path" ]; then
  export LD_PRELOAD=$libjemalloc_path
  echo "Start server with libjemalloc."
fi
#nohup gunicorn -c conf/gunicorn.conf.py main:flask_app >nohup.out &
if [ $rank_size -gt 1 ]; then
  nohup mpirun --allow-run-as-root -np $rank_size -H localhost:$rank_size python3 grps_framework/main.py >nohup.out &
else
  nohup python3 grps_framework/main.py >nohup.out &
fi

# Loop check server status.
# clear http and https proxy
unset http_proxy
unset https_proxy
url="http://${host}:${port}/grps/v1/health/live"
# echo $url
retry=0
end_time_str=$(date '+%Y%m%d_%H:%M:%S.%3N') # Used to show new log.
while true; do
  sleep 1

  # Show new logs if have.
  if [ -e "${log_dir}/grps_server.log" ]; then
    while read -r line; do
      log_time_str=$(echo "$line" | awk '{print $1"_"$2}')
      if [[ ! "$log_time_str" =~ ^[0-9]{8}_[0-9]{2}:[0-9]{2}:[0-9]{2}.[0-9]{3} ]]; then
        continue
      fi
      if [ "$log_time_str" \> "$end_time_str" -a "$line" != "" ]; then
        echo "$line"
      fi
    done <<< "$(tail -n 30 ${log_dir}/grps_server.log)"
  fi
  if [ -e "${log_dir}/grps_usr.log" ]; then
    while read -r line; do
      log_time_str=$(echo "$line" | awk '{print $1"_"$2}')
      if [[ ! "$log_time_str" =~ ^[0-9]{8}_[0-9]{2}:[0-9]{2}:[0-9]{2}.[0-9]{3} ]]; then
        continue
      fi
      if [ "$log_time_str" \> "$end_time_str" -a "$line" != "" ]; then
        echo "$line"
      fi
    done <<< "$(tail -n 30 ${log_dir}/grps_usr.log)"
  fi
  end_time_str=$(date '+%Y%m%d_%H:%M:%S.%3N') # Used to show new log.

  # Check if server started failed.
  crit_log=""
  if [ -e "${log_dir}/grps_server.log" ]; then
    crit_log=$(cat ${log_dir}/grps_server.log | grep "CRITICAL]")
  fi
  if [ "$crit_log" != "" ]; then
    crit_time_str=$(echo "$crit_log" | awk '{print $1"_"$2}')
    if [ "$crit_time_str" \> "$start_time_str" ]; then
      echo -e "${COLORRED}[START ERROR]$COLOREND $crit_log"
      bash stop_server.sh >/dev/null 2>&1
      exit 1
    fi
  fi

  # Check if server process is died.
  if [ -e $PID ]; then
    pid=$(cat $PID | awk '{print $1}')
    if [ -z "$(ps -p $pid -o pid=)" ]; then
      echo -e "${COLORRED}[START ERROR]$COLOREND Server process is died. Maybe conf is invalid or env is not satisfied."
      bash stop_server.sh >/dev/null 2>&1
      tail -30 nohup.out
      if [ -e "${log_dir}/grps_server.log" ]; then
        crit_log=$(cat ${log_dir}/grps_server.log | grep "FATAL]")
        echo "$crit_log"
      fi
      exit 1
    fi
  else
    if [ "$retry" -ge 30 ]; then
      echo -e "${COLORRED}[START ERROR]$COLOREND No server process. Maybe conf is invalid or env is not satisfied or you can try again."
      tail -30 nohup.out
      exit 1
    fi
  fi

  # Check if server is ready.
  ret_code=$(curl -s --connect-timeout 0 --max-time 1 -o /dev/null -w "%{http_code}" "${url}")
  if [ "$ret_code" = "200" ]; then
    # Show last logs.
    if [ -e "${log_dir}/grps_server.log" ]; then
      while read -r line; do
        log_time_str=$(echo "$line" | awk '{print $1"_"$2}')
        if [[ ! "$log_time_str" =~ ^[0-9]{8}_[0-9]{2}:[0-9]{2}:[0-9]{2}.[0-9]{3} ]]; then
          continue
        fi
        if [ "$log_time_str" \> "$end_time_str" -a "$line" != "" ]; then
          echo "$line"
        fi
      done <<< "$(tail -n 30 ${log_dir}/grps_server.log)"
    fi
    if [ -e "${log_dir}/grps_usr.log" ]; then
      while read -r line; do
        log_time_str=$(echo "$line" | awk '{print $1"_"$2}')
        if [[ ! "$log_time_str" =~ ^[0-9]{8}_[0-9]{2}:[0-9]{2}:[0-9]{2}.[0-9]{3} ]]; then
          continue
        fi
        if [ "$log_time_str" \> "$end_time_str" -a "$line" != "" ]; then
          echo "$line"
        fi
      done <<< "$(tail -n 30 ${log_dir}/grps_usr.log)"
    fi
    echo -e "${COLORGREEN}[START SUCCESS]$COLOREND"
    echo -e "${COLORGREEN}>> You can access http://${host}:${port}/ to observe server status.$COLOREND"
    echo -e "${COLORGREEN}>> You can check service info by \`grpst ps\` cmd.$COLOREND"
    echo -e "${COLORGREEN}>> Will show logs after 5s, you can stop showing by \`Ctrl+C\`.$COLOREND"
    sleep 5
    exit 0
  fi

  # Retry.
  retry=$((retry + 1))
  if [ "$retry" -gt $timeout ]; then
    echo -e "${COLORRED}[START TIMEOUT]$COLOREND"
    bash stop_server.sh >/dev/null 2>&1
    exit 1
  fi
done
