#!/usr/bin/env bash
# Unattended in-place upgrade for a CoderAI orchestrator, for a host that should
# follow the `production` branch without someone logging in.
#
# Driven by a systemd timer (see packaging/linux/systemd/) or cron. One run:
#
#   1. take a lock, so two timers can never upgrade the same image at once;
#   2. remember the version baked into the image, and tag that image as the
#      rollback point;
#   3. run the host runner's `--upgrade`;
#   4. read the baked version again. CHANGED is the signal to restart -- not the
#      runner's exit code, which reports 0 both when it upgraded and when the
#      image was already current, so an exit-code-driven loop would restart the
#      service on every tick;
#   5. restart only if the service is idle, else defer to the next tick (the new
#      code sits in the image until then, which is safe). A host that is busy at
#      EVERY tick would otherwise never upgrade at all, so after MAX_DEFERRALS
#      busy ticks the script books a FORCED upgrade for the next midnight and
#      takes it, work in flight or not. Observed before that existed: 24 skips in
#      one day and two versions behind, because the busy path exited before the
#      counter the force depended on was ever incremented;
#   6. verify health after the restart, and ROLL BACK to the remembered image if
#      it does not come up.
#
# MAX_DEFERRALS=0 keeps the old meaning: wait for an idle moment, however long,
# and never book a forced restart. Anything above 0 is a promise that the upgrade
# WILL land, at midnight at the latest.
#
# Exit 0: nothing to do, upgraded and healthy, or deliberately deferred.
# Exit 1: the upgrade or the restart failed -- and then the message says whether
#         a rollback happened, because that is the only thing an operator who
#         reads one line needs to know.
set -uo pipefail

CONF="${CODERAI_AUTOUPGRADE_CONF:-${XDG_CONFIG_HOME:-$HOME/.config}/coderai/autoupgrade.conf}"

# Every knob below is `VAR="${VAR:-default}"`, which reads as "the environment may
# override it" -- but the conf is sourced FIRST and assigns unconditionally
# (`MAX_DEFERRALS="0"`), so the conf silently won and `MAX_DEFERRALS=1
# coderai-autoupgrade` did exactly nothing, with no line in the log to say why.
# An operator passing a variable on the command line is overriding the unattended
# policy deliberately and for one run, so the environment has to win. Snapshot what
# it set, source the conf, put the snapshot back.
AUTOUPGRADE_TUNABLES="ENGINE IMAGE_TAG UPGRADE_REF RUNNER RESTART_CMD HEALTH_URL
    HEALTH_TIMEOUT ADMIN_TOKEN METRICS_URL STATUS_URL REQUIRE_IDLE_CHECK
    MAX_DEFERRALS CONTAINER_NAME STATE_DIR LOG ROLLBACK_TAG LOG_MAX_LINES"
for _t in $AUTOUPGRADE_TUNABLES; do
    eval "[ -n \"\${$_t:-}\" ] && _ENV_$_t=\"\$$_t\""
done
# shellcheck disable=SC1090
[ -f "$CONF" ] && . "$CONF"
for _t in $AUTOUPGRADE_TUNABLES; do
    eval "[ -n \"\${_ENV_$_t:-}\" ] && $_t=\"\$_ENV_$_t\""
done
unset _t

# Which container engine. A rootless podman host is the common case for this
# script, so prefer podman when both are installed.
if [ -z "${ENGINE:-}" ]; then
    if command -v podman >/dev/null 2>&1; then ENGINE=podman
    elif command -v docker >/dev/null 2>&1; then ENGINE=docker
    else echo "autoupgrade: neither podman nor docker on PATH" >&2; exit 1; fi
