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

# Safe upgrade workflow:
# 1) backup
# 2) upgrade dependencies/components
# 3) run healthcheck
# 4) optional rollback (config restore) on failure
#
# Usage:
#   ./upgrade.sh
#   ./upgrade.sh --model qwen2.5-coder:7b
#   ./upgrade.sh --skip-model-pull

WORKDIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
MODEL="${MODEL:-qwen2.5-coder:7b}"
SKIP_MODEL_PULL=0
BACKUP_DIR="${BACKUP_DIR:-$HOME/openclaude-backups}"
LAST_BACKUP=""

while [[ $# -gt 0 ]]; do
  case "$1" in
    --model)
      MODEL="$2"
      shift 2
      ;;
    --skip-model-pull)
      SKIP_MODEL_PULL=1
      shift
      ;;
    *)
      echo "Unknown argument: $1"
      exit 2
      ;;
  esac
done

step() { echo; echo "==> $*"; }
fail() { echo "ERROR: $*" >&2; exit 1; }

run_backup() {
  step "Creating backup before upgrade"
  pushd "$WORKDIR" >/dev/null
  ./backup.sh --output-dir "$BACKUP_DIR"
  popd >/dev/null
  LAST_BACKUP="$(ls -dt "$BACKUP_DIR"/backup_* 2>/dev/null | head -n 1 || true)"
  [[ -n "$LAST_BACKUP" ]] || fail "Could not find created backup directory"
  echo "Backup: $LAST_BACKUP"
}

rollback_configs() {
  if [[ -z "${LAST_BACKUP:-}" ]]; then
    echo "No backup found, cannot rollback."
    return
  fi
  step "Rolling back configs from backup"
  pushd "$WORKDIR" >/dev/null
  ./restore.sh --from "$LAST_BACKUP"
  popd >/dev/null
}

upgrade_system() {
  step "Updating apt packages"
  sudo apt update
  sudo apt install -y curl git docker.io docker-compose-plugin python3 python3-venv python3-pip
}

upgrade_openclaude() {
  step "Updating OpenClaude CLI"
  if command -v npm >/dev/null 2>&1; then
    npm install -g @gitlawb/openclaude
  else
    echo "npm not found, skipping OpenClaude update"
  fi
}

upgrade_ollama() {
  step "Ensuring Ollama is installed and running"
  if ! command -v ollama >/dev/null 2>&1; then
    curl -fsSL https://ollama.com/install.sh | sh
  fi
  sudo systemctl enable ollama
  sudo systemctl restart ollama
}

upgrade_models() {
  if [[ "$SKIP_MODEL_PULL" -eq 1 ]]; then
    step "Skipping model pull (--skip-model-pull)"
    return
  fi
  step "Refreshing model: $MODEL"
  ollama pull "$MODEL"
  ollama pull nomic-embed-text || true
}

upgrade_qdrant() {
  step "Refreshing Qdrant container"
  pushd "$WORKDIR" >/dev/null
  docker compose pull
  docker compose up -d
  popd >/dev/null
}

upgrade_python_deps() {
  step "Refreshing Python virtual environment dependencies"
  pushd "$WORKDIR" >/dev/null
  if [[ ! -d ".venv" ]]; then
    python3 -m venv .venv
  fi
  source .venv/bin/activate
  pip install --upgrade pip
  pip install -r requirements.txt
  deactivate || true
  popd >/dev/null
}

run_checks() {
  step "Running healthcheck (full)"
  pushd "$WORKDIR" >/dev/null
  ./healthcheck.sh --mode full --model "$MODEL"
  popd >/dev/null
}

main() {
  run_backup

  if ! {
    upgrade_system
    upgrade_openclaude
    upgrade_ollama
    upgrade_models
    upgrade_qdrant
    upgrade_python_deps
    run_checks
  }; then
    echo
    echo "Upgrade failed. Starting config rollback..."
    rollback_configs
    fail "Upgrade aborted. Config rollback executed from: $LAST_BACKUP"
  fi

  step "Upgrade completed successfully"
  echo "Backup kept at: $LAST_BACKUP"
  echo "Tip: run ./status.sh for quick post-upgrade overview."
}

main "$@"