#!/usr/bin/env bash
# ===========================================================================
# SearchAI Inference Server — installer (Linux systemd).
# Installs the single self-contained prebuilt server binary from the artifact
# bucket, fetches models (4B text+vision by default; add asr/tts/voice), writes
# config + systemd, and starts the service. The SIMD kernel is embedded in the
# binary — there is NO companion .so and NO JVM/-D flags.
#
#   sudo API_KEY='...' ./install-searchai.sh                              # single node, 4B+vision
#   sudo API_KEY='...' MODELS='4b asr tts voice' ./install-searchai.sh    # + media
#   sudo API_KEY='...' BACKEND=auto ./install-searchai.sh                 # NVIDIA GPU box: use GPU when present
#   sudo API_KEY='...' BACKEND=cuda ./install-searchai.sh                 # require the GPU (refuse CPU fallback)
#   sudo API_KEY='...' CLUSTER_ADVERTISE='10.0.0.10:8081' CLUSTER_SECRET='...' ./install-searchai.sh   # seed
#   sudo API_KEY='...' CLUSTER_ADVERTISE='10.0.0.11:8081' CLUSTER_SEED='10.0.0.10:8081' \
#        CLUSTER_SECRET='...' ./install-searchai.sh                       # joiner
# ===========================================================================
set -Eeuo pipefail
[[ "${EUID}" -eq 0 ]] || { echo "Run as root." >&2; exit 1; }

MODEL_ID="${MODEL_ID:-q35-4b}"
MODELS="${MODELS:-4b}"                 # fetch-models.sh groups: 4b asr tts voice 2b 0.8b 9b all
API_KEY="${API_KEY:-}"
# Artifact source: public bucket (HTTPS, no AWS credentials needed).
ARTIFACT_BUCKET="${ARTIFACT_BUCKET:-s3://searchai-inference-server}"
# Public HTTPS base for the curl fallback (CDN-cached; internet downloads).
# The aws-CLI path above stays S3-direct, which is free from in-region EC2.
PUB_BASE="${ARTIFACT_URL:-https://inference-server.searchblox.com}"
# fetch_artifact <key-under-bucket> <dest>: aws CLI when present, else public HTTPS
fetch_artifact() {
  local key="$1" dest="$2"
  if command -v aws >/dev/null 2>&1 && aws s3 cp "${ARTIFACT_BUCKET}/${key}" "$dest" --only-show-errors 2>/dev/null; then
    return 0
  fi
  curl -fL --retry 3 -o "$dest" "${PUB_BASE}/${key}"
}
SERVER_BINARY="${SERVER_BINARY:-}"
SERVER_HOST="${SERVER_HOST:-0.0.0.0}"; SERVER_PORT="${SERVER_PORT:-8081}"
KV_CACHE="${KV_CACHE:-q8}"; MAX_INFLIGHT="${MAX_INFLIGHT:-4}"
# min-free-mb default scales with RAM: 4096 was 25% of a 16GB box and shed
# load spuriously right after each completed request; ~1/8 of RAM capped at
# 4096 keeps the OOM guard without starving small boxes.
if [[ -z "${MIN_FREE_MB:-}" ]]; then
  TOTAL_RAM_MB="$(awk '/MemTotal:/ {print int($2/1024)}' /proc/meminfo 2>/dev/null || echo 32768)"
  MIN_FREE_MB=$(( TOTAL_RAM_MB / 8 )); (( MIN_FREE_MB > 4096 )) && MIN_FREE_MB=4096
fi
MAX_LOADED_MODELS="${MAX_LOADED_MODELS:-0}"
# BACKEND: cpu (default) | auto (use an NVIDIA GPU when present, else CPU) |
# cuda (require the GPU; startup fails without it). auto/cuda additionally
# install the GPU acceleration add-on from build/linux-amd64-cuda/.
BACKEND="${BACKEND:-cpu}"
MAX_CONTEXT="${MAX_CONTEXT:-8192}"; MAX_TOKENS="${MAX_TOKENS:-6144}"; MAX_BODY_MB="${MAX_BODY_MB:-64}"
MIN_RAM_GB="${MIN_RAM_GB:-16}"   # lower (e.g. 8) for small test boxes; the startup gate refuses below this
SAI_THREADS_VALUE="${SAI_THREADS_VALUE:-}"
CLUSTER_ADVERTISE="${CLUSTER_ADVERTISE:-}"; CLUSTER_SEED="${CLUSTER_SEED:-}"
CLUSTER_PEERS="${CLUSTER_PEERS:-}"; CLUSTER_SECRET="${CLUSTER_SECRET:-}"; CLUSTER_PREFIX="${CLUSTER_PREFIX:-true}"