fi
IMAGE_TAG="${IMAGE_TAG:-ghcr.io/nextime/coderai:latest}"
UPGRADE_REF="${UPGRADE_REF:-production}"
RUNNER="${RUNNER:-$(command -v coderai-docker || true)}"
RESTART_CMD="${RESTART_CMD:-systemctl --user restart coderai.service}"
HEALTH_URL="${HEALTH_URL:-http://127.0.0.1:8777/healthz}"
HEALTH_TIMEOUT="${HEALTH_TIMEOUT:-300}"
# The idle check. An upgrade is SKIPPED outright while CoderAI is serving
# anything -- a request on a local engine or, the case that matters most here,
# one being served by a rented RunPod pod, because interrupting that both loses
# the work and keeps paying for the pod that was doing it.
#
# Both probes are derived from HEALTH_URL so there is one address to get right:
#   <base>/metrics            -> coderai_engine_inflight, every request the front
#                                has in flight, pod-served ones included
#   <base>/v1/runpod/status   -> summary.inflight, the pods' own count
# The higher of the two wins. They overlap; neither is trusted to be complete.
ADMIN_TOKEN="${ADMIN_TOKEN:-}"
_base="${HEALTH_URL%/healthz}"; _base="${_base%/}"
METRICS_URL="${METRICS_URL:-$_base/metrics}"
STATUS_URL="${STATUS_URL:-$_base/v1/runpod/status}"
# Fail CLOSED: if neither probe can be read, assume it is serving and do nothing.
# A silent upgrade that interrupts production is worse than a skipped one, and an
# unreadable probe usually means a wrong token or URL, not an idle server.
REQUIRE_IDLE_CHECK="${REQUIRE_IDLE_CHECK:-1}"
# Which container is the server, for the one question the HTTP probes cannot
# answer: is there a server at all? See container_state().
CONTAINER_NAME="${CONTAINER_NAME:-coderai}"
# Once the image HAS been upgraded, the restart can be deferred this many ticks
# before going ahead anyway. 0 (the default) means never interrupt a request:
# the new code waits in the image and every run says so.
MAX_DEFERRALS="${MAX_DEFERRALS:-0}"
STATE_DIR="${STATE_DIR:-${XDG_STATE_HOME:-$HOME/.local/state}/coderai}"
LOG="${LOG:-$STATE_DIR/autoupgrade.log}"
ROLLBACK_TAG="${ROLLBACK_TAG:-coderai-autoupgrade:rollback}"

mkdir -p "$STATE_DIR" || { echo "autoupgrade: cannot create $STATE_DIR" >&2; exit 1; }
mkdir -p "$(dirname "$LOG")" 2>/dev/null
PENDING="$STATE_DIR/restart-pending"
DEFERRALS="$STATE_DIR/deferrals"
# When a host is NEVER idle, the upgrade is armed for the next midnight and
# this file holds that deadline (epoch seconds). The hourly timer ticks at
# 00:00, so the arming run and the firing run are the ordinary schedule --
# no second timer to install, nothing to clean up if the host reboots.
FORCE_AT="$STATE_DIR/force-upgrade-at"

say(){ printf '%s %s\n' "$(date -Is)" "$*" | tee -a "$LOG"; }

# Running hourly for years, this is the one file that grows without bound. Keep
# the tail and nothing else: the journal has the per-run copy, this log exists so
# the HISTORY of upgrades survives a journal rotation.
LOG_MAX_LINES="${LOG_MAX_LINES:-5000}"
if [ -f "$LOG" ] && [ "$(wc -l < "$LOG" 2>/dev/null || echo 0)" -gt "$LOG_MAX_LINES" ]; then
    tail -n "$((LOG_MAX_LINES / 2))" "$LOG" > "$LOG.trim" 2>/dev/null \
        && mv "$LOG.trim" "$LOG"
fi

[ -n "$RUNNER" ] || { say "FAIL: no coderai-docker on PATH (set RUNNER= in $CONF)"; exit 1; }
[ -x "$RUNNER" ] || { say "FAIL: RUNNER '$RUNNER' is not executable"; exit 1; }

