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

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
MODEL="${1:-qwen38}"
case "$MODEL" in
  qwen38) MODEL_PATH="$ROOT/models/Qwen3.8-27B-Q6_K.gguf"; EXPECTED_SIZE=22563871392 ;;
  qwen36) MODEL_PATH="$ROOT/models/Qwen3.6-35B-A3B-UD-Q5_K_M.gguf"; EXPECTED_SIZE=26456194016 ;;
  *) echo "Usage: $0 [qwen38|qwen36]" >&2; exit 2 ;;
esac
[[ -f "$MODEL_PATH" && "$(stat -c %s "$MODEL_PATH")" == "$EXPECTED_SIZE" ]] || {
  echo "Model missing or incomplete: $MODEL_PATH" >&2
  exit 1
}

BIN="$ROOT/runtime/b11177/llama-server"
[[ -x "$BIN" ]] || BIN="$ROOT/runtime/b11177/build/bin/llama-server"
[[ -x "$BIN" ]] || { echo 'llama-server missing. Run scripts/install.sh.' >&2; exit 1; }

mkdir -p "$ROOT/secrets"
KEY_FILE="$ROOT/secrets/api_keys.txt"
if [[ ! -f "$KEY_FILE" ]]; then
  umask 077
  python -c 'import secrets; print(secrets.token_urlsafe(32))' > "$KEY_FILE"
fi

export LD_LIBRARY_PATH="$ROOT/runtime/b11177:/opt/conda/lib:${LD_LIBRARY_PATH:-}"
export OMP_NUM_THREADS=4
export OMP_THREAD_LIMIT=4
export OMP_DYNAMIC=FALSE
export MKL_NUM_THREADS=1
export OPENBLAS_NUM_THREADS=1
export NUMEXPR_NUM_THREADS=1
export BLIS_NUM_THREADS=1
export VECLIB_MAXIMUM_THREADS=1
CPU_SET="$(python -c 'import os; print(",".join(map(str, sorted(os.sched_getaffinity(0))[:32])))')"
SPEC_ARGS=()
if [[ "${LLM_SPEC_MTP:-0}" == 1 ]]; then
  [[ "$MODEL" == qwen38 ]] || { echo 'MTP draft is only configured for qwen38.' >&2; exit 1; }
  DRAFT="$ROOT/models/mtp-Qwen3.8-27B-Q8_0.gguf"
  [[ -f "$DRAFT" && "$(stat -c %s "$DRAFT")" == 3164006816 ]] || {
    echo 'MTP draft missing. Run WITH_MTP=1 bash scripts/install.sh.' >&2
    exit 1
  }
  SPEC_ARGS=(--spec-type draft-mtp --model-draft "$DRAFT" --spec-draft-n-max 2)
fi
exec taskset -c "$CPU_SET" "$BIN" \
  --model "$MODEL_PATH" \
  --alias "$MODEL" \
  --host "${LLM_HOST:-127.0.0.1}" \
  --port "${LLM_PORT:-8080}" \
  --api-key-file "$KEY_FILE" \
  --n-gpu-layers all \
  --fit off \
  --flash-attn on \
  --cache-type-k q8_0 \
  --cache-type-v q8_0 \
  --ctx-size "${LLM_CONTEXT:-65536}" \
  --parallel 4 \
  --threads 16 \
  --threads-batch 16 \
  --threads-http 4 \
  --cont-batching \
  --reasoning-effort "${LLM_REASONING_EFFORT:-medium}" \
  --batch-size 2048 \
  --ubatch-size 512 \
  --jinja \
  "${SPEC_ARGS[@]}"