#!/usr/bin/env bash
# Purpose: the reduced path for Tracks X and M, where vLLM's connectors are not available:
#          run llama-server with prompt caching, several slots and a quantised key-value
#          cache, then save a slot's cache to disk and bring it back, which is llama.cpp's
#          version of a disk tier
# Platform: strix, mac (and anywhere else llama.cpp runs; the flags are identical)
# Minimum memory: 16 GB; a four-bit 8B model with a 32,768-token cache fits comfortably
# Assumes: llama-server on PATH or in LLAMA_BIN, built for your accelerator as in Part 6;
#          a .env copied from env-example.txt with LLAMA_MODEL pointing at a GGUF file;
#          curl on PATH for the slot subcommands
#
# Usage: bash llama-cache-tiers.sh serve            start the server with the tiers on
#        bash llama-cache-tiers.sh slots            show what each slot is holding
#        bash llama-cache-tiers.sh save 0 name.bin  write slot 0's cache to disk
#        bash llama-cache-tiers.sh restore 0 name.bin  read it back into slot 0
#        bash llama-cache-tiers.sh erase 0          discard slot 0's cache
set -euo pipefail

HERE="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
if [ -f "${HERE}/.env" ]; then
    set -a
    # shellcheck disable=SC1091  # written by the reader from env-example.txt
    . "${HERE}/.env"
    set +a
fi

ACTION="${1:-serve}"
SERVE_HOST="${SERVE_HOST:-127.0.0.1}"
LLAMA_PORT="${LLAMA_PORT:-8080}"
LLAMA_MODEL="${LLAMA_MODEL:-}"
LLAMA_PARALLEL="${LLAMA_PARALLEL:-2}"
LLAMA_SLOT_SAVE_PATH="${LLAMA_SLOT_SAVE_PATH:-./slot-cache}"
LLAMA_CACHE_TYPE="${LLAMA_CACHE_TYPE:-q8_0}"
SERVED_NAME="${SERVED_NAME:-local-chat}"
CTX="${CTX:-8192}"
LLAMA_BIN="${LLAMA_BIN:-}"
BASE="http://${SERVE_HOST}:${LLAMA_PORT}"

fail() { printf '%s\n' "$*" >&2; exit 1; }

SERVER="llama-server"
if [ -n "$LLAMA_BIN" ] && [ -x "${LLAMA_BIN}/llama-server" ]; then
    SERVER="${LLAMA_BIN}/llama-server"
fi

case "$ACTION" in
    serve)
        command -v "$SERVER" >/dev/null 2>&1 || [ -x "$SERVER" ] \
            || fail "llama-server is not on PATH and LLAMA_BIN does not contain it."
        [ -n "$LLAMA_MODEL" ] || fail "LLAMA_MODEL is not set in .env."
        [ -f "$LLAMA_MODEL" ] || fail "LLAMA_MODEL (${LLAMA_MODEL}) is not a file."
        mkdir -p "$LLAMA_SLOT_SAVE_PATH"

        cat <<INFO
==> llama-server with the cache tiers this track can reach
    model            ${LLAMA_MODEL}
    served as        ${SERVED_NAME}
    listening on     ${BASE}
    context          ${CTX}
    slots            ${LLAMA_PARALLEL}
    key-value type   ${LLAMA_CACHE_TYPE} for both keys and values
    slot save path   ${LLAMA_SLOT_SAVE_PATH}

    What each of these does, from the server README:
      --cache-prompt     prompt caching, which is enabled by default; passed explicitly
                         here so the run is self-documenting
      --parallel         the number of server slots, which is also the number of distinct
                         conversations that can stay warm at once
      -ctk / -ctv        the key and value cache data type; q8_0 halves the cache against
                         the f16 default, and unlike prompt caching this DOES change
                         outputs, so do not carry it into a comparison against an f16 run
      --slot-save-path   where "save slot kv cache" writes to; without it the save and
                         restore endpoints are disabled
      --metrics --slots  the two endpoints this lab reads

    The README also documents --cache-reuse, "min chunk size to attempt reusing from the
    cache via KV shifting", which is not used here and is not yet in this course's
    captured command reference. Check your build's --help before adding it.

INFO
        exec "$SERVER" \
            --model "$LLAMA_MODEL" \
            --alias "$SERVED_NAME" \
            --host "$SERVE_HOST" \
            --port "$LLAMA_PORT" \
            --ctx-size "$CTX" \
            --parallel "$LLAMA_PARALLEL" \
            --n-gpu-layers 999 \
            --cache-prompt \
            --cache-type-k "$LLAMA_CACHE_TYPE" \
            --cache-type-v "$LLAMA_CACHE_TYPE" \
            --slot-save-path "$LLAMA_SLOT_SAVE_PATH" \
            --metrics \
            --slots
        ;;

    slots)
        command -v curl >/dev/null 2>&1 || fail "curl is not on PATH."
        printf '==> What each slot is holding\n'
        curl -fsS "${BASE}/slots"
        printf '\n'
        ;;

    save|restore)
        command -v curl >/dev/null 2>&1 || fail "curl is not on PATH."
        SLOT="${2:-}"
        FILENAME="${3:-}"
        [ -n "$SLOT" ] || fail "Usage: bash llama-cache-tiers.sh ${ACTION} <slot> <filename>"
        [ -n "$FILENAME" ] || fail "Usage: bash llama-cache-tiers.sh ${ACTION} <slot> <filename>"
        printf '==> %s slot %s as %s\n' "$ACTION" "$SLOT" "$FILENAME"
        START="$(date +%s%N)"
        curl -fsS -X POST "${BASE}/slots/${SLOT}?action=${ACTION}" \
            -H "Content-Type: application/json" \
            -d "{\"filename\": \"${FILENAME}\"}"
        END="$(date +%s%N)"
        printf '\n    took %s ms\n' "$(( (END - START) / 1000000 ))"
        printf '    That figure is the disk tier''s latency for this conversation. Compare\n'
        printf '    it against the time the same prompt takes to prefill from cold.\n'
        ;;

    erase)
        command -v curl >/dev/null 2>&1 || fail "curl is not on PATH."
        SLOT="${2:-}"
        [ -n "$SLOT" ] || fail "Usage: bash llama-cache-tiers.sh erase <slot>"
        printf '==> Erasing slot %s. This discards its cache and cannot be undone.\n' "$SLOT"
        printf '    Check "bash llama-cache-tiers.sh slots" first if you are unsure.\n'
        curl -fsS -X POST "${BASE}/slots/${SLOT}?action=erase"
        printf '\n'
        ;;

    *)
        fail "Unknown action '${ACTION}'. Use: serve, slots, save, restore, erase."
        ;;
esac