# Echo the settings that were actually resolved, every run. The config file is
# SOURCED, so an unquoted multi-word value silently keeps only its first word --
# `RESTART_CMD=systemctl --user restart coderai.service` assigns "systemctl" and
# runs the rest. That mangling is invisible until it restarts nothing, the health
# check fails and a perfectly good upgrade gets rolled back. One line makes it
# obvious instead.
IDLE_CHECK=off; [ -n "$ADMIN_TOKEN" ] && IDLE_CHECK=on
say "config: engine=$ENGINE image=$IMAGE_TAG ref=$UPGRADE_REF container=$CONTAINER_NAME"
say "config: restart=[$RESTART_CMD] health=$HEALTH_URL idle-check=$IDLE_CHECK"
# The two knobs that decide whether a run acts or waits. Neither was in the log,
# so "it skipped again" and "it was told to skip" looked identical -- and with the
# conf overriding the environment, an operator could be looking at the opposite of
# what they typed.
say "config: require-idle=$REQUIRE_IDLE_CHECK max-deferrals=$MAX_DEFERRALS"
if [ -f "$FORCE_AT" ]; then
    _at="$(cat "$FORCE_AT" 2>/dev/null || echo 0)"
    say "config: a FORCED upgrade is booked for $(date -d "@$_at" '+%F %T' 2>/dev/null || echo "$_at")"
fi

# ---------------------------------------------------------------------- the lock
# flock on a dedicated fd, non-blocking: a run that overlaps the previous one is
# a no-op, not a queue. The timer fires again soon enough.
exec 9>"$STATE_DIR/autoupgrade.lock"
if ! flock -n 9; then
    say "another autoupgrade run holds the lock; nothing done"
    exit 0
fi

baked_version(){
    "$ENGINE" run --rm --entrypoint grep "$IMAGE_TAG" -h "__version__" \
        /opt/coderai/app/codai/__init__.py 2>/dev/null \
        | tr -d ' "' | cut -d= -f2 | head -1
}

# "running", "stopped", or "unknown" when the engine itself could not be asked.
#
# Both probes are HTTP, so a server that is DOWN looks exactly like a wrong token:
# both are unreadable, inflight() says "unknown", and the fail-closed guard refuses
# to touch the install that most needs touching. That is not hypothetical -- zeiss
# came back from a hard reset with no container at all, and every hourly run from
# then on skipped, so the host sat on old code and stayed DOWN until a human
# noticed. A container that is not running cannot be serving anything, and the
# restart that follows an upgrade is also what brings it back.
#
# Only a POSITIVE answer counts. If the engine cannot be asked (daemon down,
# wrong socket) this says "unknown" and the fail-closed guard still applies: the
# point is to recognise an absent server, not to assume one.
container_state(){
    local names rc
    names="$("$ENGINE" ps --format '{{.Names}}' 2>/dev/null)"; rc=$?
    [ "$rc" -eq 0 ] || { echo unknown; return; }
    printf '%s\n' "$names" | grep -qx "$CONTAINER_NAME" && { echo running; return; }
    # An install that names its container something else is still covered: any
    # running container of the image being upgraded is the server.
    [ -n "$("$ENGINE" ps --quiet --filter "ancestor=$IMAGE_TAG" 2>/dev/null)" ] \
        && { echo running; return; }
    echo stopped
}

