#!/usr/bin/env bash
# ling3-run.sh — one-command server launcher for bloomer010 Ling-3.0 GGUFs
#
# Picks a quant to fit available memory, applies the source model's recommended
# sampling, and optionally enables the bundled MTP drafter (flash only).
#
# Requires a llama.cpp build with bailingmoe3 support (upstream PR #26608 or
# newer; older builds: https://github.com/aetherbird/llama.cpp/tree/bailingmoe3-support).
#
# Usage:
#   ./ling3-run.sh flash                    # auto quant, auto memory detection
#   ./ling3-run.sh flash Q4_K_S             # explicit quant
#   ./ling3-run.sh flash --mtp              # force MTP drafter on
#   ./ling3-run.sh tiny
#   ./ling3-run.sh flash --ctx 65536 --port 8081
#
# Env overrides: LLAMA_BIN_DIR (path to llama-server), VRAM_GB / RAM_GB (skip detection)

set -euo pipefail

MODEL="${1:-flash}"; shift || true
QUANT="auto"; MTP="auto"; CTX="32768"; PORT="8080"
while [ $# -gt 0 ]; do
  case "$1" in
    --mtp) MTP="on"; shift ;;
    --no-mtp) MTP="off"; shift ;;
    --ctx) CTX="$2"; shift 2 ;;
    --port) PORT="$2"; shift 2 ;;
    *) QUANT="$1"; shift ;;
  esac
done

REPO="bloomer010/Ling-3.0-${MODEL}-GGUF"
case "$MODEL" in
  flash)
    # ladder: (min usable total memory GB -> quant), weights + context headroom
    LADDER=(
      "192 UD-Q8_K_XL"
      "136 Q8_0"
      "124 UD-Q6_K_XL"
      "80  Q5_K_M"
      "64  Q4_K_M"
      "56  MXFP4_MOE"   # Blackwell/GB10 native; others dequant fallback
      "56  Q4_K_S"
      "48  Q3_K_M"
      "32  UD-Q2_K_XL"
      "24  IQ1_M"
    )
    TEMP="0.6"   # source model card recommendation
    ;;
  tiny)
    LADDER=(
      "16 BF16"
      "12 UD-Q8_K_XL"
      "10 Q8_0"
      "8  UD-Q6_K_XL"
      "6  Q4_K_M"
      "5  MXFP4_MOE"
      "4  Q3_K_M"
      "3  IQ2_M"
    )
    TEMP="1.0"   # source model card recommendation
    ;;
  *) echo "unknown model '$MODEL' (flash|tiny)" >&2; exit 1 ;;
esac

# --- memory detection -------------------------------------------------------
detect_mem() {
  if [ -n "${VRAM_GB:-}" ] && [ -n "${RAM_GB:-}" ]; then
    echo $(( VRAM_GB > RAM_GB ? VRAM_GB : RAM_GB ))
    return
  fi
  local vram ram
  vram=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits 2>/dev/null \
         | awk '{s+=$1} END {printf "%d", s/1024}' || echo 0)
  ram=$(free -g | awk '/^Mem:/ {printf "%d", $2}')
  echo $(( vram > ram ? vram : ram ))
}

TOTAL_MEM=$(detect_mem)
# reserve headroom for KV cache, runtime, and OS
TOTAL_MEM=$((TOTAL_MEM > 10 ? TOTAL_MEM - 8 : TOTAL_MEM))

if [ "$QUANT" = "auto" ]; then
  for rung in "${LADDER[@]}"; do
    read -r need q <<< "$rung"
    if [ "$TOTAL_MEM" -ge "$need" ]; then QUANT="$q"; break; fi
  done
  [ "$QUANT" = "auto" ] && QUANT="$(read -r _ q <<< "${LADDER[-1]}"; echo "$q")"
  echo "[ling3] ${TOTAL_MEM} GB detected -> ${QUANT} (override: ./ling3-run.sh ${MODEL} <quant>)"
fi

# --- binary -----------------------------------------------------------------
BIN="${LLAMA_BIN_DIR:-}/llama-server"
command -v "$BIN" >/dev/null 2>&1 || BIN="llama-server"
command -v "$BIN" >/dev/null 2>&1 || { echo "llama-server not found (set LLAMA_BIN_DIR)" >&2; exit 1; }

# --- MTP: on for flash unless disabled; never for tiny (no bundled MTP) ----
SPEC_ARGS=()
if [ "$MODEL" = "flash" ] && { [ "$MTP" = "on" ] || [ "$MTP" = "auto" ]; }; then
  SPEC_ARGS=(--spec-type draft-mtp)
fi

exec "$BIN" \
  -hf "${REPO}:${QUANT}" \
  --host 0.0.0.0 --port "$PORT" \
  -c "$CTX" -ngl auto -fa on --jinja \
  --temp "$TEMP" --top-p 0.95 --top-k 20 \
  "${SPEC_ARGS[@]}" "$@"