#!/usr/bin/env bash
# Start Prism llama-server with Ternary-Bonsai defaults.
set -euo pipefail

MODEL_PATH="${MODEL_PATH:-/models/Ternary-Bonsai-27B-Q2_0.gguf}"
HOST="${HOST:-0.0.0.0}"
PORT="${PORT:-8080}"
CTX_SIZE="${CTX_SIZE:-16384}"
TEMP="${TEMP:-0.7}"
TOP_P="${TOP_P:-0.95}"
TOP_K="${TOP_K:-20}"
NGL="${NGL:-0}"
THREADS="${THREADS:-0}"
KV4="${KV4:-0}"
EXTRA_ARGS="${EXTRA_ARGS:-}"

# Prefer vulkan binary tree if present and NGL > 0
BIN=/usr/local/bin/llama-server
if [ -x /opt/bonsai/bin/vulkan/llama-server ] && [ "${NGL}" != "0" ]; then
  export LD_LIBRARY_PATH="/opt/bonsai/bin/vulkan${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
  BIN=/opt/bonsai/bin/vulkan/llama-server
elif [ -x /opt/bonsai/bin/cpu/llama-server ]; then
  export LD_LIBRARY_PATH="/opt/bonsai/bin/cpu${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
  BIN=/opt/bonsai/bin/cpu/llama-server
fi

if [ ! -f "$MODEL_PATH" ]; then
  echo "ERROR: model not found at $MODEL_PATH"
  echo "Run:  ./scripts/download-bonsai-model.sh"
  echo "Or set MODEL_PATH to your GGUF file."
  exit 1
fi

# Optional vision projector (27B VLM)
MMPROJ_ARGS=()
if [ -n "${MMPROJ_PATH:-}" ] && [ -f "$MMPROJ_PATH" ]; then
  MMPROJ_ARGS=(--mmproj "$MMPROJ_PATH")
else
  # Auto-detect mmproj next to the model
  _dir="$(dirname "$MODEL_PATH")"
  for _mp in "$_dir"/*mmproj*.gguf; do
    if [ -f "$_mp" ]; then
      MMPROJ_ARGS=(--mmproj "$_mp")
      break
    fi
  done
fi

# Optional 4-bit KV cache (long context on limited RAM)
KV_ARGS=()
if [ "$KV4" = "1" ]; then
  KV_ARGS=(--cache-type-k q4_0 --cache-type-v q4_0)
fi

# Thread count: 0 = leave llama.cpp default (usually all physical cores)
THREAD_ARGS=()
if [ "$THREADS" != "0" ] && [ -n "$THREADS" ]; then
  THREAD_ARGS=(-t "$THREADS")
fi

echo "=== Ternary-Bonsai llama-server ==="
echo "  model:   $MODEL_PATH"
echo "  binary:  $BIN"
echo "  ngl:     $NGL  (0=CPU; for Vulkan image try 99 with /dev/dri)"
echo "  context: $CTX_SIZE"
echo "  api:     http://${HOST}:${PORT}/v1/chat/completions"
[ "${#MMPROJ_ARGS[@]}" -gt 0 ] && echo "  vision:  ${MMPROJ_ARGS[1]}"
[ "$KV4" = "1" ] && echo "  kv:      q4_0"
echo ""

# shellcheck disable=SC2086
exec "$BIN" \
  -m "$MODEL_PATH" \
  --host "$HOST" \
  --port "$PORT" \
  -ngl "$NGL" \
  -fa on \
  -c "$CTX_SIZE" \
  --temp "$TEMP" \
  --top-p "$TOP_P" \
  --top-k "$TOP_K" \
  --min-p 0 \
  --jinja \
  "${THREAD_ARGS[@]}" \
  "${MMPROJ_ARGS[@]}" \
  "${KV_ARGS[@]}" \
  $EXTRA_ARGS \
  "$@"