#!/bin/bash
# service-ple: idempotent installer. Running it twice changes nothing the second time: every file is compared (sha256) before it is
# replaced (atomically, via a temporary file and rename), systemd is reloaded and services restarted only when something they use changed,
# and the last line says what changed ("no changes" otherwise).
#
#   install.sh server --table /path/ple-fp8.safetensors [--port 18460] [--allow "192.0.2.10 192.0.2.11"]
#   install.sh agent  [--server HOST:18460] [--table /path/ple-fp8.safetensors] [--rows 16777216]
#   install.sh plugin --python /path/to/venv/bin/python
# From PyPI, without the repository: `uv tool install service-ple` (or pipx), then `sudo "$(command -v service-ple-install)" agent ...`.
# Common options: -n (dry run: report what would change), --no-restart (install but leave the restart pending), --build DIR.
# ROOT=/some/dir installs under that root and skips systemctl/nft (for tests).
set -euo pipefail

usage() { sed -n '2,11p' "$0" | sed 's/^# \{0,1\}//'; exit ${1:-2}; }
[ $# -ge 1 ] || usage
if [ "${1:-}" = --help ] || [ "${1:-}" = -h ]; then usage 0; fi
MODE=$1; shift
SRC=$(cd "$(dirname "$0")/.." && pwd)
# installed from PyPI (`pip install service-ple`, this file as bin/service-ple-install): the binaries are already built, next to it
PREBUILT=""; [ -x "$(dirname "$0")/service-ple-agent" ] && [ ! -f "$SRC/CMakeLists.txt" ] && PREBUILT=$(cd "$(dirname "$0")" && pwd)
ROOT=${ROOT:-}
DRY=0 RESTART=1 BUILD=${PREBUILT:-$SRC/build}
TABLE="" PORT=18460 ALLOW="" SERVER="" ROWS=16777216 PY=""
while [ $# -gt 0 ]; do
  case $1 in
    -n) DRY=1 ;;
    --no-restart) RESTART=0 ;;
    --build) BUILD=$2; shift ;;
    --table) TABLE=$2; shift ;;
    --port) PORT=$2; shift ;;
    --allow) ALLOW=$2; shift ;;
    --server) SERVER=$2; shift ;;
    --rows) ROWS=$2; shift ;;
    --python) PY=$2; shift ;;
    *) usage ;;
  esac
  shift
done

[[ $PORT =~ ^[0-9]+$ ]] && [ "$PORT" -ge 1 ] && [ "$PORT" -le 65535 ] || { echo "--port: not a port number: $PORT" >&2; exit 2; }
[[ $ROWS =~ ^[0-9]+$ ]] && [ "$ROWS" -ge 1 ] || { echo "--rows: not a positive number: $ROWS" >&2; exit 2; }

CHANGES=()
note() { CHANGES+=("$1"); }
sys() { if [ -z "$ROOT" ] && [ $DRY = 0 ]; then "$@"; fi; }   # real system actions only (a failure stops the installer)

# put_file SRC DST MODE: install only if the content differs; returns 0 if it changed (or would change)
put_file() {
  local src=$1 dst=$ROOT$2 mode=$3
  if [ -f "$dst" ] && cmp -s "$src" "$dst"; then return 1; fi
  if [ $DRY = 0 ]; then
    mkdir -p "$(dirname "$dst")"
    # explicit exit: put_file is called in && / || lists, where bash ignores set -e (a failed install would pass as "changed")
    install -m "$mode" "$src" "$dst.tmp.$$" && mv -f "$dst.tmp.$$" "$dst" || { echo "ERROR: cannot install $dst" >&2; exit 1; }
  fi
  note "$2"
  return 0
}

# put_text DST MODE <<<content: same, for generated files
put_text() {
  local t; t=$(mktemp); cat > "$t"
  local r=1; put_file "$t" "$1" "$2" && r=0
  rm -f "$t"; return $r
}

build() {
  [ -n "$PREBUILT" ] && return 0
  cmake -S "$SRC" -B "$BUILD" -DCMAKE_BUILD_TYPE=RelWithDebInfo >/dev/null
  cmake --build "$BUILD" -j"$(nproc)" >/dev/null
}

# enable + (re)start a unit only when needed
apply_unit() {
  local unit=$1 unit_changed=$2 needs_restart=$3
  [ "$unit_changed" = 1 ] && sys systemctl daemon-reload
  if [ -z "$ROOT" ] && ! systemctl is-enabled -q "$unit" 2>/dev/null; then
    [ $DRY = 0 ] && systemctl enable -q "$unit"; note "enable $unit"
  fi
  if [ -z "$ROOT" ] && ! systemctl is-active -q "$unit"; then
    [ $DRY = 0 ] && systemctl start "$unit"; note "start $unit"
  elif [ "$needs_restart" = 1 ]; then
    if [ $RESTART = 1 ]; then sys systemctl restart "$unit"; note "restart $unit"; else note "RESTART PENDING: $unit"; fi
  fi
}

case $MODE in
  server)
    [ -n "$TABLE" ] || { echo "server: --table is required" >&2; exit 2; }
    build
    bin=0 unit=0 fw=0
    put_file "$BUILD/service-ple-server" /usr/local/bin/service-ple-server 755 && bin=1
    if [ -n "$ALLOW" ]; then
      # one nft transaction: the old table is replaced with no moment without rules ("table" first so the delete works when it is absent)
      put_text /etc/service-ple/server.nft 644 <<E && fw=1
