#!/usr/bin/env bash
# Run a web service (any container image) as a GA4GH TES task and make it reachable at
# 127.0.0.1:<HUB_PORT> on this VM through the tunnel hub (tes-hub.sh in this directory).
#
# A service is a directory next to tes-hub/ with two files:
#   service.conf  NAME, IMAGE, PORT (port in the pod); optional HUB_PORT (port on this VM,
#                 default PORT), POD_SHELL (/bin/sh), RAM_GB (1), CPU_CORES (1). Any other
#                 KEY=value is passed to the pod as an environment variable.
#   pod.sh        runs in the task after common.sh: starts the service on 127.0.0.1:$PORT,
#                 calls start_tunnel "R:127.0.0.1:$HUB_PORT:127.0.0.1:$PORT", stays in the
#                 foreground. Generated state goes to <service>/.state/.
#
# Usage: tes-service.sh <service-dir> <command>
#   up         start: hub user (once), task (unless running), keep-alive loop; waits for the tunnel
#   down       stop: keep-alive loop, lock the old pod out of the hub (new password + hub
#              restart, because cancel does not stop the pod on TESK), cancel
#   restart    down + up
#   status     task, tunnel and keep-alive loop
#   logs       task logs (TESK fills them when the task ends)
#   selftest   short task without tunnel (pod.sh sees SELFTEST=1); fails unless SELFTEST_OK
#   print      the TES request JSON
#   setup | submit | cancel | watch    single steps, used by up/down

set -euo pipefail

