#!/usr/bin/env bash
# -E 让 ERR trap 能在函数内部生效，保证不可逆区间的异常一定触发回滚
set -Eeuo pipefail

# ============================================
# Vaultwarden 安全升级脚本
#
# 设计原则：
#   1. 先拉镜像再动服务 —— 网络失败时服务零影响
#   2. 冷备份（停容器 + WAL checkpoint）—— 保证快照一致
#   3. 升级前后比对数据指纹 —— 数据异常自动回滚
#   4. 回滚不删数据 —— 旧数据目录改名保留，不 rm
#
# 用法:
#   bash upgrade.sh              升级到 GitHub 最新 release
#   bash upgrade.sh --check      只检查，不做任何改动
#   bash upgrade.sh --version 1.37.0
#   bash upgrade.sh --rollback   回滚到上次升级前的状态
#   bash upgrade.sh --yes        跳过交互确认（用于自动化）
# ============================================

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

# ===== 输出样式 =====
if [ -t 1 ]; then
    RED=$'\033[0;31m'; GREEN=$'\033[0;32m'; YELLOW=$'\033[1;33m'
    CYAN=$'\033[0;36m'; BOLD=$'\033[1m';    NC=$'\033[0m'
else
    RED=''; GREEN=''; YELLOW=''; CYAN=''; BOLD=''; NC=''
fi

log()  { echo -e "${GREEN}[✓]${NC} $*"; }
info() { echo -e "${CYAN}[i]${NC} $*"; }
warn() { echo -e "${YELLOW}[!]${NC} $*"; }
error(){ echo -e "${RED}[✗]${NC} $*" >&2; }
step() { echo ""; echo -e "${BOLD}${CYAN}━━━ $* ━━━${NC}"; }

# ===== 全局状态 =====
GITHUB_REPO="dani-garcia/vaultwarden"
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
STATE_FILE="$SCRIPT_DIR/.upgrade-state"
MODE="upgrade"          # upgrade | check | rollback
TARGET_VERSION=""
ASSUME_YES=0
SAFETY_STAGE="none"     # none → stopped → armed → done，决定异常退出时如何兜底
BACKUP_PATH=""
ROLLBACK_TAG=""
OLD_IMAGE_REF=""
COUNTS_BEFORE=""

# =============================================================
#  参数解析
# =============================================================
usage() {
    cat <<'EOF'
Vaultwarden 安全升级脚本

用法: bash upgrade.sh [选项]

选项:
  -c, --check           只做预检和版本对比，不做任何改动
  -v, --version <ver>   指定目标版本（如 1.37.0），默认取 GitHub 最新 release
  -r, --rollback        回滚到上次升级前的状态（读取 .upgrade-state）
  -y, --yes             跳过交互确认
  -h, --help            显示本帮助

示例:
  bash upgrade.sh --check          # 先看看会发生什么
  bash upgrade.sh                  # 正式升级
  bash upgrade.sh --rollback       # 升级后发现问题，回滚
EOF
}

parse_args() {
    while [ $# -gt 0 ]; do
        case "$1" in
            -c|--check)    MODE="check" ;;
            -r|--rollback) MODE="rollback" ;;
            -y|--yes)      ASSUME_YES=1 ;;
            -v|--version)
                shift
                [ $# -gt 0 ] || { error "--version 需要一个参数"; exit 1; }
                TARGET_VERSION="${1#v}"
                ;;
            -h|--help)     usage; exit 0 ;;
            *) error "未知参数: $1"; echo ""; usage; exit 1 ;;
        esac
        shift
    done
}

confirm() {
    [ "$ASSUME_YES" -eq 1 ] && return 0
    local prompt="$1" answer
    echo ""
    read -r -p "${prompt} [y/N] " answer
    [[ "$answer" =~ ^[Yy]$ ]]
}

