#!/usr/bin/env bash
# Purpose: upgrade the container image behind the gateway without finding out in
#          production whether the new one works - pull it, resolve it to a digest, run it
#          beside the live one on a spare port, check it, swap it in, re-check the
#          gateway, and put the old digest back automatically if the re-check fails
# Platform: spark, strix, nvidia (Docker Engine with the Compose plugin), and mac where
#           Docker Desktop is installed. A native install has no image to swap; the same
#           procedure applies with a second binary and is described on the page.
# Minimum memory: 8 GB. The canary is started without loading a model, so it costs almost
#           nothing; --deep asks it to answer a real request and needs room for a second
#           copy of the model, which most machines do not have.
# Assumes: run from the gateway directory built in Part 9, holding compose.yaml and a .env
#          with LLAMA_SWAP_IMAGE in it. This script edits exactly two keys in .env,
#          LLAMA_SWAP_IMAGE and LLAMA_SWAP_IMAGE_PREVIOUS, and touches nothing else. It
#          never prints the contents of .env, which holds your keys.
#
# Usage: bash upgrade-engine.sh --image ghcr.io/example/engine:sometag
#        bash upgrade-engine.sh --image ghcr.io/example/engine:sometag --deep
#        bash upgrade-engine.sh --rollback
set -euo pipefail

COMPOSE_FILE="compose.yaml"
SERVICE="swap"
IMAGE_KEY="LLAMA_SWAP_IMAGE"
CANARY_PORT="9293"
CHANGELOG=""
NEW_IMAGE=""
DEEP=0
ROLLBACK=0
ASSUME_YES=0

usage() { sed -n '2,20p' "$0"; }

while [ $# -gt 0 ]; do
    case "$1" in
        --image)      NEW_IMAGE="${2:?--image needs an image reference}"; shift 2 ;;
        --compose)    COMPOSE_FILE="${2:?--compose needs a file}"; shift 2 ;;
        --service)    SERVICE="${2:?--service needs a compose service name}"; shift 2 ;;
        --key)        IMAGE_KEY="${2:?--key needs an .env variable name}"; shift 2 ;;
        --port)       CANARY_PORT="${2:?--port needs a port}"; shift 2 ;;
        --changelog)  CHANGELOG="${2:?--changelog needs a file}"; shift 2 ;;
        --deep)       DEEP=1; shift ;;
        --rollback)   ROLLBACK=1; shift ;;
        --yes)        ASSUME_YES=1; shift ;;
        -h|--help)    usage; exit 0 ;;
        *) echo "unknown argument: $1" >&2; exit 2 ;;
    esac
done

command -v docker >/dev/null 2>&1 || { echo "docker is not on PATH." >&2; exit 1; }
command -v curl >/dev/null 2>&1 || { echo "curl is not on PATH." >&2; exit 1; }
[ -f "$COMPOSE_FILE" ] || { echo "no such compose file: ${COMPOSE_FILE}" >&2; exit 1; }
[ -f ".env" ] || { echo "no .env beside ${COMPOSE_FILE}; run this from the gateway directory." >&2; exit 1; }

# Read one key out of .env without sourcing the whole file, so that a stray line in a file
# holding your keys cannot execute.
env_value() {
    grep -E "^$1=" .env | tail -n 1 | cut -d= -f2- || true
}

# Replace one key in .env, or append it. A temporary file and a move, so an interrupted
# write cannot leave you with half a .env.
set_env_value() {
    local key="$1" value="$2" tmp
    tmp="$(mktemp)"
    if grep -qE "^${key}=" .env; then
        sed "s|^${key}=.*|${key}=${value}|" .env > "$tmp"
    else
        cp .env "$tmp"
        printf '%s=%s\n' "$key" "$value" >> "$tmp"
    fi
    mv "$tmp" .env
}

confirm() {
    [ "$ASSUME_YES" -eq 1 ] && return 0
    printf '%s [y/N] ' "$1"
    read -r answer
    case "$answer" in y|Y|yes|YES) return 0 ;; *) return 1 ;; esac
}

record() {
    # record <action> <from> <to> <result>
    [ -n "$CHANGELOG" ] || return 0
    printf '{"lab":"part-23/upgrades","action":"%s","service":"%s","from":"%s","to":"%s","result":"%s","at":"%s"}\n' \
        "$1" "$SERVICE" "$2" "$3" "$4" "$(date -u +%Y-%m-%dT%H:%M:%SZ)" >> "$CHANGELOG"
    echo "    recorded in ${CHANGELOG}"
}

CURRENT="$(env_value "$IMAGE_KEY")"
PREVIOUS="$(env_value "${IMAGE_KEY}_PREVIOUS")"

if [ "$ROLLBACK" -eq 1 ]; then
    [ -n "$PREVIOUS" ] || { echo "No ${IMAGE_KEY}_PREVIOUS in .env; nothing to roll back to." >&2; exit 1; }
    echo "==> Rolling back ${SERVICE}"
    echo "    from ${CURRENT}"
    echo "    to   ${PREVIOUS}"
    confirm "Proceed?" || { echo "Nothing changed."; exit 0; }
    set_env_value "$IMAGE_KEY" "$PREVIOUS"
    set_env_value "${IMAGE_KEY}_PREVIOUS" "$CURRENT"
    docker compose -f "$COMPOSE_FILE" up -d "$SERVICE"
    echo "==> Rolled back. Run check-gateway.sh to confirm."
    record rollback "$CURRENT" "$PREVIOUS" "applied"
    exit 0
