#!/usr/bin/env bash
# In-image code upgrader for the CoderAI distributable container.
#
# Runs INSIDE the container (driven by the host runner `coderai-docker --upgrade`,
# i.e. run_oci.sh). It fetches the configured branch of the CoderAI git repo,
# compares its product version (codai/__init__.py __version__) with the version
# baked into this image, and — when the remote is newer (or --force) — replaces
# the in-image application tree at /opt/coderai/app with the fetched code.
#
# The image ships WITHOUT a .git dir (it's stripped at build time), so this does
# a fresh shallow clone of the target ref rather than a `git pull`. After it
# returns, the host runner `docker commit`s the container back onto the SAME
# image tag, so the updated code persists with no Dockerfile rebuild and no new
# overlay image.
#
# Contract with the host runner (exit codes):
#   0   -> code was updated; the host should commit the container.
#   10  -> already up to date (nothing changed); host must NOT commit.
#   !=0 -> error; host must NOT commit.
#
# Configuration (all optional; the host runner passes these as env):
#   CODERAI_UPGRADE_REPO   git URL to fetch (https or ssh). Default: the nexlab
#                          HTTPS repo, or its SSH form when a key is provided.
#   CODERAI_UPGRADE_REF    branch/tag/commit to fetch. Default: production.
#   CODERAI_UPGRADE_FORCE  "1" to upgrade even if not strictly newer.
#   CODERAI_UPGRADE_SSH_KEY  path (inside the container) to an SSH private key;
#                          when set, git uses it (StrictHostKeyChecking=accept-new).
set -euo pipefail

APP_DIR="/opt/coderai/app"
DEFAULT_HTTPS_REPO="https://git.nexlab.net/nexlab/coderai.git"
DEFAULT_SSH_REPO="git@git.nexlab.net:nexlab/coderai.git"

REF="${CODERAI_UPGRADE_REF:-production}"
FORCE="${CODERAI_UPGRADE_FORCE:-0}"
SSH_KEY="${CODERAI_UPGRADE_SSH_KEY:-}"
# When the fetched code changes its declared dependencies, re-run pip so new
# packages land in the image's python env. Set CODERAI_UPGRADE_SKIP_PIP=1 to skip
# (e.g. an offline/air-gapped host that only wants the code refresh).
SKIP_PIP="${CODERAI_UPGRADE_SKIP_PIP:-0}"
PYBIN="/opt/coderai/python/bin/python3"

# Pick the default repo form to match the auth method: SSH URL when a key was
# supplied, HTTPS otherwise. An explicit CODERAI_UPGRADE_REPO always wins.
if [[ -n "${CODERAI_UPGRADE_REPO:-}" ]]; then
  REPO="$CODERAI_UPGRADE_REPO"
elif [[ -n "$SSH_KEY" ]]; then
  REPO="$DEFAULT_SSH_REPO"
else
  REPO="$DEFAULT_HTTPS_REPO"
fi

log(){ printf '[upgrade] %s\n' "$*" >&2; }
die(){ printf '[upgrade] error: %s\n' "$*" >&2; exit 1; }

command -v git >/dev/null 2>&1 || die "git is not available in this image"

# Read __version__ = "x.y.z" from a codai/__init__.py file.
read_version() {
  local f="$1"
  [[ -f "$f" ]] || return 1
  sed -n 's/^__version__ *= *["'"'"']\([^"'"'"']*\)["'"'"'].*/\1/p' "$f" | head -1
}

