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

ENV_FILE="/etc/default/model-gateway"

usage() {
  cat <<'EOF'
Usage:
  sudo model-gateway-set-llama-server /path/to/llama-server

Configure model-gateway to use an externally installed llama.cpp llama-server.
NVIDIA driver, CUDA runtime, and llama.cpp are prerequisites and are not
installed by the uni-model-infer deb package.
EOF
}

if [[ "${EUID:-$(id -u)}" -ne 0 ]]; then
  echo "Please run as root: sudo model-gateway-set-llama-server /path/to/llama-server" >&2
  exit 1
fi

if [[ $# -ne 1 || "$1" == "-h" || "$1" == "--help" ]]; then
  usage
  exit $([[ $# -eq 1 && ( "$1" == "-h" || "$1" == "--help" ) ]] && echo 0 || echo 2)
fi

LLAMA_SERVER="$1"
if [[ ! -x "$LLAMA_SERVER" ]]; then
  echo "llama-server is not executable: $LLAMA_SERVER" >&2
  exit 2
fi

if command -v realpath >/dev/null 2>&1; then
  LLAMA_SERVER_ABS="$(realpath -m "$LLAMA_SERVER")"
else
  LLAMA_SERVER_ABS="$(readlink -m "$LLAMA_SERVER")"
fi
BIN_DIR="$(cd "$(dirname "$LLAMA_SERVER_ABS")" && pwd)"

build_lib_path() {
  local bin_dir="$1"
  local -a parts=(
    "$bin_dir"
    "$bin_dir/../lib64"
    "$bin_dir/../lib"
    "$bin_dir/../src"
    "$bin_dir/../ggml/src"
    "$bin_dir/../ggml/src/ggml-cpu"
    "$bin_dir/../ggml/src/ggml-cuda"
    "$bin_dir/../common"
    "$bin_dir/../tools/server"
    "$bin_dir/../../lib64"
    "$bin_dir/../../lib"
    "/opt/model-gateway/lib"
    "/usr/local/lib64"
    "/usr/local/lib"
    "/usr/lib64"
    "/usr/lib/x86_64-linux-gnu"
    "/usr/local/cuda/lib64"
    "/usr/local/cuda/targets/x86_64-linux/lib"
  )

  local result=""
  local item
  for item in "${parts[@]}"; do
    [[ -n "$item" ]] || continue
    result="${result:+$result:}$item"
  done
  printf '%s' "$result"
}

LIB_PATH="$(build_lib_path "$BIN_DIR")"

if command -v ldd >/dev/null 2>&1; then
  if LD_LIBRARY_PATH="$LIB_PATH:${LD_LIBRARY_PATH:-}" ldd "$LLAMA_SERVER_ABS" 2>/dev/null | grep -q "not found"; then
    echo "llama-server has missing shared libraries:" >&2
    LD_LIBRARY_PATH="$LIB_PATH:${LD_LIBRARY_PATH:-}" ldd "$LLAMA_SERVER_ABS" 2>/dev/null | grep "not found" >&2 || true
    echo "Install the matching CUDA/llama.cpp runtime first, then rerun this command." >&2
    exit 2
  fi
fi

escape_env_value() {
  printf '%s' "$1" | sed 's/\\/\\\\/g; s/"/\\"/g'
}

install -d -m 0755 "$(dirname "$ENV_FILE")"
cat > "$ENV_FILE" <<EOF
# Managed by model-gateway-set-llama-server.
MODEL_GATEWAY_LLAMA_SERVER="$(escape_env_value "$LLAMA_SERVER_ABS")"
MODEL_GATEWAY_LLAMA_LIBRARY_PATH="$(escape_env_value "$LIB_PATH")"
EOF

if command -v systemctl >/dev/null 2>&1 && [[ -d /run/systemd/system ]]; then
  systemctl daemon-reload || true
  systemctl restart model-gateway || true
fi

cat <<MSG
Configured model-gateway external llama-server:
  $LLAMA_SERVER_ABS

Check:
  systemctl status model-gateway --no-pager
  journalctl -u model-gateway -n 50 --no-pager
MSG