# =============================================================
#  环境加载与预检
# =============================================================
load_env() {
    [ -f .env ] || { error ".env 不存在，请先运行 deploy.sh"; exit 1; }
    sed -i 's/\r$//' .env
    set -a; source .env; set +a

    DATA_DIR="${VAULTWARDEN_DATA_DIR:-/var/lib/vaultwarden}"
    BACKUP_BASE="${BACKUP_DIR:-/var/backups/vaultwarden}"
    PORT="${VAULTWARDEN_PORT:-8080}"
    DOMAIN="${VAULTWARDEN_DOMAIN:-}"
    CURRENT_IMAGE="${VAULTWARDEN_IMAGE:-vaultwarden/server:latest}"
    DB_FILE="$DATA_DIR/db.sqlite3"
}

preflight() {
    step "环境预检"

    [ "$(id -u)" -eq 0 ] || { error "需要 root 权限运行"; exit 1; }

    local missing=0
    for cmd in docker curl tar; do
        command -v "$cmd" >/dev/null 2>&1 || { error "缺少命令: $cmd"; missing=1; }
    done
    docker compose version >/dev/null 2>&1 || { error "docker compose 不可用"; missing=1; }
    [ "$missing" -eq 0 ] || exit 1

    [ -f docker-compose.yml ] || { error "docker-compose.yml 不存在"; exit 1; }
    [ -d "$DATA_DIR" ] || { error "数据目录不存在: $DATA_DIR"; exit 1; }
    [ -f "$DB_FILE" ] || { error "数据库不存在: $DB_FILE"; exit 1; }

    # 磁盘空间：要求可用空间 ≥ 数据目录的 3 倍（备份 + 解压余量），至少 500MB
    mkdir -p "$BACKUP_BASE"
    local data_kb avail_kb need_kb
    data_kb=$(du -sk "$DATA_DIR" | cut -f1)
    avail_kb=$(df -Pk "$BACKUP_BASE" | awk 'NR==2{print $4}')
    [ -n "$avail_kb" ] || { error "无法检测 $BACKUP_BASE 的可用空间"; exit 1; }
    need_kb=$(( data_kb * 3 ))
    [ "$need_kb" -lt 512000 ] && need_kb=512000
    if [ "$avail_kb" -lt "$need_kb" ]; then
        error "磁盘空间不足: 可用 $((avail_kb/1024))MB，需要 $((need_kb/1024))MB"
        exit 1
    fi

    log "root 权限、依赖命令、数据目录、磁盘空间 均正常"
    info "数据目录 $DATA_DIR ($((data_kb/1024))MB)，可用空间 $((avail_kb/1024))MB"
}

ensure_sqlite3() {
    command -v sqlite3 >/dev/null 2>&1 && return 0
    warn "sqlite3 未安装，正在安装（用于 WAL checkpoint 和完整性校验）..."
    if command -v apt-get >/dev/null 2>&1; then
        DEBIAN_FRONTEND=noninteractive apt-get update -qq
        DEBIAN_FRONTEND=noninteractive apt-get install -y -qq sqlite3
    elif command -v yum >/dev/null 2>&1; then
        yum install -y -q sqlite
    else
        error "无法自动安装 sqlite3，请手动安装后重试"
        exit 1
    fi
    command -v sqlite3 >/dev/null 2>&1 || { error "sqlite3 安装失败"; exit 1; }
    log "sqlite3 已安装"
}

# =============================================================
#  版本解析
# =============================================================
resolve_target_version() {
    if [ -n "$TARGET_VERSION" ]; then
        info "使用指定版本: $TARGET_VERSION"
        return 0
    fi
    info "查询 GitHub 最新 release..."
    TARGET_VERSION=$(curl -sf --max-time 20 \
        "https://api.github.com/repos/${GITHUB_REPO}/releases/latest" \
        | grep -o '"tag_name"[[:space:]]*:[[:space:]]*"[^"]*"' \
        | head -1 | sed 's/.*"\([^"]*\)"$/\1/' | sed 's/^v//') || true

    if [ -z "$TARGET_VERSION" ]; then
        error "无法从 GitHub 获取最新版本（网络问题？）"
        error "请手动指定: bash upgrade.sh --version 1.37.0"
        exit 1
    fi
    log "GitHub 最新 release: $TARGET_VERSION"
}

# 从 compose 项目推导容器 ID（而非硬编码容器名），停止状态也能拿到
compose_container_id() {
    docker compose ps -a -q vaultwarden 2>/dev/null | head -1 || true
}