_probe_metrics(){
    # Two numbers from one scrape, as "<inflight> <jobs>":
    #   coderai_engine_inflight  -- requests in flight, pod-served ones included
    #   coderai_jobs_active      -- running/queued TASKS across the engines
    # The second is what catches training: a LoRA run is started by a request
    # that returns immediately and then works for hours, so inflight is 0 the
    # whole time. An older front serves no jobs gauge; "-" then says so rather
    # than claiming zero.
    curl -fsS --max-time 10 -H "Authorization: Bearer $ADMIN_TOKEN" "$METRICS_URL" 2>/dev/null \
        | awk '/^coderai_engine_inflight[ {]/ { n += $NF; seen = 1 }
               /^coderai_jobs_active[ {]/     { j = $NF; jseen = 1 }
               END { if (!seen) { print "unknown unknown"; exit }
                     printf "%d %s\n", n, (jseen ? sprintf("%d", j) : "-") }'
}

_probe_runpod(){
    # What the pods themselves report. Present even if the front's counter is not.
    curl -fsS --max-time 10 -H "Authorization: Bearer $ADMIN_TOKEN" "$STATUS_URL" 2>/dev/null \
        | python3 -c 'import json,sys
try: print(int((json.load(sys.stdin).get("summary") or {}).get("inflight") or 0))
except Exception: print("unknown")' 2>/dev/null
}

# Echoes "<count> <detail>", or "unknown <why>" when nothing could be read.
#
# The MAXIMUM of the sources, never the sum: a generation appears as both an
# in-flight request and a running task, so adding them would double-count and
# the numbers in the log would not match anything an operator can see.
inflight(){
    local out m j r best detail
    # Cheapest and most certain signal first, and the only one that works when the
    # server is down: no container, no work. Checked before the token, because a
    # host with no server configured for the probes is the same situation.
    if [ "$(container_state)" = stopped ]; then
        echo "0 no '$CONTAINER_NAME' container running — nothing can be in flight"
        return
    fi
    [ -n "$ADMIN_TOKEN" ] || { echo "unknown no ADMIN_TOKEN configured"; return; }
    out="$(_probe_metrics)"; m="${out%% *}"; j="${out##* }"
    r="$(_probe_runpod)"
    case "$m" in ''|unknown) m=unknown ;; *[!0-9]*) m=unknown ;; esac
    case "$j" in ''|-|unknown) j=unknown ;; *[!0-9]*) j=unknown ;; esac
    case "$r" in ''|unknown) r=unknown ;; *[!0-9]*) r=unknown ;; esac
    if [ "$m" = unknown ] && [ "$r" = unknown ]; then
        echo "unknown neither $METRICS_URL nor $STATUS_URL could be read"
        return
    fi
    best=0; detail=""
    [ "$m" != unknown ] && { best="$m"; detail="engines=$m"; }
    if [ "$j" != unknown ]; then
        detail="${detail:+$detail }jobs=$j"
        [ "$j" -gt "$best" ] && best="$j"
    else
        # Only worth saying once the rest is known to work: it means the front
        # predates coderai_jobs_active, so training is invisible to this check.
        detail="${detail:+$detail }jobs=? (front too old)"
    fi
    if [ "$r" != unknown ]; then
        detail="${detail:+$detail }runpod=$r"
        [ "$r" -gt "$best" ] && best="$r"
    fi
    echo "$best $detail"
}

# After MAX_DEFERRALS busy ticks, stop waiting for a window that is not coming
# and book one. Midnight local, because that is the cheapest hour to interrupt
# on a business-hours workload and it is a time an operator can predict without
# reading this file.
arm_midnight_force(){
    local n="$1" at
    # Tomorrow 00:00 local. `date -d` is GNU; fall back to "now + the seconds
    # left in the day" so this still works on a busybox host.
    at="$(date -d 'tomorrow 00:00' +%s 2>/dev/null || true)"
    if [ -z "$at" ]; then
        at=$(( $(date +%s) + (86400 - ($(date +%H) * 3600 + $(date +%M) * 60 + $(date +%S))) ))
    fi
    echo "$at" > "$FORCE_AT"
    say "NEVER IDLE for $n run(s) — booking a FORCED upgrade for $(date -d "@$at" '+%F %T' 2>/dev/null || echo "$at")."
    say "      It will restart the service even with work in flight, because an"
    say "      install that is busy every hour otherwise never upgrades at all."
    say "      Cancel it by deleting $FORCE_AT."
}

# Is a booked forced upgrade due? Clears the booking either way it is used, so a
# deadline that has passed fires ONCE and not on every tick afterwards.
force_due(){
    [ -f "$FORCE_AT" ] || return 1
    local at now
    at="$(cat "$FORCE_AT" 2>/dev/null || echo 0)"
    case "$at" in ''|*[!0-9]*) rm -f "$FORCE_AT"; return 1 ;; esac
    now="$(date +%s)"
    [ "$now" -ge "$at" ] || return 1
    return 0
}