fi

[ -n "$NEW_IMAGE" ] || { echo "Give --image <reference>, or --rollback." >&2; usage >&2; exit 2; }

echo "==> Pulling ${NEW_IMAGE}"
docker image pull "$NEW_IMAGE" >/dev/null

# A tag moves; a digest does not. Resolve now, so that what is recorded in .env and in the
# changelog is the thing that actually ran, not a name that pointed at it this morning.
DIGEST="$(docker image inspect --format '{{if .RepoDigests}}{{index .RepoDigests 0}}{{end}}' "$NEW_IMAGE" 2>/dev/null || true)"
if [ -z "$DIGEST" ]; then
    echo "    This image has no repository digest, which means it was built locally rather"
    echo "    than pulled. Continuing with the reference you gave, which is not pinned."
    DIGEST="$NEW_IMAGE"
fi
echo "    resolved to ${DIGEST}"

echo "==> Starting a canary on port ${CANARY_PORT}, beside the live service"
CANARY="upgrade-canary-$$"
CLEANED=0
cleanup_canary() {
    [ "$CLEANED" -eq 1 ] && return 0
    CLEANED=1
    docker rm -f "$CANARY" >/dev/null 2>&1 || true
}
trap cleanup_canary EXIT

# The canary runs the same configuration as the live service, read-only, on a spare port.
# It is not given the model directory unless --deep was asked for, so it starts, answers
# and stops without ever loading weights.
CANARY_ARGS=(run -d --name "$CANARY" -p "127.0.0.1:${CANARY_PORT}:8080"
             -v "$(pwd)/llama-swap.yaml:/app/config.yaml:ro")
if [ "$DEEP" -eq 1 ]; then
    MODELS_DIR="$(env_value MODELS_DIR)"
    [ -n "$MODELS_DIR" ] || { echo "--deep needs MODELS_DIR in .env." >&2; exit 1; }
    CANARY_ARGS+=(-v "${MODELS_DIR}:/models:ro" -e "MODELS_DIR=/models" -e "LLAMA_BIN=llama-server")
fi
docker "${CANARY_ARGS[@]}" "$NEW_IMAGE" --config /app/config.yaml --listen 0.0.0.0:8080 >/dev/null

echo "==> Checking the canary"
CANARY_URL="http://127.0.0.1:${CANARY_PORT}"
OK=1
for _ in $(seq 1 30); do
    if curl -sf "${CANARY_URL}/health" >/dev/null 2>&1; then OK=0; break; fi
    sleep 2
done
if [ "$OK" -ne 0 ]; then
    echo "    FAIL  the canary never became healthy. Its log follows; nothing was changed." >&2
    docker logs --tail 40 "$CANARY" >&2 || true
    record upgrade "$CURRENT" "$DIGEST" "canary-failed"
    exit 1
fi
echo "    PASS  health endpoint answers"

if curl -sf "${CANARY_URL}/v1/models" | grep -q '"id"'; then
    echo "    PASS  the model list is published"
else
    echo "    FAIL  the model list is empty; the new image reads your configuration differently." >&2
    docker logs --tail 40 "$CANARY" >&2 || true
    record upgrade "$CURRENT" "$DIGEST" "canary-failed"
    exit 1
fi

if [ "$DEEP" -eq 1 ]; then
    echo "    ... asking the canary for a real completion; this loads a model and takes a while"
    if curl -sf --max-time 600 "${CANARY_URL}/v1/chat/completions" \
        -H "Content-Type: application/json" \
        -d '{"model":"local/chat","messages":[{"role":"user","content":"Reply with one word: ready"}],"max_tokens":8}' \
        | grep -q '"content"'; then
        echo "    PASS  the canary generated a completion"
    else
        echo "    FAIL  the canary could not generate. Nothing was changed." >&2
        docker logs --tail 40 "$CANARY" >&2 || true
        record upgrade "$CURRENT" "$DIGEST" "canary-failed"
        exit 1
    fi
fi

cleanup_canary

echo "==> Swapping the live service"
echo "    from ${CURRENT}"
echo "    to   ${DIGEST}"
confirm "Proceed?" || { echo "Nothing changed."; exit 0; }

set_env_value "${IMAGE_KEY}_PREVIOUS" "$CURRENT"
set_env_value "$IMAGE_KEY" "$DIGEST"
docker compose -f "$COMPOSE_FILE" up -d "$SERVICE"

echo "==> Re-checking the gateway"
if [ -x ./check-gateway.sh ] || [ -f ./check-gateway.sh ]; then
    if bash ./check-gateway.sh; then
        echo "==> Upgrade complete. The previous digest is in .env as ${IMAGE_KEY}_PREVIOUS."
        record upgrade "$CURRENT" "$DIGEST" "applied"
        exit 0
    fi
    echo "    The gateway check failed after the swap. Putting the old digest back." >&2
    set_env_value "$IMAGE_KEY" "$CURRENT"
    set_env_value "${IMAGE_KEY}_PREVIOUS" "$DIGEST"
    docker compose -f "$COMPOSE_FILE" up -d "$SERVICE"
    record upgrade "$CURRENT" "$DIGEST" "reverted-automatically"
    exit 1
fi

echo "    check-gateway.sh from Part 9 is not in this directory, so the swap was not"
echo "    verified. Run your own check now, and 'bash upgrade-engine.sh --rollback' if it"
echo "    is unwell."
record upgrade "$CURRENT" "$DIGEST" "applied-unverified"