# 当前容器实际运行的镜像 ID（sha256:...），拿不到则返回空
running_image_id() {
    local cid
    cid="$(compose_container_id)"
    [ -n "$cid" ] || return 0
    docker inspect "$cid" --format '{{.Image}}' 2>/dev/null || true
}

get_running_githash() {
    curl -sf --max-time 10 "http://127.0.0.1:${PORT}/api/config" 2>/dev/null \
        | grep -o '"gitHash"[[:space:]]*:[[:space:]]*"[^"]*"' \
        | sed 's/.*"\([^"]*\)"$/\1/' || true
}

# =============================================================
#  数据指纹
# =============================================================
# 统计关键表行数，用于升级前后比对。表不存在时输出 "-"
snapshot_counts() {
    local db="$1" t out=""
    for t in users ciphers folders sends attachments organizations collections; do
        local n
        n=$(sqlite3 "$db" "SELECT COUNT(*) FROM $t;" 2>/dev/null || echo "-")
        out+="${t}=${n} "
    done
    echo "${out% }"
}

print_counts() {
    local label="$1" counts="$2" kv
    echo "  ${label}:"
    for kv in $counts; do
        printf "    %-16s %s\n" "${kv%%=*}" "${kv##*=}"
    done
}

# =============================================================
#  备份
# =============================================================
cold_backup() {
    step "冷备份（服务已停止，快照一致）"

    mkdir -p "$BACKUP_BASE"

    info "执行 WAL checkpoint，将未落盘数据合并进主库..."
    sqlite3 "$DB_FILE" "PRAGMA wal_checkpoint(TRUNCATE);" >/dev/null
    log "WAL 已合并"

    info "校验数据库完整性..."
    local integrity
    integrity=$(sqlite3 "$DB_FILE" "PRAGMA integrity_check;" 2>&1 | head -1)
    if [ "$integrity" != "ok" ]; then
        error "数据库完整性检查未通过: $integrity"
        error "升级已中止，服务未改动。请先修复数据库。"
        docker compose up -d
        exit 1
    fi
    log "数据库完整性: ok"

    COUNTS_BEFORE="$(snapshot_counts "$DB_FILE")"
    print_counts "升级前数据统计" "$COUNTS_BEFORE"

    BACKUP_PATH="${BACKUP_BASE}/pre-upgrade-${TIMESTAMP}.tar.gz"
    info "打包整个数据目录 → $BACKUP_PATH"
    tar czf "$BACKUP_PATH" -C "$(dirname "$DATA_DIR")" "$(basename "$DATA_DIR")"

    # 校验备份包。注意：这里先把清单收进变量再用 here-string 匹配，
    # 不能写成 `tar tzf ... | grep -q`：grep -q 命中即退出会让 tar 收到
    # SIGPIPE，在 pipefail 下整条管道被判为失败，导致好备份被误判成坏的。
    info "校验备份包..."
    local listing
    listing="$(tar tzf "$BACKUP_PATH")" || { error "备份包损坏，无法列出内容"; exit 1; }
    grep -q 'db\.sqlite3$'  <<<"$listing" || { error "备份包中没有数据库文件"; exit 1; }
    grep -q 'rsa_key\.pem$' <<<"$listing" || warn "备份包中没有 rsa_key.pem"

    # 同时备份 .env（含 ADMIN_TOKEN）
    cp .env "${BACKUP_BASE}/env-${TIMESTAMP}.bak"
    chmod 600 "${BACKUP_BASE}/env-${TIMESTAMP}.bak"

    log "备份完成: $BACKUP_PATH ($(du -h "$BACKUP_PATH" | cut -f1))"
}

# =============================================================
#  镜像与容器
# =============================================================
tag_rollback_image() {
    step "锚定当前镜像（用于回滚）"

    local image_id
    image_id="$(running_image_id)"
    if [ -z "$image_id" ]; then
        image_id=$(docker image inspect "$CURRENT_IMAGE" --format '{{.Id}}' 2>/dev/null || true)
    fi
    [ -n "$image_id" ] || { error "找不到当前运行的镜像，无法建立回滚点"; exit 1; }

    ROLLBACK_TAG="vaultwarden/server:pre-upgrade-${TIMESTAMP}"
    docker tag "$image_id" "$ROLLBACK_TAG"
    OLD_IMAGE_REF="$CURRENT_IMAGE"

    log "当前镜像已标记为 $ROLLBACK_TAG"
    info "镜像 ID: ${image_id#sha256:}"
}