# Any work in progress -- a request OR a detached job like training -- or we
# cannot tell? Then this run does nothing. Calling it "serving" was wrong:
# the case this exists for has no request in flight at all.
guard_idle(){
    local out count detail
    out="$(inflight)"; count="${out%% *}"; detail="${out#* }"
    if [ "$count" = unknown ]; then
        if [ "$REQUIRE_IDLE_CHECK" = "1" ]; then
            say "SKIP: cannot tell whether CoderAI is serving ($detail) — refusing to"
            say "      upgrade blind. Fix the probe, or set REQUIRE_IDLE_CHECK=\"0\" to"
            say "      upgrade without the check (it will interrupt live requests)."
            return 1
        fi
        say "WARN: idle check unavailable ($detail) and REQUIRE_IDLE_CHECK=0 — continuing"
        return 0
    fi
    if [ "$count" -gt 0 ]; then
        say "SKIP: $count piece(s) of work in progress [$detail] — leaving the image alone"
        # COUNT IT. This used to return without touching the counter, and the
        # counter is only otherwise incremented in the restart step -- which is
        # never reached when we exit here. So `max-deferrals` could not fire and
        # a host busy on every hourly tick never upgraded AT ALL, which is not
        # what the timer's own comment promises ("or after MAX_DEFERRALS hours of
        # never being idle"). Observed: 24 skips in one day, two versions behind,
        # deferrals file still 0.
        local n
        n=$(cat "$DEFERRALS" 2>/dev/null || echo 0)
        case "$n" in ''|*[!0-9]*) n=0 ;; esac
        n=$((n + 1))
        echo "$n" > "$DEFERRALS"
        if [ "$MAX_DEFERRALS" -gt 0 ] && [ "$n" -ge "$MAX_DEFERRALS" ] \
                && [ ! -f "$FORCE_AT" ]; then
            arm_midnight_force "$n"
        else
            say "      busy run $n$( [ "$MAX_DEFERRALS" -gt 0 ] && printf '/%s' "$MAX_DEFERRALS" )"
        fi
        return 1
    fi
    say "idle ($detail)"
    return 0
}

healthy(){
    local deadline=$(( $(date +%s) + HEALTH_TIMEOUT ))
    while [ "$(date +%s)" -lt "$deadline" ]; do
        if curl -fsS --max-time 10 "$HEALTH_URL" 2>/dev/null | grep -q '"ok"[[:space:]]*:[[:space:]]*true'; then
            return 0
        fi
        sleep 10
    done
    return 1
}

restart_and_verify(){
    say "restarting: $RESTART_CMD"
    if ! sh -c "$RESTART_CMD" >>"$LOG" 2>&1; then
        say "WARN: the restart command returned non-zero"
    fi
    if healthy; then return 0; fi
    return 1
}

# ------------------------------------------------------------------- the upgrade
BEFORE="$(baked_version)"
[ -n "$BEFORE" ] || { say "FAIL: cannot read the version from $IMAGE_TAG (wrong tag? engine '$ENGINE'?)"; exit 1; }

if [ -f "$PENDING" ]; then
    # A previous run upgraded the image but could not restart. Do not upgrade
    # again on top of that; land what is already there first.
    say "a restart is still pending for $(cat "$PENDING" 2>/dev/null) — not upgrading further"