INSTALL_DIR=/opt/searchai; CONFIG_DIR=/etc/searchai
DATA_DIR=/var/lib/searchai; MODELS_DIR="${DATA_DIR}/models"
SERVICE_FILE=/etc/systemd/system/searchai.service
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

[[ "$MODEL_ID" =~ ^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$ && "$MODEL_ID" != *".."* ]] || { echo "Invalid MODEL_ID" >&2; exit 1; }
for v in "$API_KEY" "$CLUSTER_SECRET" "$CLUSTER_ADVERTISE" "$CLUSTER_SEED"; do
  [[ "$v" == *$'\n'* || "$v" == *$'\r'* ]] && { echo "Config values cannot contain newlines." >&2; exit 1; }
done
[[ -n "$CLUSTER_ADVERTISE" && -z "$CLUSTER_SECRET" ]] && { echo "CLUSTER_SECRET required for clusters." >&2; exit 1; }
[[ -z "$API_KEY" ]] && API_KEY="$(od -An -N32 -tx1 /dev/urandom | tr -d ' \n')"
case "$(uname -m)" in
  aarch64) PLAT=linux-arm64 ;;
  x86_64)  PLAT=linux-amd64 ;;   # one x86_64 binary runs on Intel (AMX) + AMD (VNNI) via runtime CPUID dispatch
  *) echo "Unsupported arch $(uname -m)" >&2; exit 1 ;;
esac

# runtime deps: LLVM libomp (the binary links __kmpc_* — NOT gcc libgomp),
# ca-certs, unzip, and ffmpeg (video). ffmpeg is optional (|| true); libomp required.
if command -v apt-get >/dev/null 2>&1; then
  apt-get update -qq
  DEBIAN_FRONTEND=noninteractive apt-get install -y ca-certificates curl unzip ffmpeg >/dev/null || true
  DEBIAN_FRONTEND=noninteractive apt-get install -y libomp5 || apt-get install -y libomp-dev || apt-get install -y libomp5-18 || apt-get install -y libomp5-17 || true
elif command -v dnf >/dev/null 2>&1; then
  dnf install -y ca-certificates curl unzip ffmpeg >/dev/null || true
  dnf install -y libomp || dnf install -y libomp-devel || true
fi

# Artifacts and models download over public HTTPS (curl) — no cloud CLI or
# credentials required. When an aws CLI with working credentials is already
# present (e.g. an EC2 role), it is used instead for free in-region transfer.
# Verify the LLVM OpenMP runtime is present (libomp.so.5 / libomp.so); the server
# will not start without it.
if ! ldconfig -p 2>/dev/null | grep -q 'libomp\.so'; then
  echo "WARNING: LLVM libomp not found via ldconfig — install it (e.g. 'apt-get install libomp5' or 'dnf install libomp') or the server will fail to start." >&2
fi

id searchai >/dev/null 2>&1 || useradd --system --home-dir "$DATA_DIR" --shell /usr/sbin/nologin searchai
install -d -o root -g root -m 0755 "$INSTALL_DIR/bin"
install -d -o root -g searchai -m 0750 "$CONFIG_DIR"
install -d -o searchai -g searchai -m 0750 "$DATA_DIR" "$MODELS_DIR"

# --- binary (prebuilt from the artifact bucket unless SERVER_BINARY given) ---
if [[ -z "$SERVER_BINARY" ]]; then
  SERVER_BINARY="/tmp/searchai-server.${PLAT}"
  echo "Downloading server (${PLAT}) ..."
  fetch_artifact "build/${PLAT}/searchai-server" "$SERVER_BINARY"
  if fetch_artifact "build/${PLAT}/searchai-server.sha256" /tmp/ss.sha 2>/dev/null; then
    exp="$(awk '{print $1}' /tmp/ss.sha)"; act="$(sha256sum "$SERVER_BINARY" | awk '{print $1}')"
    [[ "$exp" == "$act" ]] || { echo "binary sha256 mismatch" >&2; exit 1; }; echo "binary sha256 OK"
  fi
  chmod +x "$SERVER_BINARY"