HERE="$(cd "$(dirname "$0")" && pwd)"
[[ $# -ge 2 ]] || { sed -n '2,/^set -euo/p' "$0" | sed '$d; s/^# \{0,1\}//'; exit 1; }
if [[ -d "$1" ]]; then DIR="$(cd "$1" && pwd)"; elif [[ -d "$HERE/../$1" ]]; then DIR="$(cd "$HERE/../$1" && pwd)"
else echo "ERROR: no service directory '$1'" >&2; exit 1; fi
CMD="$2"; SELF="$0 $1"
[[ -f "$DIR/service.conf" && -f "$DIR/pod.sh" ]] || { echo "ERROR: $DIR needs service.conf and pod.sh" >&2; exit 1; }

STATE_DIR="$DIR/.state"
mkdir -p "$STATE_DIR"; chmod 700 "$STATE_DIR"
[[ -f "$STATE_DIR/config" ]] && source "$STATE_DIR/config"
source "$DIR/service.conf"
CONF_KEYS="$(sed -n 's/^[[:space:]]*\([A-Za-z_][A-Za-z0-9_]*\)=.*/\1/p' "$DIR/service.conf")"

# ---- configuration -----------------------------------------------------------------------
TES_URL="${TES_URL:-https://csc-tesk-noauth.rahtiapp.fi}"           # TESK @ ELIXIR-FI
TES_API="${TES_URL%/}/ga4gh/tes/v1"
HUB="${HUB:-$HERE/tes-hub.sh}"                                      # shared tunnel hub
: "${NAME:?NAME missing in service.conf}" "${IMAGE:?IMAGE missing in service.conf}" "${PORT:?PORT missing in service.conf}"
HUB_PORT="${HUB_PORT:-$PORT}"
POD_SHELL="${POD_SHELL:-/bin/sh}"
RAM_GB="${RAM_GB:-1}"
CPU_CORES="${CPU_CORES:-1}"
HUB_AUTH="${HUB_AUTH:-}"                                            # generated by setup
[[ "${HUB_AUTH%%:*}" == "$NAME" ]] || HUB_AUTH=""                   # e.g. .state copied from another service
WATCH_INTERVAL="${WATCH_INTERVAL:-60}"
# ------------------------------------------------------------------------------------------

die() { echo "ERROR: $*" >&2; exit 1; }
json_get() { python3 -c 'import json,sys; d=json.load(sys.stdin); print(eval(sys.argv[1], {"d": d}))' "$1"; }
ID_FILE="$STATE_DIR/task"
WATCH_PID="$STATE_DIR/watch.pid"
task_id() { [[ -s "$ID_FILE" ]] && cat "$ID_FILE" || true; }
need_id() { [[ -s "$ID_FILE" ]] || die "not submitted"; cat "$ID_FILE"; }
task_state() { curl -fsS -m 30 "$TES_API/tasks/$1?view=MINIMAL" | json_get 'd["state"]'; }
current_state() { local id; id="$(task_id)"; [[ -n "$id" ]] && { task_state "$id" 2>/dev/null || echo UNREACHABLE; } || echo NONE; }
is_done() { [[ "$1" =~ ^(NONE|COMPLETE|EXECUTOR_ERROR|SYSTEM_ERROR|CANCELED|PREEMPTED)$ ]]; }
tunnel_up() { [[ "$(ss -tln 2>/dev/null)" == *"127.0.0.1:$HUB_PORT "* ]]; }
# The pid must still be this service's loop: after a reboot, another process may have it.
watching() {
  local pid; pid="$(cat "$WATCH_PID" 2>/dev/null)" || return 1
  [[ -n "$pid" ]] && tr '\0' ' ' < "/proc/$pid/cmdline" 2>/dev/null | grep -qF -- "tes-service.sh $DIR watch"
}
print_logs() {
  json_get '"state: " + d["state"] + "\n" + "\n".join((l.get("stdout") or "") + (l.get("stderr") or "") for t in d.get("logs",[]) for l in t.get("logs",[]))'
}

# build_task [selftest]
build_task() {
  local selftest="${1:-}" k
  [[ -n "$selftest" || -n "$HUB_AUTH" ]] || die "credentials missing (run: $SELF setup)"
  local HUB_URL CHISEL_URL HUB_FINGERPRINT; eval "$("$HUB" env)"
  # service.conf values (minus the launcher's own keys) become the pod's environment.
  for k in $CONF_KEYS PORT HUB_PORT; do
    [[ $k =~ ^(NAME|IMAGE|POD_SHELL|RAM_GB|CPU_CORES|TES_URL)$ ]] || printf '%s\0%s\0' "$k" "${!k}"
  done > "$STATE_DIR/pod.env"
  export HUB_AUTH HUB_URL CHISEL_URL HUB_FINGERPRINT NAME IMAGE POD_SHELL RAM_GB CPU_CORES
  SELFTEST="${selftest:+1}" ENV_FILE="$STATE_DIR/pod.env" \
  SCRIPT="$(cat "$HERE/common.sh" "$DIR/pod.sh")" python3 - <<'PY'
import json, os
e = os.environ
selftest = e["SELFTEST"] == "1"
f = open(e["ENV_FILE"]).read().split("\0")[:-1]
env = {"ROLE": e["NAME"], **dict(zip(f[::2], f[1::2]))}
env.update({"SELFTEST": "1"} if selftest else {
    k: e[k] for k in ("HUB_AUTH", "HUB_URL", "HUB_FINGERPRINT", "CHISEL_URL")})
print(json.dumps({
    "name": e["NAME"] + ("-selftest" if selftest else ""),
    "description": f"{e['NAME']}, see tes-service.sh",
    "tags": {"service": e["NAME"], "managed-by": "tes-service.sh"},
    "resources": {"cpu_cores": int(e["CPU_CORES"]), "ram_gb": float(e["RAM_GB"]), "preemptible": False},
    "executors": [{"image": e["IMAGE"], "command": [e["POD_SHELL"], "-c", e["SCRIPT"]], "env": env}],
}, indent=2))
PY
  rm -f "$STATE_DIR/pod.env"
}

post_task() {
  build_task "$@" | curl -fsS --connect-timeout 10 -m 60 -X POST -H 'Content-Type: application/json' \
    --data-binary @- "$TES_API/tasks" | json_get 'd["id"]'
}

cmd_setup() {
  if [[ -z "$HUB_AUTH" ]]; then
    HUB_AUTH="$NAME:$(python3 -c 'import secrets; print(secrets.token_urlsafe(24))')"
    printf 'HUB_AUTH=%q\n' "$HUB_AUTH" > "$STATE_DIR/config"; chmod 600 "$STATE_DIR/config"
    echo "generated hub credentials in $STATE_DIR/config"
  fi
  "$HUB" add-user "$HUB_AUTH" "^R:127\\.0\\.0\\.1:$HUB_PORT\$"
}

# Lock out the pods of this service: a new hub password and a hub restart.
lock_out() { rm -f "$STATE_DIR/config"; HUB_AUTH=""; cmd_setup; "$HUB" restart; }

# One submission at a time (up, watch and submit may run together), and only if no task
# of this service runs, so that two tasks never compete for the port.
cmd_submit() {
  [[ "$("$HUB" status)" == *"hub running"* ]] || die "hub not running (run: $HUB start)"
  exec 9>"$STATE_DIR/submit.lock"
  flock -w 120 9 || die "another submission is running"
  local id state; state="$(current_state)"
  is_done "$state" || { echo "task $(task_id) is $state, not submitting"; flock -u 9; return 0; }
  id="$(post_task)" || { flock -u 9; die "submitting to $TES_API failed"; }
  echo "$id" > "$ID_FILE"; echo "submitted $id"
  flock -u 9
}

cmd_cancel() {
  local id; id="$(task_id)"
  [[ -n "$id" ]] && curl -fsS -m 30 -X POST "$TES_API/tasks/$id:cancel" >/dev/null && echo "canceled $id" || true
}

cmd_status() {
  local id; id="$(task_id)"
  echo "task: ${id:-NOT_SUBMITTED} ${id:+$(current_state)}"
  tunnel_up && echo "tunnel up: http://localhost:$HUB_PORT" || echo "tunnel not connected"
  watching && echo "keep-alive loop running (pid $(cat "$WATCH_PID"))" || echo "keep-alive loop not running"
}

cmd_selftest() {
  local id i out
  id="$(post_task selftest)" || die "submitting to $TES_API failed"; echo "selftest: $id"
  for ((i = 0; i < 60; i++)); do
    is_done "$(task_state "$id" 2>/dev/null || echo UNKNOWN)" && break
    sleep 10
  done
  out="$(curl -fsS -m 60 "$TES_API/tasks/$id?view=FULL" | print_logs)"; echo "$out"
  [[ "$out" == "state: COMPLETE"* && "$out" == *SELFTEST_OK* ]] || die "selftest failed"
}

cmd_watch() {
  echo $$ > "$WATCH_PID"
  trap 'rm -f "$WATCH_PID"' EXIT
  echo "$(date -Is) keeping $NAME alive on $TES_URL (every ${WATCH_INTERVAL}s)"
  local state
  while true; do
    state="$(current_state)"
    if is_done "$state"; then
      echo "$(date -Is) task $(task_id) is $state, resubmitting"
      # A canceled task's pod may still hold the port: lock it out first.
      if tunnel_up; then echo "$(date -Is) port $HUB_PORT still taken, locking the old pod out"; lock_out || true; fi
      (cmd_submit) || true
    fi
    sleep "$WATCH_INTERVAL"
  done
}

cmd_up() {
  [[ -n "$HUB_AUTH" ]] || cmd_setup
  local state i; state="$(current_state)"
  if is_done "$state"; then cmd_submit; else echo "task $(task_id) is $state, not resubmitting"; fi
  if ! watching; then
    setsid -f "$HERE/tes-service.sh" "$DIR" watch >> "$STATE_DIR/watch.log" 2>&1 < /dev/null
    echo "keep-alive loop started (log: $STATE_DIR/watch.log)"
  fi
  echo "waiting for the tunnel (usually 1-3 minutes; Ctrl-C is safe, the task keeps running)"
  for ((i = 0; i < 60; i++)); do tunnel_up && break; sleep 10; done
  tunnel_up || die "tunnel not up after 10 minutes, see: $SELF status"
  echo "$NAME is up on this VM at 127.0.0.1:$HUB_PORT. On your laptop run:"
  echo "  ssh -N -L $HUB_PORT:127.0.0.1:$HUB_PORT $(id -un)@$(hostname -f)"
  echo "and open http://localhost:$HUB_PORT"
}

cmd_down() {
  if watching; then kill "$(cat "$WATCH_PID")" && rm -f "$WATCH_PID" && echo "keep-alive loop stopped"; fi
  # Cancel marks the task CANCELED but the pod keeps running: lock it out of the hub first,
  # so that an interrupted down never leaves the old pod connected.
  lock_out
  cmd_cancel
}

case "$CMD" in
  up)       cmd_up ;;
  down)     cmd_down ;;
  restart)  cmd_down; cmd_up ;;
  setup)    cmd_setup ;;
  submit)   cmd_submit ;;
  status)   cmd_status ;;
  logs)     id="$(need_id)"; curl -fsS -m 60 "$TES_API/tasks/$id?view=FULL" | print_logs ;;
  cancel)   cmd_cancel ;;
  watch)    cmd_watch ;;
  selftest) cmd_selftest ;;
  print)    build_task ;;
  *) die "unknown command '$CMD' (run $0 without arguments for help)" ;;
esac