#!/usr/bin/env bash
set -euo pipefail

usage() {
  cat <<'USAGE'
Usage:
  bash ./scripts/update_openclaw.sh [options]

Options:
  --channel <stable|beta|dev>   Update channel (default: stable)
  --tag <dist-tag|version>      One-off npm tag/version (example: beta, 2026.2.24)
  --skip-guard-stop             Do not stop local model guard LaunchAgents before upgrade
  --skip-guard-reinstall        Do not reinstall local model guard LaunchAgents after upgrade
  -h, --help                    Show this help

Examples:
  bash ./scripts/update_openclaw.sh
  bash ./scripts/update_openclaw.sh --channel beta
  bash ./scripts/update_openclaw.sh --tag 2026.2.24
USAGE
}

now_ts() { date '+%Y-%m-%d %H:%M:%S%z'; }
log()  { printf '[%s] [update-openclaw] %s\n' "$(now_ts)" "$*"; }
warn() { printf '[%s] [update-openclaw] WARN: %s\n' "$(now_ts)" "$*" >&2; }
err()  { printf '[%s] [update-openclaw] ERROR: %s\n' "$(now_ts)" "$*" >&2; }

require_cmd() {
  if ! command -v "$1" >/dev/null 2>&1; then
    err "Required command missing: $1"
    exit 1
  fi
}

run_with_timeout() {
  local timeout_s="$1"
  shift

  local start_ts elapsed pid
  start_ts="$(date +%s)"
  "$@" &
  pid=$!

  while kill -0 "$pid" 2>/dev/null; do
    elapsed="$(( $(date +%s) - start_ts ))"
    if (( elapsed >= timeout_s )); then
      warn "Command timed out after ${timeout_s}s: $*"
      kill -INT "$pid" 2>/dev/null || true
      sleep 2
      kill -TERM "$pid" 2>/dev/null || true
      sleep 1
      kill -KILL "$pid" 2>/dev/null || true
      wait "$pid" 2>/dev/null || true
      return 124
    fi
    sleep 1
  done

  wait "$pid"
}

extract_version() {
  local raw
  raw="$("$1" --version 2>&1 || true)"
  printf '%s\n' "$raw" | grep -Eo '[0-9]{4}\.[0-9]+\.[0-9]+' | head -n1
}

print_guard_status() {
  local label="$1"
  launchctl print "gui/$(id -u)/$label" 2>/dev/null | \
    grep -E 'state =|last exit code|runs =|path =|program =' || true
}

is_launchagent_loaded() {
  local label="$1"
  launchctl print "gui/$(id -u)/$label" >/dev/null 2>&1
}

is_gateway_launchagent_running() {
  local label="$1"
  launchctl print "gui/$(id -u)/$label" 2>/dev/null | grep -q 'state = running'
}

