#!/usr/bin/env bash
# GGUF quantization workflow: convert, quantize, serve, verify.
# Requires llama.cpp built with `make` (or `cmake`) and a HF checkpoint.
set -euo pipefail

LLAMA_CPP_DIR="${LLAMA_CPP_DIR:-$HOME/llama.cpp}"
HF_MODEL="${1:?Usage: $0 <path/to/hf-model-dir>}"
OUT_BASE="${2:-model}"

CONVERT_PY="$LLAMA_CPP_DIR/convert_hf_to_gguf.py"
QUANTIZE_BIN="$LLAMA_CPP_DIR/llama-quantize"
SERVER_BIN="$LLAMA_CPP_DIR/llama-server"

for bin in "$CONVERT_PY" "$QUANTIZE_BIN" "$SERVER_BIN"; do
    if [[ ! -x "$bin" && ! -f "$bin" ]]; then
        echo "Missing: $bin" >&2
        echo "Build llama.cpp first: cd $LLAMA_CPP_DIR && make" >&2
        exit 1
    fi
done

echo "==> 1/4 Converting HF checkpoint to FP16 GGUF"
F16_FILE="${OUT_BASE}-f16.gguf"
python "$CONVERT_PY" "$HF_MODEL" --outfile "$F16_FILE" --outtype f16

echo "==> 2/4 Quantizing to Q4_K_M (best size/quality balance)"
Q4_FILE="${OUT_BASE}-Q4_K_M.gguf"
"$QUANTIZE_BIN" "$F16_FILE" "$Q4_FILE" Q4_K_M

echo "==> 3/4 Sizes:"
ls -lh "$F16_FILE" "$Q4_FILE" | awk '{print "  " $9 " -> " $5}'

echo "==> 4/4 Serving with llama-server (OpenAI-compatible endpoint)"
# -ngl 999 offloads all layers to GPU if available; set 0 for CPU-only.
"$SERVER_BIN" \
  -m "$Q4_FILE" \
  --ctx-size 8192 \
  -ngl 999 \
  --host 127.0.0.1 \
  --port 8080 \
  &

SERVER_PID=$!
sleep 2

echo "==> Smoke test: POST /v1/chat/completions"
curl -s http://127.0.0.1:8080/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "local-model",
    "messages": [
      {"role": "system", "content": "You are a terse assistant."},
      {"role": "user", "content": "Explain quantization in one sentence."}
    ],
    "max_tokens": 64
  }' | python3 -m json.tool

kill "$SERVER_PID" 2>/dev/null || true
echo "==> Done. Files: $F16_FILE, $Q4_FILE"