pull_target_image() {
    step "拉取目标镜像（此步骤不影响运行中的服务）"

    local target="vaultwarden/server:${TARGET_VERSION}"
    info "docker pull $target ..."
    if ! docker pull "$target"; then
        error "镜像拉取失败，服务未做任何改动，可安全重试"
        exit 1
    fi
    log "镜像已就绪: $target"
}

switch_image_in_env() {
    local new_ref="$1"
    cp .env ".env.bak-${TIMESTAMP}"
    chmod 600 ".env.bak-${TIMESTAMP}"
    if grep -q '^VAULTWARDEN_IMAGE=' .env; then
        sed -i "s|^VAULTWARDEN_IMAGE=.*|VAULTWARDEN_IMAGE=${new_ref}|" .env
    else
        printf '\nVAULTWARDEN_IMAGE=%s\n' "$new_ref" >> .env
    fi
    # 必须同步更新导出的变量：load_env 里 `set -a; source .env` 已经把旧值导出到
    # 本进程环境，而 docker compose 对环境变量的优先级高于 .env 文件 —— 只改文件
    # 的话 compose 仍会读到旧镜像，表现为「容器重建了但版本没变」。
    export VAULTWARDEN_IMAGE="$new_ref"
    info ".env 中 VAULTWARDEN_IMAGE → $new_ref"
}

wait_healthy() {
    local max_wait="${1:-60}" i
    info "等待服务就绪（最多 ${max_wait}s）..."
    for i in $(seq 1 "$max_wait"); do
        if curl -sf --max-time 3 "http://127.0.0.1:${PORT}/alive" >/dev/null 2>&1; then
            log "服务已就绪（${i}s）"
            return 0
        fi
        sleep 1
    done
    return 1
}

# =============================================================
#  验证
# =============================================================
# 确认容器真的跑在目标镜像上，防止「改了配置但容器没重建」这类静默失败
verify_running_image() {
    local expect="$1" cid actual_ref actual_id expect_id
    cid="$(compose_container_id)"
    [ -n "$cid" ] || { error "找不到容器"; return 1; }

    actual_ref=$(docker inspect "$cid" --format '{{.Config.Image}}' 2>/dev/null || true)
    actual_id=$(docker inspect "$cid" --format '{{.Image}}' 2>/dev/null || true)
    expect_id=$(docker image inspect "$expect" --format '{{.Id}}' 2>/dev/null || true)

    if [ "$actual_id" != "$expect_id" ]; then
        error "容器未运行在目标镜像上"
        error "  期望: $expect ($expect_id)"
        error "  实际: $actual_ref ($actual_id)"
        return 1
    fi
    log "容器镜像确认: $actual_ref"
    return 0
}

