#!/usr/bin/env bash
# Post-update: restart service, health-check, rollback on failure.
set -euo pipefail

ROOT="$(cd "$(dirname "$0")/.." && pwd)"
cd "$ROOT"

PENDING="${ROOT}/data/.pending_restart.json"
LISTEN="${WESP_LISTEN:-0.0.0.0:80}"
HEALTH_URL="${WESP_HEALTH_URL:-http://127.0.0.1/api/health}"
HEALTH_TIMEOUT="${WESP_HEALTH_TIMEOUT_SEC:-60}"
HEALTH_INTERVAL="${WESP_HEALTH_INTERVAL_SEC:-2}"

wesp_python() {
  local venv_dir py
  if [[ -n "${WESP_PYTHON:-}" && ( -x "${WESP_PYTHON}" || -f "${WESP_PYTHON}" ) ]]; then
    echo "${WESP_PYTHON}"
    return
  fi
  venv_dir="${WESP_VENV:-/opt/wes}"
  for py in python python3 python3.13 python3.12 python3.11; do
    if [[ -x "${venv_dir}/bin/${py}" ]]; then
      echo "${venv_dir}/bin/${py}"
      return
    fi
  done
  if [[ -x "${ROOT}/../wes_mac/bin/python" ]]; then
    echo "${ROOT}/../wes_mac/bin/python"
    return
  fi
  if [[ -x "${ROOT}/.venv/bin/python" ]]; then
    echo "${ROOT}/.venv/bin/python"
    return
  fi
  command -v python3
}

wesp_waitress() {
  local venv_dir
  venv_dir="${WESP_VENV:-/opt/wes}"
  if [[ -x "${venv_dir}/bin/waitress-serve" ]]; then
    echo "${venv_dir}/bin/waitress-serve"
    return
  fi
  local py
  py="$(wesp_python)"
  if [[ -x "${py%/*}/waitress-serve" ]]; then
    echo "${py%/*}/waitress-serve"
    return
  fi
  if [[ -x "${ROOT}/../wes_mac/bin/waitress-serve" ]]; then
    echo "${ROOT}/../wes_mac/bin/waitress-serve"
    return
  fi
  if [[ -x "${ROOT}/.venv/bin/waitress-serve" ]]; then
    echo "${ROOT}/.venv/bin/waitress-serve"
    return
  fi
  command -v waitress-serve
}

stop_service() {
  pkill -f "waitress-serve.*wsgi:app" 2>/dev/null || true
  sleep 2
}

start_service() {
  if [[ "${WESP_SKIP_RESTART:-}" == "1" ]]; then
    return 0
  fi
  LOG="${ROOT}/data/logs/wesp-stdout.log"
  mkdir -p "${ROOT}/data/logs"
  export WESP_LISTEN="$LISTEN"
  nohup bash "${ROOT}/scripts/wesp_waitress_serve.sh" >>"$LOG" 2>&1 &
  disown || true
}

health_ok() {
  if [[ -n "${WESP_HEALTH_CMD:-}" ]]; then
    bash -c "$WESP_HEALTH_CMD"
    return $?
  fi
  curl -sf "$HEALTH_URL" | grep -q '"ok"[[:space:]]*:[[:space:]]*true'
}

mark_success() {
  local PY
  PY="$(wesp_python)"
  "$PY" -c "
import sys
from pathlib import Path
root = sys.argv[1]
sys.path.insert(0, root)
from app.services.update_state_store import mark_success, read_update_state
state = read_update_state(root)
mark_success(root, target_version=str(state.get('target_version') or '?'))
Path(root, 'data', '.pending_restart.json').unlink(missing_ok=True)
" "$ROOT"
}

rollback_and_restart() {
  local PY
  PY="$(wesp_python)"
  "$PY" "${ROOT}/scripts/rollback_update.py" --root "$ROOT" || true
  stop_service
  start_service
}

restart_systemd_unit() {
  local UNIT="${WESP_SYSTEMD_UNIT:?}"
  echo "post_update: systemctl restart ${UNIT}" >&2
  systemctl restart "${UNIT}" || true
}

if [[ -n "${WESP_SYSTEMD_UNIT:-}" ]]; then
  restart_systemd_unit
else
  stop_service
  start_service
fi

deadline=$((SECONDS + HEALTH_TIMEOUT))
while (( SECONDS < deadline )); do
  if health_ok; then
    mark_success
    exit 0
  fi
  sleep "$HEALTH_INTERVAL"
done

rollback_and_restart
exit 1