#!/usr/bin/env bash
set -euo pipefail

role="${ARGUS_ROLE:-}"
if [[ -z "${role}" ]]; then
  echo "ERROR: ARGUS_ROLE must be set to head|worker" >&2
  exit 1
fi
if [[ "${role}" != "head" && "${role}" != "worker" ]]; then
  echo "ERROR: invalid ARGUS_ROLE=${role} (expected head|worker)" >&2
  exit 1
fi

shared_root="${ARGUS_SHARED_ROOT:-/private}"
cluster_name="${ARGUS_CLUSTER_NAME:-argus-ray}"
head_ip_file="${ARGUS_HEAD_IP_FILE:-${shared_root}/ray/discovery/${cluster_name}/head.json}"

ray_port="${ARGUS_RAY_PORT:-6379}"
dashboard_port="${ARGUS_DASHBOARD_PORT:-8265}"

ttl_s="${ARGUS_TTL_S:-60}"
refresh_s="${ARGUS_REFRESH_S:-10}"
poll_s="${ARGUS_POLL_S:-5}"

log_dir="${ARGUS_LOG_DIR:-${shared_root}/common/logs}"
mkdir -p "${log_dir}"

node_ip="${ARGUS_NODE_IP:-}"
if [[ -z "${node_ip}" ]]; then
  # Prefer the first IP returned by `hostname -i`.
  node_ip="$(hostname -i | awk '{print $1}')"
fi
if [[ -z "${node_ip}" ]]; then
  echo "ERROR: failed to determine container IP; set ARGUS_NODE_IP" >&2
  exit 1
fi

export PYTHONPATH="/opt/argus/py:${PYTHONPATH:-}"
export ARGUS_SHARED_ROOT="${shared_root}"
export ARGUS_CLUSTER_NAME="${cluster_name}"
export ARGUS_HEAD_IP_FILE="${head_ip_file}"
export ARGUS_NODE_IP="${node_ip}"
export ARGUS_RAY_PORT="${ray_port}"
export ARGUS_DASHBOARD_PORT="${dashboard_port}"
export ARGUS_TTL_S="${ttl_s}"
export ARGUS_REFRESH_S="${refresh_s}"
export ARGUS_POLL_S="${poll_s}"
export ARGUS_LOG_DIR="${log_dir}"

supervisor_conf="/tmp/supervisord.conf"
cat >"${supervisor_conf}" <<'CONF'
[supervisord]
nodaemon=true
pidfile=/tmp/supervisord.pid
logfile=/dev/null
childlogdir=/tmp
user=root

[supervisorctl]
serverurl=unix:///tmp/supervisor.sock

[unix_http_server]
file=/tmp/supervisor.sock
chmod=0700
CONF

append_program() {
  local name="$1"
  local cmd="$2"
  local stdout="$3"
  local stderr="$4"
  cat >>"${supervisor_conf}" <<CONF

[program:${name}]
command=${cmd}
autorestart=true
startsecs=1
startretries=999
stdout_logfile=${stdout}
stdout_logfile_maxbytes=0
stderr_logfile=${stderr}
stderr_logfile_maxbytes=0
CONF
}

if [[ "${role}" == "head" ]]; then
  append_program "argus_head_ray" "/usr/local/bin/argus-head-ray.sh" "/dev/fd/1" "/dev/fd/2"
  append_program "argus_head_publisher" "/usr/local/bin/argus-head-publisher.sh" "${log_dir}/argus_head_publisher.log" "${log_dir}/argus_head_publisher.err.log"
else
  append_program "argus_worker_watchdog" "/usr/local/bin/argus-worker-watchdog.sh" "${log_dir}/argus_worker_watchdog.log" "${log_dir}/argus_worker_watchdog.err.log"
fi

exec supervisord -n -c "${supervisor_conf}"
