#!/usr/bin/env bash
# Purpose: turn the adapter this part trained into a served alias: run Part 13's
#          export-gguf.sh to merge, convert, quantise and smoke-test it, then write the
#          llama-swap and LiteLLM snippets that give the result a name next to the base
#          model, so the comparison in this part's lab can address both at once
# Platform: spark, strix, nvidia (and mac, where the adapter comes from mlx_lm.lora and is
#           fused with mlx_lm.fuse first; see the lab page for that path)
# Minimum memory: 16 GB; the merge inside export-gguf.sh needs one bfloat16 copy of the base
# Assumes: Part 13's export-gguf.sh and merge-adapter.py in $PART13_DIR; a llama.cpp
#          checkout holding convert_hf_to_gguf.py at $LLAMA_CPP; llama-quantize and
#          llama-cli on PATH; and the gateway from Part 9, whose llama-swap.yaml and
#          litellm-config.yaml this script writes snippets for rather than editing.
#
# Usage: bash export-and-alias.sh ADAPTER_DIR ALIAS [QUANT]
#   ADAPTER_DIR  what train-agent-sft.py saved, e.g. runs/agent-qwen3-4b
#   ALIAS        the gateway name the fine-tune will answer to, e.g. local/agent-tuned
#   QUANT        a type llama-quantize accepts; defaults to Q4_K_M
#
# Environment: PART13_DIR (default ../part-13-supervised-fine-tuning), LLAMA_CPP
#              (default $HOME/llama.cpp), OUT_DIR (default $HOME/models), SNIPPET_DIR
#              (default ./gateway-snippets), CTX (default 16384), SKIP_EXPORT=1 to write
#              only the snippets for a GGUF you already have.
set -euo pipefail

ADAPTER="${1:-}"
ALIAS="${2:-}"
QUANT="${3:-Q4_K_M}"

PART13_DIR="${PART13_DIR:-../part-13-supervised-fine-tuning}"
LLAMA_CPP="${LLAMA_CPP:-$HOME/llama.cpp}"
OUT_DIR="${OUT_DIR:-$HOME/models}"
SNIPPET_DIR="${SNIPPET_DIR:-./gateway-snippets}"
CTX="${CTX:-16384}"

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

[[ -n "$ADAPTER" ]] || die "usage: bash export-and-alias.sh ADAPTER_DIR ALIAS [QUANT]"
[[ -n "$ALIAS" ]] || die "give the gateway alias as the second argument, e.g. local/agent-tuned"
case "$ALIAS" in
  */*) : ;;
  *) die "the alias should look like local/agent-tuned so it reads as a role, not a file" ;;
esac

NAME="$(basename "$ADAPTER")"
GGUF="$OUT_DIR/$NAME-$QUANT.gguf"

if [[ "${SKIP_EXPORT:-0}" == "1" ]]; then
  echo "==> 1/3 Skipping the export (SKIP_EXPORT=1); expecting $GGUF"
else
  [[ -f "$ADAPTER/adapter_config.json" ]] || die "$ADAPTER has no adapter_config.json; point at what train-agent-sft.py saved"
  [[ -f "$PART13_DIR/export-gguf.sh" ]] || die "export-gguf.sh not found under $PART13_DIR; set PART13_DIR to Part 13's lab directory"
  [[ -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"
  command -v llama-quantize >/dev/null || die "llama-quantize is not on PATH; build it as in Part 6"

  echo "==> 1/3 Merging, converting, quantising and smoke-testing with Part 13's script"
  echo "    Nothing about that step is specific to agents: an adapter is an adapter, and"
  echo "    the reason it is one script for both parts is that it should stay one script."
  LLAMA_CPP="$LLAMA_CPP" OUT_DIR="$OUT_DIR" bash "$PART13_DIR/export-gguf.sh" "$ADAPTER" "$QUANT"
fi

[[ -f "$GGUF" ]] || die "expected $GGUF and it is not there; read the export output above"
SIZE=$(wc -c < "$GGUF" | tr -d ' ')
echo "    quantised export: $GGUF ($SIZE bytes)"

echo "==> 2/3 Writing the gateway snippets"
mkdir -p "$SNIPPET_DIR"
SAFE_NAME="${ALIAS//\//-}"
SWAP_SNIPPET="$SNIPPET_DIR/llama-swap-$SAFE_NAME.yaml"
LITELLM_SNIPPET="$SNIPPET_DIR/litellm-$SAFE_NAME.yaml"

# The chat template travels inside the GGUF, and --jinja is what makes llama-server use it
# rather than a built-in guess. For a tool-calling fine-tune that flag is the whole point:
# without it the model is served with a template that does not know what a tool call is.
cat > "$SWAP_SNIPPET" <<EOF
# Merge this block into the models: section of the llama-swap.yaml you wrote in Part 9.
# It sits next to the base model rather than replacing it, because this part's comparison
# addresses both aliases in the same run.
  $ALIAS:
    name: Agent fine-tune ($NAME)
    description: $NAME exported at $QUANT by Part 27's export-and-alias.sh
    cmd: |
      \${server}
      --model $GGUF
      --alias $ALIAS
      --ctx-size $CTX
      --jinja
    checkEndpoint: /health
EOF

cat > "$LITELLM_SNIPPET" <<EOF
# Merge this entry into the model_list of the litellm-config.yaml you wrote in Part 9.
model_list:
  - model_name: $ALIAS
    litellm_params:
      model: openai/$ALIAS
      api_base: os.environ/LLAMA_SWAP_BASE_URL
      api_key: os.environ/LOCAL_API_KEY
EOF

echo "    $SWAP_SNIPPET"
echo "    $LITELLM_SNIPPET"

echo "==> 3/3 What to do with them"
if [[ -n "${APPEND_TO:-}" ]]; then
  [[ -f "$APPEND_TO" ]] || die "APPEND_TO=$APPEND_TO does not exist"
  if grep -qF "$ALIAS:" "$APPEND_TO"; then
    die "$APPEND_TO already mentions $ALIAS; edit it by hand rather than appending twice"
  fi
  cp "$APPEND_TO" "$APPEND_TO.before-$SAFE_NAME"
  # Only the block itself, without the two comment lines that explain it.
  tail -n +4 "$SWAP_SNIPPET" >> "$APPEND_TO"
  echo "    appended the llama-swap block to $APPEND_TO"
  echo "    the previous version is $APPEND_TO.before-$SAFE_NAME"
  echo "    the LiteLLM entry is still yours to merge: $LITELLM_SNIPPET"
else
  echo "    Merge both snippets into your Part 9 gateway configuration, then reload it."
  echo "    Set APPEND_TO=/path/to/llama-swap.yaml to have this script append the first"
  echo "    one for you; it takes a copy first and refuses if the alias is already there."
fi

echo
echo "Then check that the alias answers, and that it answers with a tool call:"
echo "  curl -s \"\$GATEWAY/v1/models\" | grep -o '$ALIAS'"
echo "  python3 ../part-24-tools-mcp-and-the-agent-loop/tool-call-reliability.py \\"
echo "      --base-url \"\$GATEWAY/v1\" --model $ALIAS --repeat 1"
echo "A fine-tune that serves but does not call tools is almost always the template: check"
echo "that --jinja is on and that the tokeniser files were saved with the adapter."