# $1: strict(默认) —— 要求新版 prelogin 路由存在（升级后）
#     lenient      —— 该路由 404 属正常（回滚到 1.36.0 之前的版本后）
verify_endpoints() {
    local strict="${1:-strict}"
    step "接口验证"
    local failed=0

    # 1. /alive
    if curl -sf --max-time 5 "http://127.0.0.1:${PORT}/alive" >/dev/null; then
        log "/alive 正常"
    else
        error "/alive 无响应"; failed=1
    fi

    # 2. 新版 prelogin 路由（1.36.0+ 才有，是本次升级的核心目标）
    local code
    code=$(curl -s -o /dev/null -w '%{http_code}' --max-time 8 \
        -X POST -H 'Content-Type: application/json' \
        -d '{"email":"healthcheck@example.invalid"}' \
        "http://127.0.0.1:${PORT}/identity/accounts/prelogin/password" || echo "000")
    if [ "$code" = "200" ]; then
        log "/identity/accounts/prelogin/password → 200（新版客户端可登录）"
    elif [ "$strict" = "lenient" ]; then
        info "/identity/accounts/prelogin/password → $code（该版本无此路由，回滚后属预期）"
    else
        error "/identity/accounts/prelogin/password → $code（期望 200）"; failed=1
    fi

    # 3. 旧版 prelogin 路由（兼容老客户端）
    code=$(curl -s -o /dev/null -w '%{http_code}' --max-time 8 \
        -X POST -H 'Content-Type: application/json' \
        -d '{"email":"healthcheck@example.invalid"}' \
        "http://127.0.0.1:${PORT}/identity/accounts/prelogin" || echo "000")
    if [ "$code" = "200" ]; then
        log "/identity/accounts/prelogin → 200（旧客户端兼容）"
    else
        warn "/identity/accounts/prelogin → $code"
    fi

    # 4. 经 Nginx 的外部访问（失败只告警，属于代理层问题，不触发回滚）
    if [ -n "$DOMAIN" ]; then
        code=$(curl -s -o /dev/null -w '%{http_code}' --max-time 10 \
            "https://${DOMAIN}/api/config" || echo "000")
        if [ "$code" = "200" ]; then
            log "https://${DOMAIN}/api/config → 200（Nginx 反代正常）"
        else
            warn "https://${DOMAIN}/api/config → $code（检查 Nginx/证书，不影响本次升级判定）"
        fi
    fi

    return $failed
}

verify_data_integrity() {
    step "数据完整性验证"

    local integrity
    integrity=$(sqlite3 "$DB_FILE" "PRAGMA integrity_check;" 2>&1 | head -1)
    if [ "$integrity" != "ok" ]; then
        error "升级后数据库完整性检查失败: $integrity"
        return 1
    fi
    log "数据库完整性: ok"

    local counts_after
    counts_after="$(snapshot_counts "$DB_FILE")"
    print_counts "升级后数据统计" "$counts_after"

    # 逐表比对：行数只允许持平或增加，减少视为数据丢失
    local kv table before after lost=0
    for kv in $COUNTS_BEFORE; do
        table="${kv%%=*}"; before="${kv##*=}"
        after=$(echo "$counts_after" | tr ' ' '\n' | grep "^${table}=" | cut -d= -f2 || true)
        # 升级前该表就不存在 / 升级后读不到 → 跳过比对
        if [ "$before" = "-" ] || [ -z "$after" ] || [ "$after" = "-" ]; then
            continue
        fi
        if [ "$after" -lt "$before" ]; then
            error "表 $table 行数减少: $before → $after"
            lost=1
        fi
    done

    if [ "$lost" -eq 1 ]; then
        error "检测到数据丢失！"
        return 1
    fi
    log "所有表行数无减少，数据完好"
    return 0
}

# =============================================================
#  回滚
# =============================================================
save_state() {
    cat > "$STATE_FILE" <<EOF
# Vaultwarden 升级状态记录 —— 由 upgrade.sh 自动生成
UPGRADE_TIMESTAMP=${TIMESTAMP}
BACKUP_PATH=${BACKUP_PATH}
ROLLBACK_TAG=${ROLLBACK_TAG}
OLD_IMAGE_REF=${OLD_IMAGE_REF}
NEW_IMAGE_REF=vaultwarden/server:${TARGET_VERSION}
DATA_DIR=${DATA_DIR}
ENV_BACKUP=${SCRIPT_DIR}/.env.bak-${TIMESTAMP}
EOF
    chmod 600 "$STATE_FILE"
}

# 恢复数据目录：旧目录改名保留，绝不 rm
restore_data_dir() {
    local backup_tar="$1" ts="$2"
    local parent failed_dir
    parent="$(dirname "$DATA_DIR")"
    failed_dir="${DATA_DIR}.failed-${ts}"

    if [ -d "$DATA_DIR" ]; then
        mv "$DATA_DIR" "$failed_dir"
        warn "升级后的数据目录已保留为: $failed_dir（未删除）"
    fi
    tar xzf "$backup_tar" -C "$parent"
    [ -f "$DB_FILE" ] || { error "恢复失败：$DB_FILE 不存在"; return 1; }
    log "数据目录已从备份恢复"
}

