#!/usr/bin/env bash
# Llama Champion + Off the Grid + Well-Tuned
# Run Nemotron 4B GGUF + fine-tuned LoRA adapter via llama.cpp
#
# Terminal 1:
#   ./scripts/start_llamacpp.sh
# Terminal 2:
#   PMS_INFERENCE_BACKEND=llamacpp python app.py

set -euo pipefail

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
GGUF_DIR="${PMS_GGUF_DIR:-${ROOT}/models/nemotron}"
GGUF_FILE="${PMS_GGUF_FILE:-NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf}"
GGUF_REPO="${PMS_GGUF_MODEL:-nvidia/NVIDIA-Nemotron-3-Nano-4B-GGUF}"
LORA_GGUF="${PMS_LLAMACPP_LORA:-${GGUF_DIR}/plane-mode-study-coach-lora.gguf}"
LORA_SCALE="${PMS_LLAMACPP_LORA_SCALE:-1.0}"
PORT="${PMS_LLAMACPP_PORT:-8080}"
USE_FT="${PMS_USE_FINETUNED:-true}"

mkdir -p "${GGUF_DIR}"

if [[ ! -f "${GGUF_DIR}/${GGUF_FILE}" ]]; then
  echo "Downloading base Nemotron 4B GGUF from ${GGUF_REPO}..."
  python3 - <<PY
from huggingface_hub import hf_hub_download
from pathlib import Path
path = hf_hub_download(
    repo_id="${GGUF_REPO}",
    filename="${GGUF_FILE}",
    local_dir="${GGUF_DIR}",
)
print("Saved", path)
PY
fi

if [[ "${USE_FT}" == "true" || "${USE_FT}" == "1" ]]; then
  if [[ ! -f "${LORA_GGUF}" ]]; then
    echo "Fine-tuned LoRA GGUF missing — converting from Hugging Face adapter..."
    python3 "${ROOT}/scripts/export_lora_gguf.py" --out "${LORA_GGUF}"
  fi
fi

export PMS_INFERENCE_BACKEND=llamacpp
export PMS_LLAMACPP_URL="http://127.0.0.1:${PORT}/v1"
export PMS_LLAMACPP_LORA="${LORA_GGUF}"

LORA_ARGS=()
if [[ -f "${LORA_GGUF}" ]]; then
  LORA_ARGS=(--lora-scaled "${LORA_GGUF}:${LORA_SCALE}")
  echo "Fine-tuned LoRA: ${LORA_GGUF} (scale=${LORA_SCALE})"
fi

if ! command -v llama-server >/dev/null 2>&1; then
  VENDOR_SERVER="${ROOT}/vendor/llama.cpp/build/bin/llama-server"
  if [[ -x "${VENDOR_SERVER}" ]]; then
    export PATH="${ROOT}/vendor/llama.cpp/build/bin:${PATH}"
  else
    echo "ERROR: llama-server not found. Build it with:"
    echo "  CC=gcc CXX=g++ cmake -S vendor/llama.cpp -B vendor/llama.cpp/build -DCMAKE_BUILD_TYPE=Release"
    echo "  cmake --build vendor/llama.cpp/build -j --target llama-server"
    echo "  Or install from https://github.com/ggml-org/llama.cpp"
    exit 1
  fi
fi

NGPU="${PMS_LLAMACPP_GPU_LAYERS:-}"
if [[ -z "${NGPU}" ]]; then
  if command -v nvidia-smi >/dev/null 2>&1; then
    NGPU=99
  else
    NGPU=0
  fi
fi

echo "Starting llama-server on port ${PORT} (n-gpu-layers=${NGPU})..."
echo "  base: ${GGUF_DIR}/${GGUF_FILE}"
exec llama-server \
  --model "${GGUF_DIR}/${GGUF_FILE}" \
  "${LORA_ARGS[@]}" \
  --host 0.0.0.0 \
  --port "${PORT}" \
  --n-gpu-layers "${NGPU}" \
  --ctx-size 8192