fix_gateway_launchagent_log_paths_for_external_state() {
  local label="$1"
  local plist="$HOME/Library/LaunchAgents/$label.plist"
  local state_dir_link_target=""

  if [[ ! -f "$plist" ]]; then
    return 0
  fi

  # launchd may fail with EX_CONFIG when stdout/stderr paths point into
  # ~/.openclaw on an external-volume symlink. Force local /tmp logs.
  if [[ -L "$HOME/.openclaw" ]]; then
    state_dir_link_target="$(readlink "$HOME/.openclaw" || true)"
  fi
  if [[ "$state_dir_link_target" != /Volumes/* ]]; then
    return 0
  fi

  if ! command -v /usr/libexec/PlistBuddy >/dev/null 2>&1; then
    warn "PlistBuddy not found; skipping LaunchAgent log-path patch."
    return 0
  fi

  local desired_out="/tmp/openclaw-gateway.launchd.log"
  local desired_err="/tmp/openclaw-gateway.launchd.err.log"
  local current_out=""
  local current_err=""

  current_out="$(/usr/libexec/PlistBuddy -c 'Print :StandardOutPath' "$plist" 2>/dev/null || true)"
  current_err="$(/usr/libexec/PlistBuddy -c 'Print :StandardErrorPath' "$plist" 2>/dev/null || true)"

  if [[ "$current_out" == "$desired_out" && "$current_err" == "$desired_err" ]]; then
    return 0
  fi

  log "Patching gateway LaunchAgent logs to local /tmp paths (external state dir detected)"
  /usr/libexec/PlistBuddy -c "Set :StandardOutPath $desired_out" "$plist"
  /usr/libexec/PlistBuddy -c "Set :StandardErrorPath $desired_err" "$plist"

  launchctl bootout "gui/$(id -u)/$label" 2>/dev/null || true
  launchctl bootstrap "gui/$(id -u)" "$plist"
  launchctl kickstart -k "gui/$(id -u)/$label"
}

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ROOT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"

CHANNEL="stable"
TAG=""
STOP_GUARDS="true"
REINSTALL_GUARDS="true"
UPDATE_MAX_SECONDS="${UPDATE_MAX_SECONDS:-600}"

while (( "$#" )); do
  case "$1" in
    --channel)
      shift
      if [[ $# -eq 0 ]]; then
        err "--channel requires a value"
        usage
        exit 1
      fi
      CHANNEL="$1"
      ;;
    --tag)
      shift
      if [[ $# -eq 0 ]]; then
        err "--tag requires a value"
        usage
        exit 1
      fi
      TAG="$1"
      ;;
    --skip-guard-stop)
      STOP_GUARDS="false"
      ;;
    --skip-guard-reinstall)
      REINSTALL_GUARDS="false"
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      err "Unknown argument: $1"
      usage
      exit 1
      ;;
  esac
  shift
done

case "$CHANNEL" in
  stable|beta|dev) ;;
  *)
    err "Invalid --channel '$CHANNEL'. Expected stable, beta, or dev."
    exit 1
    ;;
esac

if ! [[ "$UPDATE_MAX_SECONDS" =~ ^[0-9]+$ ]] || [[ "$UPDATE_MAX_SECONDS" -lt 1 ]]; then
  err "UPDATE_MAX_SECONDS must be a positive integer."
  exit 1
fi

require_cmd openclaw
require_cmd launchctl
require_cmd cp
require_cmd mkdir
require_cmd date

if ! command -v npm >/dev/null 2>&1; then
  warn "npm not found; fallback npm upgrade path will be unavailable."
fi

ts="$(date +%Y%m%d-%H%M%S)"
backup_dir="$HOME/.openclaw/backups/update-$ts"
mkdir -p "$backup_dir"

config_file="$HOME/.openclaw/openclaw.json"
sessions_file="$HOME/.openclaw/agents/main/sessions/sessions.json"

log "Backing up OpenClaw files to: $backup_dir"
if [[ -f "$config_file" ]]; then
  cp "$config_file" "$backup_dir/openclaw.json"
else
  warn "Config not found: $config_file"
fi
if [[ -f "$sessions_file" ]]; then
  cp "$sessions_file" "$backup_dir/sessions.json"
else
  warn "Sessions file not found: $sessions_file"
fi

if [[ "$STOP_GUARDS" == "true" ]]; then
  log "Stopping local model guard LaunchAgents to avoid config races during update"
  launchctl bootout "gui/$(id -u)/ai.openclaw.local.model-budget-guard" 2>/dev/null || true
  launchctl bootout "gui/$(id -u)/ai.openclaw.local.model-schedule-guard" 2>/dev/null || true
fi

log "Stopping OpenClaw gateway before update"
if ! openclaw gateway stop; then
  warn "Gateway stop returned non-zero (it may already be stopped). Continuing."
fi

current_version="$(extract_version openclaw || true)"
if [[ -n "$current_version" ]]; then
  log "Current OpenClaw version: $current_version"
else
  warn "Could not parse current OpenClaw version from 'openclaw --version'."
fi

if command -v npm >/dev/null 2>&1; then
  latest_npm_version="$(npm view openclaw version 2>/dev/null || true)"
  if [[ -n "$latest_npm_version" ]]; then
    log "Latest npm OpenClaw version: $latest_npm_version"
  fi
fi

log "Running doctor fix for schema/config migration"
if ! openclaw doctor --fix --non-interactive; then
  warn "Doctor fix returned non-zero. Continuing with upgrade."
fi

update_cmd=(openclaw update --yes --no-restart --channel "$CHANNEL")
if [[ -n "$TAG" ]]; then
  update_cmd+=(--tag "$TAG")
fi

log "Running OpenClaw updater: ${update_cmd[*]}"
if ! run_with_timeout "$UPDATE_MAX_SECONDS" "${update_cmd[@]}"; then
  warn "Built-in updater failed."
  if ! command -v npm >/dev/null 2>&1; then
    err "npm not available for fallback update."
    exit 1
  fi

  pkg="openclaw@latest"
  if [[ -n "$TAG" ]]; then
    pkg="openclaw@$TAG"
  fi

  log "Falling back to npm global install: npm install -g $pkg"
  npm install -g "$pkg"
fi

new_version="$(extract_version openclaw || true)"
if [[ -n "$new_version" ]]; then
  log "Updated OpenClaw version: $new_version"
else
  warn "Could not parse updated version from 'openclaw --version'."
fi

log "Starting OpenClaw gateway after update"
if ! openclaw gateway start; then
  err "Gateway failed to start after update."
  exit 1
fi

# Some OpenClaw versions return success from "gateway start" even when the
# LaunchAgent is not installed/loaded; enforce a service-managed gateway.
gateway_label="ai.openclaw.gateway"
if ! launchctl print "gui/$(id -u)/$gateway_label" >/dev/null 2>&1; then
  warn "Gateway LaunchAgent is not loaded after start; installing and retrying."
  if ! openclaw gateway install --force; then
    err "Failed to install gateway service."
    exit 1
  fi
  if ! openclaw gateway start; then
    err "Gateway failed to start after installing service."
    exit 1
  fi
  if ! launchctl print "gui/$(id -u)/$gateway_label" >/dev/null 2>&1; then
    err "Gateway LaunchAgent is still not loaded after install/start."
    exit 1
  fi
fi

fix_gateway_launchagent_log_paths_for_external_state "$gateway_label"

if ! is_gateway_launchagent_running "$gateway_label"; then
  warn "Gateway LaunchAgent is loaded but not running; attempting kickstart."
  launchctl kickstart -k "gui/$(id -u)/$gateway_label" || true
  sleep 2
fi
if ! is_gateway_launchagent_running "$gateway_label"; then
  err "Gateway LaunchAgent is still not running."
  launchctl print "gui/$(id -u)/$gateway_label" | grep -E 'state =|last exit code|runs =|path =|program =' || true
  exit 1
fi

log "Running post-upgrade health checks"
openclaw doctor --non-interactive || true
openclaw gateway status || true
openclaw models status || true

if [[ "$REINSTALL_GUARDS" == "true" ]]; then
  installer="$SCRIPT_DIR/install_local_model_guardrails.sh"
  if [[ -x "$installer" ]]; then
    log "Reinstalling local model guard LaunchAgents"
    bash "$installer"

    # Run both periodic guards immediately after reinstall so profile/budget
    # state is refreshed in this update run (not only on the next interval).
    launchctl kickstart -k "gui/$(id -u)/ai.openclaw.local.model-schedule-guard" 2>/dev/null || true
    launchctl kickstart -k "gui/$(id -u)/ai.openclaw.local.model-budget-guard" 2>/dev/null || true
    sleep 2

    if ! is_launchagent_loaded "ai.openclaw.local.model-schedule-guard"; then
      err "Schedule guard LaunchAgent is not loaded after reinstall."
      exit 1
    fi
    if ! is_launchagent_loaded "ai.openclaw.local.model-budget-guard"; then
      err "Budget guard LaunchAgent is not loaded after reinstall."
      exit 1
    fi
  else
    warn "Guard installer missing or not executable: $installer"
  fi
fi

log "Local guard status summary"
print_guard_status "ai.openclaw.local.model-budget-guard"
print_guard_status "ai.openclaw.local.model-schedule-guard"

if [[ -x "$SCRIPT_DIR/model_profile_switch.sh" ]]; then
  log "Profile status"
  bash "$SCRIPT_DIR/model_profile_switch.sh" status || true
fi

cat <<MSG

Upgrade complete.
- Backup directory: $backup_dir
- Channel: $CHANNEL
- Tag override: ${TAG:-none}
- Current version: ${new_version:-unknown}

If needed, inspect logs:
  tail -n 80 /tmp/openclaw-model-budget-guard.log /tmp/openclaw-model-budget-guard.err.log
  tail -n 80 /tmp/openclaw-model-schedule-guard.log /tmp/openclaw-model-schedule-guard.err.log
MSG