# 服务已停、但备份尚未就绪时的紧急恢复：此阶段 .env 未改动，直接拉起原版本即可
emergency_start_old() {
    step "紧急恢复"
    warn "升级在备份阶段中止，数据未被改动，正在拉起原服务..."
    docker compose up -d || true
    if wait_healthy 60; then
        log "原服务已恢复运行，本次升级未造成任何影响"
    else
        error "原服务恢复失败！请手动执行: cd $SCRIPT_DIR && docker compose up -d"
        error "查看日志: docker compose logs --tail 100"
    fi
}

# 统一的退出兜底：确保任何异常路径都不会把服务丢在停机状态
on_exit() {
    local code=$?
    trap - EXIT ERR
    [ "$code" -eq 0 ] && return 0
    case "$SAFETY_STAGE" in
        stopped) emergency_start_old ;;
        armed)   do_rollback "脚本异常中止（退出码 $code）" ;;
        *)       : ;;   # 尚未动过服务，无需处理
    esac
    return 0
}

do_rollback() {
    trap - ERR EXIT     # 先解除，避免回滚过程自身出错导致递归
    SAFETY_STAGE="done"
    set +e
    local reason="$1"
    step "自动回滚"
    error "触发原因: $reason"

    docker compose down || true
    restore_data_dir "$BACKUP_PATH" "$TIMESTAMP" || {
        error "自动恢复失败！请手动处理："
        error "  tar xzf $BACKUP_PATH -C $(dirname "$DATA_DIR")"
        exit 1
    }
    switch_image_in_env "$ROLLBACK_TAG"
    docker compose up -d --force-recreate

    if wait_healthy 60; then
        log "回滚完成，已恢复到升级前状态"
    else
        error "回滚后服务仍未就绪，请检查: docker compose logs --tail 100"
    fi
    echo ""
    warn "升级失败已回滚。备份保留在: $BACKUP_PATH"
    exit 1
}

manual_rollback() {
    step "手动回滚"

    [ -f "$STATE_FILE" ] || { error "找不到升级记录 $STATE_FILE，无法自动回滚"; exit 1; }
    # shellcheck disable=SC1090
    source "$STATE_FILE"
    DB_FILE="$DATA_DIR/db.sqlite3"   # 状态文件里的 DATA_DIR 可能与 .env 不同，重新推导

    echo ""
    info "将回滚到:"
    info "  镜像:   $ROLLBACK_TAG"
    info "  备份:   $BACKUP_PATH"
    info "  时间:   $UPGRADE_TIMESTAMP"
    echo ""

    [ -f "$BACKUP_PATH" ] || { error "备份文件不存在: $BACKUP_PATH"; exit 1; }
    docker image inspect "$ROLLBACK_TAG" >/dev/null 2>&1 \
        || { error "回滚镜像不存在: $ROLLBACK_TAG"; exit 1; }

    warn "回滚会丢弃升级后产生的所有数据变更（新增/修改的密码条目）"
    confirm "确认回滚？" || { info "已取消"; exit 0; }

    docker compose down || true
    restore_data_dir "$BACKUP_PATH" "$(date +%Y%m%d_%H%M%S)" || exit 1
    switch_image_in_env "$ROLLBACK_TAG"
    docker compose up -d --force-recreate

    if wait_healthy 60; then
        log "回滚完成"
        verify_endpoints lenient || true
    else
        error "服务未就绪，请检查: docker compose logs --tail 100"
        exit 1
    fi
}

# =============================================================
#  主流程
# =============================================================
show_plan() {
    local githash
    githash="$(get_running_githash)"

    echo ""
    echo -e "${BOLD}══════════ 升级计划 ══════════${NC}"
    echo "  域名:         ${DOMAIN:-未配置}"
    echo "  当前镜像:     $CURRENT_IMAGE"
    echo "  当前 gitHash: ${githash:-未知（服务未响应？）}"
    echo "  目标版本:     vaultwarden/server:${TARGET_VERSION}"
    echo "  数据目录:     $DATA_DIR"
    echo "  备份目录:     $BACKUP_BASE"
    echo ""
    echo "  执行步骤:"
    echo "    1. 锚定当前镜像为回滚点"
    echo "    2. 拉取新镜像（失败则零影响退出）"
    echo "    3. 停止容器"
    echo "    4. WAL checkpoint + 完整性校验 + 冷备份"
    echo "    5. 切换镜像并启动"
    echo "    6. 验证接口 + 比对数据行数"
    echo "    7. 任一验证失败 → 自动回滚"
    echo -e "${BOLD}═════════════════════════════${NC}"
}