# generated by install.sh: port $PORT only from the listed hosts
table inet service_ple
delete table inet service_ple
table inet service_ple {
	chain input {
		type filter hook input priority filter; policy accept;
		tcp dport $PORT iif "lo" accept
		tcp dport $PORT ip saddr { $(echo $ALLOW | sed 's/ /, /g') } accept
		tcp dport $PORT counter drop
	}
}
E
      NFT_PRE="ExecStartPre=/usr/sbin/nft -f /etc/service-ple/server.nft"
    else
      NFT_PRE=""
    fi
    # 1 GiB pages for the whole table, or none: a partial reservation would stay pinned and unused while the server falls back to THP
    [ -f "$TABLE" ] || { echo "server: $TABLE does not exist" >&2; exit 2; }
    PAGES=$(( ($(stat -L -c %s "$TABLE") >> 30) + 1 ))
    put_text /etc/systemd/system/service-ple.service 644 <<E && unit=1
# generated by install.sh
[Unit]
Description=service-ple server: PLE table in RAM, served over TCP
After=network-online.target local-fs.target
Wants=network-online.target

[Service]
# \$\$ = a literal \$ for the shell (systemd would expand \$H itself)
Environment=H=/sys/kernel/mm/hugepages/hugepages-1048576kB/nr_hugepages
$NFT_PRE
ExecStartPre=-/bin/sh -c "echo 1 > /proc/sys/vm/compact_memory; echo $PAGES > \$\$H; [ \$\$(cat \$\$H) -ge $PAGES ] || echo 0 > \$\$H"
ExecStart=/usr/local/bin/service-ple-server --fichero $TABLE --puerto $PORT --mlock
ExecStopPost=-/bin/sh -c "echo 0 > \$\$H"
LimitMEMLOCK=infinity
Restart=on-failure
RestartSec=30

[Install]
WantedBy=multi-user.target
E
    apply_unit service-ple.service $unit $(( bin | unit | fw ))
    ;;
  agent)
    [ -n "$SERVER$TABLE" ] || { echo "agent: --server, --table or both" >&2; exit 2; }
    build
    bin=0 unit=0
    put_file "$BUILD/service-ple-agent" /usr/local/bin/service-ple-agent 755 && bin=1
    # the client library for C++ engines (vLLM gets its own copy from the vllm-service-ple package); not in the PyPI package
    if [ -z "$PREBUILT" ]; then
      [ -f "$BUILD/libserviceple.so" ] || { echo "agent: $BUILD/libserviceple.so was not built" >&2; exit 1; }
      put_file "$BUILD/libserviceple.so" /usr/local/lib/libserviceple.so 644 || true
    fi
    args="--socket /run/service-ple/agente.sock --filas $ROWS"
    [ -n "$SERVER" ] && args="--servidor $SERVER $args"
    [ -n "$TABLE" ] && args="$args --fichero $TABLE"
    MOUNTS=""; [ -n "$TABLE" ] && MOUNTS="RequiresMountsFor=$(dirname "$TABLE")"   # the table's file system mounted first (e.g. a nofail entry)
    put_text /etc/systemd/system/service-ple-agent.service 644 <<E && unit=1
# generated by install.sh
[Unit]
Description=service-ple agent: shared PLE row cache for the engines of this machine
After=network-online.target
Wants=network-online.target
$MOUNTS

[Service]
ExecStart=/usr/local/bin/service-ple-agent $args
Restart=on-failure
RestartSec=10

[Install]
WantedBy=multi-user.target
E
    apply_unit service-ple-agent.service $unit $(( bin | unit ))
    ;;
  plugin)
    [ -n "$PY" ] || { echo "plugin: --python is required" >&2; exit 2; }
    [ -z "$PREBUILT" ] || { echo "plugin: from PyPI, run: $PY -m pip install vllm-service-ple (or uv pip install --python $PY vllm-service-ple)" >&2; exit 2; }
    build
    site=$("$PY" -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')
    # the package (with the library inside) is installed only if its content hash differs from the installed one
    stage=$(mktemp -d); trap 'rm -rf "$stage"' EXIT
    cp -r "$SRC/clients/vllm/." "$stage/"; rm -rf "$stage"/vllm_service_ple/__pycache__
    cp "$BUILD/libserviceple.so" "$stage/vllm_service_ple/"
    want=$(cd "$stage" && find . -type f -print0 | sort -z | xargs -0 sha256sum | sha256sum | cut -c1-16)
    have=$(cat "$ROOT$site/vllm_service_ple/.content-hash" 2>/dev/null || true)
    if [ "$want" != "$have" ]; then
      if [ $DRY = 0 ]; then
        if "$PY" -m pip --version >/dev/null 2>&1; then
          "$PY" -m pip install -q --no-deps --no-build-isolation ${ROOT:+--root "$ROOT"} "$stage"
        elif command -v uv >/dev/null && [ -z "$ROOT" ]; then   # venv made with uv: no pip inside
          uv pip install -q --python "$PY" --no-deps --reinstall "$stage"
        else
          echo "plugin: $PY has no pip (and uv is missing, or ROOT is set)" >&2; exit 1
        fi
        echo "$want" > "$ROOT$site/vllm_service_ple/.content-hash"   # marker of what is installed (pip does not copy it)
      fi
      note "plugin vllm-service-ple in $site ($want)"
    fi
    ;;
  *) usage ;;
esac

if [ ${#CHANGES[@]} = 0 ]; then echo "service-ple install ($MODE): no changes"
else
  [ $DRY = 1 ] && pre="would change" || pre="changed"
  echo "service-ple install ($MODE): $pre: ${CHANGES[*]}"
fi