fi
[[ -x "$SERVER_BINARY" ]] || { echo "binary missing: $SERVER_BINARY" >&2; exit 1; }
install -o root -g root -m 0755 "$SERVER_BINARY" "${INSTALL_DIR}/bin/searchai-server"

# --- GPU acceleration add-on (BACKEND=auto|cuda) -----------------------------
GPU_ADAPTER_PATH=""
if [[ "$BACKEND" != "cpu" ]]; then
  case "$BACKEND" in auto|cuda) ;; *) echo "Invalid BACKEND='$BACKEND' (cpu|auto|cuda)" >&2; exit 1;; esac
  [[ "$PLAT" == "linux-amd64" ]] || { echo "BACKEND=$BACKEND requires an x86_64 host (this is $PLAT)." >&2; exit 1; }
  # Driver check: the NVIDIA driver is a host prerequisite (kernel module +
  # possible reboot) and is NOT installed by this script.
  if ! command -v nvidia-smi >/dev/null 2>&1 || ! nvidia-smi >/dev/null 2>&1; then
    if [[ "$BACKEND" == "cuda" ]]; then
      echo "BACKEND=cuda but no working NVIDIA driver (nvidia-smi failed). Install the driver first." >&2; exit 1
    fi
    echo "WARNING: no working NVIDIA driver detected — BACKEND=auto will serve on CPU until one is installed." >&2
  fi
  echo "Downloading GPU acceleration add-on ..."
  install -d -o root -g root -m 0755 "${INSTALL_DIR}/lib/gpu"
  fetch_artifact "build/linux-amd64-cuda/searchai-gpu-adapter.tgz" /tmp/sgpu.tgz
  if fetch_artifact "build/linux-amd64-cuda/searchai-gpu-adapter.tgz.sha256" /tmp/sgpu.sha 2>/dev/null; then
    exp="$(awk '{print $1}' /tmp/sgpu.sha)"; act="$(sha256sum /tmp/sgpu.tgz | awk '{print $1}')"
    [[ "$exp" == "$act" ]] || { echo "GPU add-on sha256 mismatch" >&2; exit 1; }; echo "GPU add-on sha256 OK"
  fi
  tar xzf /tmp/sgpu.tgz -C "${INSTALL_DIR}/lib/gpu"
  GPU_ADAPTER_PATH="${INSTALL_DIR}/lib/gpu/libsearchai_llamacpp.so"
  [[ -f "$GPU_ADAPTER_PATH" ]] || { echo "GPU add-on missing adapter library after extract" >&2; exit 1; }
fi

# --- image-editing add-on (opt-in: MODELS contains "image") ------------------
# Instruction image editing (Qwen-Image-Edit class). GPU strongly recommended:
# a CPU-only edit takes minutes. Weights stream from host RAM by default so a
# 24 GB GPU fits the 20B editor.
IMG_ADAPTER_PATH=""
if [[ " ${MODELS} " == *" image "* ]]; then
  [[ "$PLAT" == "linux-amd64" ]] || { echo "image editing requires an x86_64 host (this is $PLAT)." >&2; exit 1; }
  echo "Downloading image-editing add-on ..."
  install -d -o root -g root -m 0755 "${INSTALL_DIR}/lib/img"
  fetch_artifact "build/linux-amd64-cuda/searchai-img-adapter.tgz" /tmp/simg.tgz
  if fetch_artifact "build/linux-amd64-cuda/searchai-img-adapter.tgz.sha256" /tmp/simg.sha 2>/dev/null; then
    exp="$(awk '{print $1}' /tmp/simg.sha)"; act="$(sha256sum /tmp/simg.tgz | awk '{print $1}')"
    [[ "$exp" == "$act" ]] || { echo "image add-on sha256 mismatch" >&2; exit 1; }; echo "image add-on sha256 OK"
  fi
  tar xzf /tmp/simg.tgz -C "${INSTALL_DIR}/lib/img"
  IMG_ADAPTER_PATH="${INSTALL_DIR}/lib/img/libsearchai_sdcpp.so"
  [[ -f "$IMG_ADAPTER_PATH" ]] || { echo "image add-on missing adapter library after extract" >&2; exit 1; }