else
    # Nothing is touched while CoderAI is serving. Checked BEFORE the upgrade and
    # not only before the restart, so a busy host never even reaches the state
    # where new code sits in the image waiting for a window.
    if force_due; then
        FORCED=1
        say "FORCED UPGRADE (booked because this host was never idle) — proceeding"
        say "      with work possibly in flight. Requests being served WILL be"
        say "      interrupted by the restart; that is the trade this booking made."
    else
        FORCED=0
        guard_idle || exit 0
    fi
    say "image $IMAGE_TAG carries $BEFORE; checking '$UPGRADE_REF'"
    # The rollback point is a TAG on the current image, so it costs nothing and
    # survives the commit the upgrade makes onto $IMAGE_TAG.
    "$ENGINE" tag "$IMAGE_TAG" "$ROLLBACK_TAG" >/dev/null 2>&1 \
        || say "WARN: could not tag the rollback point; continuing without one"

    # CONTAINER_ENGINE is how the runner picks podman over its docker default. It
    # has to be told: on a rootless-podman host the runner would otherwise look
    # for a docker socket that is not there, and the upgrade would fail on every
    # tick with nothing to show for it but a non-zero exit.
    CONTAINER_ENGINE="$ENGINE" CODERAI_UPGRADE_REF="$UPGRADE_REF" \
        "$RUNNER" --upgrade --upgrade-ref "$UPGRADE_REF" "$IMAGE_TAG" >>"$LOG" 2>&1
    rc=$?
    [ "$rc" -eq 0 ] || { say "FAIL: the upgrade run exited $rc; image left as it was"; exit 1; }

    AFTER="$(baked_version)"
    if [ "$AFTER" = "$BEFORE" ]; then
        say "already current at $BEFORE; nothing to restart"
        rm -f "$DEFERRALS"
        exit 0
    fi
    say "image upgraded $BEFORE -> $AFTER"
    printf '%s -> %s\n' "$BEFORE" "$AFTER" > "$PENDING"
fi

# ------------------------------------------------------------- restart, or defer
# The upgrade itself takes a minute, so ask again: a request may have arrived
# while it ran, and the restart is the part that would kill it.
N=$(cat "$DEFERRALS" 2>/dev/null || echo 0)
OUT="$(inflight)"; BUSY="${OUT%% *}"; WHAT="${OUT#* }"
if [ "${FORCED:-0}" = 1 ]; then
    # The whole point of the booking. Deferring here would put us back where the
    # booking was made: waiting for a window this host does not have.
    [ "$BUSY" = unknown ] || [ "$BUSY" -eq 0 ] 2>/dev/null \
        || say "forced: restarting with work in flight [$WHAT]"
elif [ "$BUSY" = unknown ] || { [ "$BUSY" -gt 0 ] 2>/dev/null; }; then
    if [ "$N" -lt "$MAX_DEFERRALS" ]; then
        echo $((N + 1)) > "$DEFERRALS"
        say "work in progress [$WHAT] — deferring the restart ($((N + 1))/$MAX_DEFERRALS)"
        exit 0
    fi
    if [ "$MAX_DEFERRALS" -le 0 ]; then
        echo $((N + 1)) > "$DEFERRALS"
        say "work in progress [$WHAT] — the new code is in the image and waits for an idle"
        say "      moment ($((N + 1)) run(s) so far). Set MAX_DEFERRALS>0 to restart anyway."
        exit 0
    fi
    say "work in progress [$WHAT] but $N deferrals already — restarting anyway"
fi

if restart_and_verify; then
    say "OK: running $(cat "$PENDING" 2>/dev/null | sed 's/.*-> //') and healthy"
    # The booking and the busy-run count both end here: the upgrade landed, so
    # the next busy stretch starts counting from zero rather than inheriting a
    # deadline from the last one.
    rm -f "$PENDING" "$DEFERRALS" "$FORCE_AT"
    exit 0
fi

# ------------------------------------------------------------------- the rollback
say "FAIL: the service did not become healthy within ${HEALTH_TIMEOUT}s after the upgrade"
if "$ENGINE" image inspect "$ROLLBACK_TAG" >/dev/null 2>&1; then
    say "rolling back $IMAGE_TAG to the pre-upgrade image"
    "$ENGINE" tag "$ROLLBACK_TAG" "$IMAGE_TAG" >/dev/null 2>&1
    if restart_and_verify; then
        say "FAIL: upgrade reverted — back on $BEFORE and healthy. The upgrade is NOT applied."
        rm -f "$PENDING" "$DEFERRALS"
        exit 1
    fi
    say "FAIL: the service is down and the rollback did not bring it back. Needs a human."
    exit 1
fi
say "FAIL: no rollback image available; the service is down. Needs a human."
exit 1
