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

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"

export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}"

GGUF_REPO="${GGUF_REPO:-deepreinforce-ai/Ornith-1.0-35B-GGUF}"
QUANT="${QUANT:-Q4_K_M}"
HOST="${HOST:-0.0.0.0}"
PORT="${PORT:-8002}"
CTX_SIZE="${CTX_SIZE:-131072}"
PARALLEL="${PARALLEL:-16}"
BATCH_SIZE="${BATCH_SIZE:-4096}"
UBATCH_SIZE="${UBATCH_SIZE:-512}"
CONT_BATCHING="${CONT_BATCHING:-1}"
METRICS="${METRICS:-1}"
BACKEND_SAMPLING="${BACKEND_SAMPLING:-0}"
CACHE_RAM="${CACHE_RAM:-}"
FLASH_ATTN="${FLASH_ATTN:-on}"
REASONING="${REASONING:-off}"
REASONING_FORMAT="${REASONING_FORMAT:-deepseek}"
REASONING_BUDGET="${REASONING_BUDGET:-}"
REASONING_BUDGET_MESSAGE="${REASONING_BUDGET_MESSAGE:-}"
SPEC_TYPE="${SPEC_TYPE:-}"
SPEC_DRAFT_MODEL="${SPEC_DRAFT_MODEL:-}"
SPEC_DRAFT_HF="${SPEC_DRAFT_HF:-}"
SPEC_DRAFT_NGL="${SPEC_DRAFT_NGL:-99}"
SPEC_DRAFT_N_MAX="${SPEC_DRAFT_N_MAX:-}"
SPEC_DRAFT_N_MIN="${SPEC_DRAFT_N_MIN:-}"
SPEC_DRAFT_P_MIN="${SPEC_DRAFT_P_MIN:-}"
SPEC_DRAFT_DEVICE="${SPEC_DRAFT_DEVICE:-}"
SPEC_DRAFT_BACKEND_SAMPLING="${SPEC_DRAFT_BACKEND_SAMPLING:-0}"
MTP_PER_HEAD_CONTEXT="${MTP_PER_HEAD_CONTEXT:-0}"
ALIAS="${ALIAS:-Ornith-1.0-35B-GGUF-${QUANT}}"
LLAMA_SERVER_BIN="${LLAMA_SERVER_BIN:-${REPO_ROOT}/vendor/llama.cpp/build-cuda/bin/llama-server}"
LLAMACPP_CUDA_HOME="${LLAMACPP_CUDA_HOME:-${REPO_ROOT}/.venv-build/lib/python3.12/site-packages/nvidia/cu13}"
MODEL_PATH="${MODEL_PATH:-}"

if [[ "${MTP_PER_HEAD_CONTEXT}" == "1" ]]; then
  export LLAMA_MTP_PER_HEAD_CONTEXT=1
fi

if [[ -x "${LLAMA_SERVER_BIN}" ]]; then
  export LD_LIBRARY_PATH="${REPO_ROOT}/vendor/llama.cpp/build-cuda/bin:${LLAMACPP_CUDA_HOME}/lib64:${LLAMACPP_CUDA_HOME}/lib:${LD_LIBRARY_PATH:-}"
else
  LLAMA_SERVER_BIN="llama-server"
fi

if [[ -z "${MODEL_PATH}" && "${GGUF_REPO}" == "deepreinforce-ai/Ornith-1.0-35B-GGUF" ]]; then
  MODEL_PATH="$(find "${HF_HOME:-${HOME}/.cache/huggingface}/hub/models--deepreinforce-ai--Ornith-1.0-35B-GGUF/snapshots" \
    -name "ornith-1.0-35b-${QUANT}.gguf" -print -quit 2>/dev/null || true)"
fi

if [[ -z "${MODEL_PATH}" && -f "${REPO_ROOT}/artifacts/quant/ornith-1.0-35b-${QUANT}.gguf" ]]; then
  MODEL_PATH="${REPO_ROOT}/artifacts/quant/ornith-1.0-35b-${QUANT}.gguf"
fi

model_args=(-hf "${GGUF_REPO}:${QUANT}")
if [[ -n "${MODEL_PATH}" ]]; then
  model_args=(-m "${MODEL_PATH}")
fi

server_args=(
  "${model_args[@]}"
  --alias "${ALIAS}"
  --host "${HOST}"
  --port "${PORT}"
  -c "${CTX_SIZE}"
  -np "${PARALLEL}"
  -b "${BATCH_SIZE}"
  -ub "${UBATCH_SIZE}"
  -ngl 99
  --flash-attn "${FLASH_ATTN}"
  --reasoning "${REASONING}"
  --reasoning-format "${REASONING_FORMAT}"
)

if [[ -n "${REASONING_BUDGET}" ]]; then
  server_args+=(--reasoning-budget "${REASONING_BUDGET}")
fi

if [[ -n "${REASONING_BUDGET_MESSAGE}" ]]; then
  server_args+=(--reasoning-budget-message "${REASONING_BUDGET_MESSAGE}")
fi

if [[ -n "${SPEC_TYPE}" ]]; then
  server_args+=(--spec-type "${SPEC_TYPE}")
fi

if [[ -n "${SPEC_DRAFT_MODEL}" ]]; then
  server_args+=(--spec-draft-model "${SPEC_DRAFT_MODEL}")
fi

if [[ -n "${SPEC_DRAFT_HF}" ]]; then
  server_args+=(--spec-draft-hf "${SPEC_DRAFT_HF}")
fi

if [[ -n "${SPEC_TYPE}" || -n "${SPEC_DRAFT_MODEL}" || -n "${SPEC_DRAFT_HF}" ]]; then
  server_args+=(--spec-draft-ngl "${SPEC_DRAFT_NGL}")
fi

if [[ -n "${SPEC_DRAFT_N_MAX}" ]]; then
  server_args+=(--spec-draft-n-max "${SPEC_DRAFT_N_MAX}")
fi

if [[ -n "${SPEC_DRAFT_N_MIN}" ]]; then
  server_args+=(--spec-draft-n-min "${SPEC_DRAFT_N_MIN}")
fi

if [[ -n "${SPEC_DRAFT_P_MIN}" ]]; then
  server_args+=(--spec-draft-p-min "${SPEC_DRAFT_P_MIN}")
fi

if [[ -n "${SPEC_DRAFT_DEVICE}" ]]; then
  server_args+=(--spec-draft-device "${SPEC_DRAFT_DEVICE}")
fi

if [[ -n "${SPEC_TYPE}" || -n "${SPEC_DRAFT_MODEL}" || -n "${SPEC_DRAFT_HF}" ]]; then
  if [[ "${SPEC_DRAFT_BACKEND_SAMPLING}" == "1" ]]; then
    server_args+=(--spec-draft-backend-sampling)
  else
    server_args+=(--no-spec-draft-backend-sampling)
  fi
fi

if [[ "${CONT_BATCHING}" == "1" ]]; then
  server_args+=(--cont-batching)
else
  server_args+=(--no-cont-batching)
fi

if [[ "${METRICS}" == "1" ]]; then
  server_args+=(--metrics)
fi

if [[ "${BACKEND_SAMPLING}" == "1" ]]; then
  server_args+=(--backend-sampling)
fi

if [[ -n "${CACHE_RAM}" ]]; then
  server_args+=(--cache-ram "${CACHE_RAM}")
fi

exec "${LLAMA_SERVER_BIN}" \
  "${server_args[@]}"