fi

# --- models (via fetch-models.sh from the shared bucket) ---------------------
echo "Fetching models: ${MODELS} -> ${MODELS_DIR}"
MODELS_DIR="$MODELS_DIR" bash "${HERE}/fetch-models.sh" --dest "$MODELS_DIR" $MODELS
chown -R searchai:searchai "$MODELS_DIR"
[[ -f "${MODELS_DIR}/${MODEL_ID}.gguf" ]] || { echo "default model ${MODEL_ID}.gguf missing after fetch" >&2; exit 1; }

# enable media in config only when the model files are present
have() { [[ -f "${MODELS_DIR}/$1" ]]; }
VISION=$(have "${MODEL_ID}-mmproj.gguf" && echo true || echo false)

# --- config -----------------------------------------------------------------
CFG="${CONFIG_DIR}/server.properties"
{
  echo "server.host=${SERVER_HOST}"
  echo "server.port=${SERVER_PORT}"
  echo "server.api-key=${API_KEY}"
  echo
  echo "backend=${BACKEND}"
  [[ -n "$GPU_ADAPTER_PATH" ]] && echo "gpu.adapter-path=${GPU_ADAPTER_PATH}"
  echo "model=${MODEL_ID}"
  echo "models-dir=${MODELS_DIR}"
  echo "kv-cache=${KV_CACHE}"
  echo "max-context=${MAX_CONTEXT}"
  echo "max-tokens=${MAX_TOKENS}"
  echo "max-inflight=${MAX_INFLIGHT}"
  echo "queue-timeout-ms=120000"
  echo "min-free-mb=${MIN_FREE_MB}"
  echo "min-ram-gb=${MIN_RAM_GB}"
  echo "max-loaded-models=${MAX_LOADED_MODELS}"
  echo "max-connections=512"
  echo "read-timeout-s=30"
  echo "max-body-mb=${MAX_BODY_MB}"
  echo "tuning.batch=true"
  echo "enable-vision=${VISION}"
  have "${MODEL_ID}-mmproj.gguf" && echo "mmproj-path=${MODELS_DIR}/${MODEL_ID}-mmproj.gguf"
  if command -v ffmpeg >/dev/null 2>&1; then echo "video-ffmpeg=$(command -v ffmpeg)"; echo "video-fps=1"; echo "video-max-frames=16"; fi
  have "qwen3-asr-1.7b.gguf"  && { echo "asr-model=qwen3-asr-1.7b"; echo "asr-mmproj=mmproj-qwen3-asr.gguf"; }
  have "qwen-talker-1.7b.gguf" && { echo "tts-talker=qwen-talker-1.7b.gguf"; echo "tts-codec=qwen-tts-codec.gguf"; }
  if [[ -n "$IMG_ADAPTER_PATH" ]] && have "qwen-image-edit-2511-Q4_K_M.gguf"; then
    echo "image-adapter-path=${IMG_ADAPTER_PATH}"
    echo "image-diffusion=qwen-image-edit-2511-Q4_K_M.gguf"
    echo "image-llm=qwen_2.5_vl_7b.safetensors"
    echo "image-vae=qwen_image_vae.safetensors"
  fi
  if [[ -n "$CLUSTER_ADVERTISE" ]]; then
    echo "cluster.advertise=${CLUSTER_ADVERTISE}"; echo "cluster.secret=${CLUSTER_SECRET}"; echo "cluster.prefix=${CLUSTER_PREFIX}"
  fi
  [[ -n "$CLUSTER_SEED"  ]] && echo "cluster.seed=${CLUSTER_SEED}"
  [[ -n "$CLUSTER_PEERS" ]] && echo "cluster.peers=${CLUSTER_PEERS}"
} > "$CFG"
chown root:searchai "$CFG"; chmod 0640 "$CFG"