# Compare two dotted numeric versions. Prints: -1 if a<b, 0 if a==b, 1 if a>b.
# Non-numeric components sort as 0. Missing components are treated as 0.
vercmp() {
  local a="$1" b="$2" IFS=.
  local -a A=($a) B=($b)
  local n=$(( ${#A[@]} > ${#B[@]} ? ${#A[@]} : ${#B[@]} ))
  local i x y
  for (( i=0; i<n; i++ )); do
    x=${A[i]:-0}; y=${B[i]:-0}
    x=$((10#${x//[^0-9]/0} + 0)) 2>/dev/null || x=0
    y=$((10#${y//[^0-9]/0} + 0)) 2>/dev/null || y=0
    if (( x < y )); then echo -1; return; fi
    if (( x > y )); then echo 1; return; fi
  done
  echo 0
}

# Emit the normalized, sorted set of requirement specs a tree declares: the base
# requirements.txt plus the container extras (packaging/common/requirements-oci.txt,
# minus its build-time `-r` self-include). Comments and surrounding whitespace are
# stripped so the diff below reflects real dependency changes, not formatting.
reqs_content() {
  local d="$1"
  { [[ -f "$d/requirements.txt" ]] && cat "$d/requirements.txt"
    [[ -f "$d/packaging/common/requirements-oci.txt" ]] \
      && grep -vE '^[[:space:]]*-r[[:space:]]' "$d/packaging/common/requirements-oci.txt"
  } 2>/dev/null \
    | sed 's/#.*//' | sed 's/[[:space:]]*$//' \
    | grep -vE '^[[:space:]]*$' | sort -u
}

# Refresh the launcher scripts and service configs that live OUTSIDE the app tree
# (in /usr/local/bin, /etc/nginx, /etc/supervisor). The in-image code replace only
# touches /opt/coderai/app, so without this a rebuild would be the only way to
# pick up launcher/config changes — including changes to THIS upgrade script.
# Mirrors the COPY set in packaging/linux/Dockerfile.update. Each file is written
# via a temp + atomic rename so replacing the currently-running coderai-upgrade
# doesn't corrupt this in-flight process (bash keeps executing the old inode).
sync_system_files() {
  local src="$WORK/src"
  _install_file() {  # <src> <dest> <mode>
    [[ -f "$1" ]] || return 0
    local tmp; tmp="$(dirname "$2")/.$(basename "$2").upgrade.$$"
    cp "$1" "$tmp" && chmod "$3" "$tmp" && mv -f "$tmp" "$2" \
      || { rm -f "$tmp"; log "warning: could not update $2"; return 1; }
  }
  _install_file "$src/packaging/linux/launcher/coderai-oci"         /usr/local/bin/coderai             0755
  _install_file "$src/packaging/linux/launcher/with-env"            /usr/local/bin/with-env            0755
  _install_file "$src/packaging/linux/launcher/coderai-entrypoint"  /usr/local/bin/coderai-entrypoint  0755
  _install_file "$src/packaging/linux/launcher/coderai-upgrade"     /usr/local/bin/coderai-upgrade     0755
  _install_file "$src/packaging/linux/launcher/wav2lip"             /usr/local/bin/wav2lip             0755
  _install_file "$src/packaging/linux/launcher/sadtalker"           /usr/local/bin/sadtalker           0755
  _install_file "$src/packaging/linux/nginx.conf"                   /etc/nginx/nginx.conf              0644
  _install_file "$src/packaging/linux/supervisord.conf"             /etc/supervisor/supervisord.conf   0644
  _install_file "$src/packaging/linux/README-RUN.txt"              /opt/coderai/README-RUN.txt        0644
}

# Install ONLY the dependency specs that are new or version-changed relative to
# what the image already declared ($WORK/old_reqs vs $WORK/new_reqs, both sorted).
# This is deliberately a DELTA install, not a full `pip install -r requirements`:
# the repo's requirements list can contain entries that don't cleanly reinstall in
# this environment (e.g. a package with no py3.13 wheel that the image provisions
# another way). Reprocessing the whole file would fail on those pre-existing lines;
# installing only what actually changed adds new deps while leaving satisfied ones
# untouched. Returns 0 (nothing to do) when there is no delta.
pip_sync() {
  [[ -x "$PYBIN" ]] || { log "python not found at $PYBIN — skipping dependency sync"; return 1; }
  local delta; delta="$(comm -13 "$WORK/old_reqs" "$WORK/new_reqs")"
  if [[ -z "$delta" ]]; then
    log "no new or changed dependencies — nothing to install"
    return 0
  fi
  log "installing new/changed dependencies:"
  printf '%s\n' "$delta" | sed 's/^/  + /' >&2
  local pip_args=(-m pip install --no-input --disable-pip-version-check)
  # Prefer prebuilt CUDA wheels when the build left them behind (older images
  # don't have /opt/wheels; the flag is simply omitted then).
  [[ -d /opt/wheels ]] && pip_args+=(--find-links /opt/wheels)
  printf '%s\n' "$delta" | "$PYBIN" "${pip_args[@]}" -r /dev/stdin
}

CUR_VER="$(read_version "$APP_DIR/codai/__init__.py" || true)"
[[ -n "$CUR_VER" ]] || CUR_VER="0"
log "installed version: $CUR_VER"
log "source:  $REPO  (ref: $REF)"

# Auth: when an SSH key path is given, drive git through it. accept-new adds the
# host key on first contact without prompting, but still rejects a changed key.
export GIT_TERMINAL_PROMPT=0
if [[ -n "$SSH_KEY" ]]; then
  [[ -f "$SSH_KEY" ]] || die "ssh key not found in container at: $SSH_KEY"
  export GIT_SSH_COMMAND="ssh -i '$SSH_KEY' -o IdentitiesOnly=yes -o StrictHostKeyChecking=accept-new -o BatchMode=yes"
fi

WORK="$(mktemp -d /tmp/coderai-upgrade.XXXXXX)"
cleanup(){ rm -rf "$WORK"; }
trap cleanup EXIT

log "fetching '$REF'…"
# Shallow single-branch clone: minimal transfer, no history. A ref that is a tag
# or bare commit also works with --branch for tags; for a raw commit we fall back
# to a full-ish fetch.
if ! git clone --depth 1 --branch "$REF" --single-branch "$REPO" "$WORK/src" 2>"$WORK/clone.err"; then
  if grep -qiE "Could not find remote branch|not found in upstream" "$WORK/clone.err"; then
    log "ref '$REF' is not a branch/tag; trying to fetch it as a commit…"
    git init -q "$WORK/src"
    git -C "$WORK/src" remote add origin "$REPO"
    git -C "$WORK/src" fetch -q --depth 1 origin "$REF" || { cat "$WORK/clone.err" >&2; die "fetch of ref '$REF' failed"; }
    git -C "$WORK/src" checkout -q FETCH_HEAD
  else
    cat "$WORK/clone.err" >&2
    die "clone failed (auth? network? wrong repo/ref?)"
  fi
fi

NEW_VER="$(read_version "$WORK/src/codai/__init__.py" || true)"
[[ -n "$NEW_VER" ]] || die "fetched tree has no codai/__init__.py __version__"
log "fetched version:   $NEW_VER"

CMP="$(vercmp "$NEW_VER" "$CUR_VER")"
if [[ "$CMP" != "1" && "$FORCE" != "1" ]]; then
  if [[ "$CMP" == "0" ]]; then
    log "already up to date ($CUR_VER) — nothing to do (use --force to reinstall)."
  else
    log "installed version ($CUR_VER) is newer than '$REF' ($NEW_VER) — refusing to downgrade (use --force)."
  fi
  exit 10
fi

[[ "$FORCE" == "1" && "$CMP" != "1" ]] && log "forcing upgrade to $NEW_VER (was $CUR_VER)…" \
                                       || log "upgrading $CUR_VER -> $NEW_VER…"

# Snapshot the CURRENT vs fetched dependency specs before we overwrite the tree,
# so pip_sync can install just the delta (new/changed lines) after the swap.
reqs_content "$APP_DIR"   > "$WORK/old_reqs"
reqs_content "$WORK/src"  > "$WORK/new_reqs"

# Strip the parts that must never land in the read-only app tree — mirror the
# exclusions the image build applies (Dockerfile.update) so the in-image layout
# stays identical to a freshly built image.
rm -rf \
  "$WORK/src/.git" \
  "$WORK/src/venv"* \
  "$WORK/src/.venv" \
  "$WORK/src/township_output" \
  "$WORK/src/offload" \
  "$WORK/src/dist" \
  "$WORK/src/.packaging-cache"
find "$WORK/src" -type d -name __pycache__ -prune -exec rm -rf '{}' + 2>/dev/null || true

# Replace the app tree in place. rsync --delete makes /opt/coderai/app exactly
# mirror the fetched (pruned) source, so files removed upstream also disappear.
# Fall back to a copy when rsync isn't present.
if command -v rsync >/dev/null 2>&1; then
  rsync -a --delete "$WORK/src/./" "$APP_DIR/./"
else
  log "rsync not found — using cp (stale removed-upstream files may linger)"
  find "$APP_DIR" -mindepth 1 -maxdepth 1 \
       ! -name models -exec rm -rf '{}' + 2>/dev/null || true
  cp -a "$WORK/src/." "$APP_DIR/"
fi

# Keep the entrypoint expectations intact.
mkdir -p "$APP_DIR/models"
[[ -f "$APP_DIR/coderai" ]] && chmod +x "$APP_DIR/coderai" || true

# Refresh launchers + service configs from the fetched tree (incl. this script).
# Non-fatal: a failed config copy must not abort an otherwise-successful upgrade.
sync_system_files || log "warning: launcher/config refresh incomplete"

# Dependency sync: install any new/changed deps (pip_sync no-ops when there are
# none). A failure aborts with a non-zero exit so the host does NOT commit a
# half-upgraded image (new code but a missing package) — the fetched code is
# simply discarded and the running image is left untouched.
if [[ "$SKIP_PIP" == "1" ]]; then
  log "dependency sync skipped (CODERAI_UPGRADE_SKIP_PIP=1)"
elif ! pip_sync; then
  die "dependency install failed — image left unchanged (fix connectivity/requirements and retry)"
fi

FINAL_VER="$(read_version "$APP_DIR/codai/__init__.py" || echo "$NEW_VER")"
log "done — in-image code is now $FINAL_VER"
exit 0
