#!/usr/bin/env bash
# Purpose: turn a LoRA adapter into a quantised GGUF file that llama.cpp and everything built
#          on it can serve: merge into the base at bfloat16, convert, quantise, and run one
#          prompt through the result so the export is proved rather than assumed
# Platform: spark, strix, nvidia (and mac for the merge and conversion, although Track M's
#           adapters come from mlx_lm.lora and are fused with mlx_lm.fuse first)
# Minimum memory: 12 GB; the merge needs one bfloat16 copy of the base in system memory
# Assumes: a Python environment with torch, transformers and peft; a llama.cpp checkout
#          holding convert_hf_to_gguf.py; llama-quantize and llama-cli on PATH or pointed at
#          by LLAMA_QUANTIZE and LLAMA_CLI; merge-adapter.py next to this script; and disk
#          for roughly two and a half times the bfloat16 size of the model
#
# Usage: bash export-gguf.sh ADAPTER_DIR [QUANT]
#   ADAPTER_DIR  the directory train-lora.py saved, e.g. runs/format-qwen3-1.7b
#   QUANT        a type llama-quantize accepts; defaults to Q4_K_M
#
# Environment: LLAMA_CPP (default ~/llama.cpp), MERGED_DIR, OUT_DIR, LLAMA_QUANTIZE,
#              LLAMA_CLI, PROMPT, PREDICT, SKIP_RUN=1 to stop before the smoke test.
#              REUSE_MERGED=1 only after checking that the existing merge belongs to
#              this exact adapter and base; otherwise choose a fresh OUT_DIR.
set -euo pipefail

ADAPTER="${1:-}"
QUANT="${2:-Q4_K_M}"
LLAMA_CPP="${LLAMA_CPP:-$HOME/llama.cpp}"
OUT_DIR="${OUT_DIR:-$HOME/models}"
PREDICT="${PREDICT:-96}"
PROMPT="${PROMPT:-The monitoring system reports: the model gateway is unreachable from every machine we have tried this morning.}"
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

die() { echo "export-gguf: $*" >&2; exit 1; }

[[ -n "$ADAPTER" ]] || die "usage: bash export-gguf.sh ADAPTER_DIR [QUANT]"
[[ -f "$ADAPTER/adapter_config.json" ]] || die "$ADAPTER has no adapter_config.json; point at what train-lora.py saved"
[[ -f "$LLAMA_CPP/convert_hf_to_gguf.py" ]] || die "convert_hf_to_gguf.py not found under $LLAMA_CPP; set LLAMA_CPP to your llama.cpp checkout"
[[ -f "$HERE/merge-adapter.py" ]] || die "merge-adapter.py is not next to this script"
command -v python3 >/dev/null || die "python3 is not on PATH"

QUANTIZE_BIN="${LLAMA_QUANTIZE:-$(command -v llama-quantize || true)}"
CLI_BIN="${LLAMA_CLI:-$(command -v llama-cli || true)}"
[[ -n "$QUANTIZE_BIN" ]] || die "llama-quantize is not on PATH; set LLAMA_QUANTIZE to its path (built in Part 6)"

NAME="$(basename "$ADAPTER")"
MERGED_DIR="${MERGED_DIR:-$OUT_DIR/$NAME-merged}"
GGUF_BF16="$OUT_DIR/$NAME-bf16.gguf"
GGUF_QUANT="$OUT_DIR/$NAME-$QUANT.gguf"
mkdir -p "$OUT_DIR"

echo "==> 1/4 Merging the adapter into its base at bfloat16, on the CPU"
if [[ -d "$MERGED_DIR" ]]; then
  [[ "${REUSE_MERGED:-0}" == "1" ]] || die "$MERGED_DIR already exists. Choose a fresh OUT_DIR for this run; set REUSE_MERGED=1 only after verifying its adapter/base provenance."
  echo "    Reusing explicitly approved merge: $MERGED_DIR"
else
  python3 "$HERE/merge-adapter.py" --adapter "$ADAPTER" --merged-dir "$MERGED_DIR" --compare 2
fi

echo "==> 2/4 Converting to GGUF at bfloat16"
echo "    Converting at full precision first and quantising afterwards keeps information"
echo "    the quantiser can use. Converting straight to a small type throws it away."
python3 "$LLAMA_CPP/convert_hf_to_gguf.py" "$MERGED_DIR" \
  --outfile "$GGUF_BF16" \
  --outtype bf16

echo "==> 3/4 Quantising to $QUANT"
"$QUANTIZE_BIN" "$GGUF_BF16" "$GGUF_QUANT" "$QUANT"

for f in "$GGUF_BF16" "$GGUF_QUANT"; do
  SIZE=$(wc -c < "$f" | tr -d ' ')
  echo "    $(basename "$f"): $SIZE bytes"
done

if [[ "${SKIP_RUN:-0}" == "1" ]]; then
  echo "==> 4/4 Skipped the smoke test (SKIP_RUN=1)"
else
  [[ -n "$CLI_BIN" ]] || die "llama-cli is not on PATH; set LLAMA_CLI or re-run with SKIP_RUN=1"
  echo "==> 4/4 One prompt through the quantised export, greedily"
  echo "    Compare this with what merge-adapter.py printed for the same prompt. A merge"
  echo "    should match; a quantisation may differ a little. A wholesale change of format"
  echo "    means the chat template did not survive the export."
  "$CLI_BIN" \
    --model "$GGUF_QUANT" \
    --prompt "$PROMPT" \
    --predict "$PREDICT" \
    --temp 0 \
    --seed 0
fi

echo
echo "==> Done."
echo "    merged safetensors: $MERGED_DIR"
echo "    full-precision GGUF: $GGUF_BF16"
echo "    quantised GGUF:      $GGUF_QUANT"
echo "    Delete the full-precision GGUF once the quantised one has been scored, not before."