{
  echo "SAI_FAST_PREFILL=1"; echo "SAI_ROUTE_LOG=0"
  # Fast f16/q8 attention kernels (NEON/AVX2 + runtime CPU gate). Float order
  # differs from the scalar path by 1 ulp; 380-prompt pack verified
  # byte-identical outputs (2026-08-12 release gate).
  echo "SAI_FAST_ATTN=1"
  # Scale the idle request-state pool to host RAM: on <=32 GB hosts a full
  # 4 GB pool can crowd out on-demand model loads (audio/image models fail
  # to (re)load once text serving reaches steady state).
  if   [[ "$TOTAL_RAM_MB" -le 34000 ]]; then echo "SAI_POOL_MAX_MB=1024"
  elif [[ "$TOTAL_RAM_MB" -le 68000 ]]; then echo "SAI_POOL_MAX_MB=2048"
  fi
  [[ -n "$SAI_THREADS_VALUE" ]] && echo "SAI_THREADS=${SAI_THREADS_VALUE}"
  ld_paths=""
  [[ -n "$GPU_ADAPTER_PATH" ]] && ld_paths="${INSTALL_DIR}/lib/gpu"
  [[ -n "$IMG_ADAPTER_PATH" ]] && ld_paths="${ld_paths:+${ld_paths}:}${INSTALL_DIR}/lib/img"
  [[ -n "$ld_paths" ]] && echo "LD_LIBRARY_PATH=${ld_paths}"
} > "${CONFIG_DIR}/searchai.env"
chown root:searchai "${CONFIG_DIR}/searchai.env"; chmod 0640 "${CONFIG_DIR}/searchai.env"

cat > "${INSTALL_DIR}/bin/preload-model" <<'PRELOAD'
#!/usr/bin/env bash
set -Eeuo pipefail
CFG=/etc/searchai/server.properties
port="$(sed -n 's/^server.port=//p' "$CFG" | tail -1)"; key="$(sed -n 's/^server.api-key=//p' "$CFG" | tail -1)"; m="$(sed -n 's/^model=//p' "$CFG" | tail -1)"
for _ in $(seq 1 180); do curl -fsS "http://127.0.0.1:${port}/health" >/dev/null && break; sleep 1; done
curl -fsS -H "Authorization: Bearer ${key}" -H "Content-Type: application/json" -d "{\"model\":\"${m}\"}" "http://127.0.0.1:${port}/v1/models/load"
PRELOAD
chmod 0755 "${INSTALL_DIR}/bin/preload-model"

cat > "$SERVICE_FILE" <<'SERVICE'
[Unit]
Description=SearchAI Inference Server
After=network-online.target
Wants=network-online.target

[Service]
Type=simple
User=searchai
Group=searchai
WorkingDirectory=/var/lib/searchai
EnvironmentFile=/etc/searchai/searchai.env
# Allocator tuning: return freed large buffers (model-load spikes, drained
# request-state pools) to the OS promptly — keeps resident memory tracking
# actual usage on smaller hosts instead of parking at the high-water mark.
Environment=MALLOC_ARENA_MAX=2
Environment=MALLOC_MMAP_THRESHOLD_=131072
Environment=MALLOC_TRIM_THRESHOLD_=8388608
ExecStart=/opt/searchai/bin/searchai-server /etc/searchai/server.properties
ExecStartPost=/opt/searchai/bin/preload-model
Restart=on-failure
RestartSec=5
TimeoutStartSec=30min
TimeoutStopSec=30
LimitNOFILE=65536
NoNewPrivileges=true
PrivateTmp=true
ProtectHome=true
ProtectSystem=strict
ReadWritePaths=/var/lib/searchai

[Install]
WantedBy=multi-user.target
SERVICE

systemctl daemon-reload
systemctl enable --now searchai
echo
echo "Installed. Endpoint http://127.0.0.1:${SERVER_PORT}  model=${MODEL_ID}  platform=${PLAT}  backend=${BACKEND}  vision=${VISION}"
[[ "$BACKEND" != "cpu" ]] && echo "Verify GPU: curl -s http://127.0.0.1:${SERVER_PORT}/health | grep -o '\"backend\":\"[a-z]*\"' — expect \"cuda\" (auto falls back to cpu with the reason in the 'fallback' field)."
echo "API key: ${API_KEY}"
systemctl --no-pager --full status searchai || true