main() {
    parse_args "$@"

    echo -e "${CYAN}${BOLD}"
    echo "  Vaultwarden 安全升级"
    echo -e "${NC}"

    load_env

    if [ "$MODE" = "rollback" ]; then
        [ "$(id -u)" -eq 0 ] || { error "需要 root 权限"; exit 1; }
        manual_rollback
        exit 0
    fi

    preflight
    resolve_target_version
    show_plan

    # 已是目标版本？
    local cur_image_id target_image_id
    cur_image_id="$(running_image_id)"
    target_image_id=$(docker image inspect "vaultwarden/server:${TARGET_VERSION}" \
                      --format '{{.Id}}' 2>/dev/null || true)
    if [ -n "$cur_image_id" ] && [ "$cur_image_id" = "$target_image_id" ]; then
        echo ""
        log "当前运行的已经是 ${TARGET_VERSION}，无需升级"
        if [ "$MODE" = "check" ]; then
            exit 0
        fi
        confirm "仍要强制重新部署一次？" || { info "已取消"; exit 0; }
    fi

    if [ "$MODE" = "check" ]; then
        echo ""
        info "--check 模式，未做任何改动"
        exit 0
    fi

    confirm "确认开始升级？" || { info "已取消"; exit 0; }

    ensure_sqlite3
    # 先拉镜像再打回滚标签：拉取失败时服务未动，也不留下无用的标签
    pull_target_image
    tag_rollback_image

    # ===== 从这里开始服务会中断，全程由 on_exit 兜底 =====
    trap on_exit EXIT
    trap 'exit 1' ERR

    step "停止容器"
    SAFETY_STAGE="stopped"      # 异常退出 → 直接拉起原版本
    docker compose stop
    log "容器已停止"

    cold_backup
    save_state
    SAFETY_STAGE="armed"        # 备份已就位，异常退出 → 完整回滚

    step "启动新版本"
    switch_image_in_env "vaultwarden/server:${TARGET_VERSION}"
    # 必须 --force-recreate：容器处于 stopped 状态时，docker compose up -d 只会
    # 把原容器重新 start，不会因为镜像变更而重建，结果是「升级了但还跑着旧镜像」。
    docker compose up -d --force-recreate

    if ! wait_healthy 90; then
        docker compose logs --tail 40 || true
        do_rollback "服务在 90 秒内未就绪"
    fi
    if ! verify_running_image "vaultwarden/server:${TARGET_VERSION}"; then
        do_rollback "容器未运行在目标镜像上"
    fi
    if ! verify_endpoints; then
        do_rollback "接口验证失败"
    fi
    if ! verify_data_integrity; then
        do_rollback "数据完整性验证失败"
    fi

    # ===== 成功 =====
    SAFETY_STAGE="done"
    trap - EXIT ERR
    step "升级成功"
    local new_githash
    new_githash="$(get_running_githash)"
    echo ""
    echo -e "${GREEN}${BOLD}══════════ 完成 ══════════${NC}"
    echo "  运行版本:   vaultwarden/server:${TARGET_VERSION}"
    echo "  gitHash:    ${new_githash:-未知}"
    echo "  备份文件:   $BACKUP_PATH"
    echo "  回滚镜像:   $ROLLBACK_TAG"
    echo ""
    echo "  如需回滚:   bash upgrade.sh --rollback"
    echo "  查看日志:   docker compose logs --tail 50"
    echo -e "${GREEN}${BOLD}══════════════════════════${NC}"
    echo ""
    warn "请登录 Web 端确认数据无误后，再考虑清理旧镜像:"
    warn "  docker image rm $ROLLBACK_TAG"
}

main "$@"