#!/usr/bin/env bash
# Purpose: run the gateway without containers - llama-swap loading models on demand and
#          LiteLLM routing aliases in front of it - and stop both cleanly on Ctrl-C
# Platform: mac primarily, and any track that would rather not run containers. The
#           container path (compose.yaml) is the one Tracks S, X and N normally use.
# Minimum memory: 8 GB
# Assumes: llama-swap and litellm on PATH, llama-server built as in Part 6, and a .env
#          beside this script copied from env-example.txt with MODELS_DIR and LLAMA_BIN
#          filled in. No database is started here, so virtual keys and spend records are
#          unavailable on this path; the master key is the only key. Logs are written to
#          ./gateway-logs/.
set -euo pipefail

HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ENV_FILE="${ENV_FILE:-${HERE}/.env}"
LOG_DIR="${LOG_DIR:-${HERE}/gateway-logs}"

if [ ! -f "$ENV_FILE" ]; then
    echo "No .env found at ${ENV_FILE}." >&2
    echo "Copy env-example.txt to .env and fill in MODELS_DIR, LLAMA_BIN and the keys." >&2
    exit 1
fi

set -a
# shellcheck source=/dev/null
. "$ENV_FILE"
set +a

: "${GATEWAY_HOST:=127.0.0.1}"
: "${GATEWAY_PORT:=4000}"
: "${SWAP_HOST:=127.0.0.1}"
: "${SWAP_PORT:=9292}"

for required in MODELS_DIR LLAMA_BIN LITELLM_MASTER_KEY; do
    if [ -z "${!required:-}" ]; then
        echo "${required} is empty in ${ENV_FILE}. Fill it in and run this again." >&2
        exit 1
    fi
done

for binary in llama-swap litellm; do
    if ! command -v "$binary" >/dev/null 2>&1; then
        echo "${binary} is not on PATH." >&2
        echo "See the project page for how to install it on your track." >&2
        exit 1
    fi
done

if [ ! -x "$LLAMA_BIN" ]; then
    echo "LLAMA_BIN points at ${LLAMA_BIN}, which is not an executable file." >&2
    exit 1
fi

if [ ! -d "$MODELS_DIR" ]; then
    echo "MODELS_DIR points at ${MODELS_DIR}, which is not a directory." >&2
    exit 1
fi

# The router needs to know where llama-swap is. Both are exported so that
# litellm-config.yaml's os.environ/ lookups resolve.
export LLAMA_SWAP_BASE_URL="http://${SWAP_HOST}:${SWAP_PORT}/v1"
export DIRECT_BASE_URL="${DIRECT_BASE_URL:-http://${SWAP_HOST}:${SWAP_PORT}/v1}"
export LOCAL_API_KEY="${LOCAL_API_KEY:-}"

mkdir -p "$LOG_DIR"

SWAP_PID=""
GATEWAY_PID=""

stop_all() {
    echo ""
    echo "==> stopping"
    if [ -n "$GATEWAY_PID" ] && kill -0 "$GATEWAY_PID" 2>/dev/null; then
        kill "$GATEWAY_PID" 2>/dev/null || true
        wait "$GATEWAY_PID" 2>/dev/null || true
    fi
    if [ -n "$SWAP_PID" ] && kill -0 "$SWAP_PID" 2>/dev/null; then
        kill "$SWAP_PID" 2>/dev/null || true
        wait "$SWAP_PID" 2>/dev/null || true
    fi
    echo "    both processes stopped. Logs are in ${LOG_DIR}."
}
trap stop_all EXIT INT TERM

echo "==> llama-swap on http://${SWAP_HOST}:${SWAP_PORT}"
llama-swap \
    --config "${HERE}/llama-swap.yaml" \
    --listen "${SWAP_HOST}:${SWAP_PORT}" \
    >>"${LOG_DIR}/llama-swap.log" 2>&1 &
SWAP_PID=$!

# Give it a moment to bind before the router starts probing it.
for _ in $(seq 1 30); do
    if ! kill -0 "$SWAP_PID" 2>/dev/null; then
        echo "llama-swap exited immediately. Last lines of its log:" >&2
        tail -n 20 "${LOG_DIR}/llama-swap.log" >&2 || true
        exit 1
    fi
    if command -v curl >/dev/null 2>&1 &&
        curl -sf "http://${SWAP_HOST}:${SWAP_PORT}/health" >/dev/null 2>&1; then
        break
    fi
    sleep 1
done

echo "==> LiteLLM on http://${GATEWAY_HOST}:${GATEWAY_PORT}"
litellm \
    --config "${HERE}/litellm-config.yaml" \
    --host "$GATEWAY_HOST" \
    --port "$GATEWAY_PORT" \
    >>"${LOG_DIR}/litellm.log" 2>&1 &
GATEWAY_PID=$!

cat <<INFO

    gateway   http://${GATEWAY_HOST}:${GATEWAY_PORT}
    swap      http://${SWAP_HOST}:${SWAP_PORT}
    logs      ${LOG_DIR}/

    Check it with:  bash check-gateway.sh
    Stop it with:   Ctrl-C in this terminal

INFO

wait