#!/usr/bin/env bash
# restore.sh —— 灾备恢复脚本
# 用途：从 S3 下载 → age 解密 → 恢复 MySQL / PostgreSQL
# 使用方式：
#   restore.sh mysql full 20240101_020000   # 恢复 MySQL 指定全量
#   restore.sh mysql inc  20240101_021500   # 在全量基础上应用增量
#   restore.sh postgres   blog 20240101_021500
#   restore.sh marzban    20240101_020000
#
# 注意：MySQL 恢复需要先停止 MySQL 服务（由运维手动执行）
#       此脚本只做"下载 + 解密 + 准备"，不自动重启 MySQL 服务
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
source "${SCRIPT_DIR}/lib.sh"

# age 私钥（恢复时必须提供）
AGE_PRIVATE_KEY_FILE="${AGE_PRIVATE_KEY_FILE:-/etc/age-private.txt}"

if [ ! -f "${AGE_PRIVATE_KEY_FILE}" ]; then
    log_error "age 私钥文件不存在: ${AGE_PRIVATE_KEY_FILE}"
    log_error "恢复需要私钥。请将私钥文件挂载到容器，或设置 AGE_PRIVATE_KEY 环境变量"
    exit 1
fi

RESTORE_DIR="${RESTORE_DIR:-/restore}"
mkdir -p "${RESTORE_DIR}"

# ---------- age 解密 ----------
age_decrypt() {
    local src="$1"
    local dst="$2"
    age --decrypt --identity "${AGE_PRIVATE_KEY_FILE}" -o "${dst}" "${src}"
}

# ---------- 从 S3 下载并解密 ----------
download_and_decrypt() {
    local s3_key="$1"
    local local_enc="${RESTORE_DIR}/$(basename "${s3_key}")"
    local local_dec="${local_enc%.age}"

    log_info "从 S3 下载: s3://${BACKUP_S3_BUCKET}/${BACKUP_S3_PREFIX}/${s3_key}.age"
    AWS_PROFILE="${AWS_PROFILE}" aws s3 cp \
        "s3://${BACKUP_S3_BUCKET}/${BACKUP_S3_PREFIX}/${s3_key}.age" \
        "${local_enc}"

    log_info "age 解密: ${local_enc}"
    age_decrypt "${local_enc}" "${local_dec}"
    rm -f "${local_enc}"

    echo "${local_dec}"
}

# ──────────────────────────────────────────────────────
# MySQL 恢复
# ──────────────────────────────────────────────────────
restore_mysql_full() {
    local timestamp="$1"
    local s3_key="mysql/full/${timestamp}.tar.gz"
    local restore_base="${RESTORE_DIR}/mysql"
    mkdir -p "${restore_base}"

    log_info "[MySQL] 下载并解密全量备份: ${timestamp}"
    local archive
    archive="$(download_and_decrypt "${s3_key}")"

    log_info "[MySQL] 解压: ${archive}"
    tar -xzf "${archive}" -C "${restore_base}"
    rm -f "${archive}"

    local backup_dir="${restore_base}/${timestamp}"
    log_info "[MySQL] xtrabackup --prepare（应用增量前的全量 prepare）"
    xtrabackup --prepare --apply-log-only --target-dir="${backup_dir}"

    log_ok "[MySQL] 全量备份已解压并 prepare 到: ${backup_dir}"
    echo "${backup_dir}"
}

restore_mysql_inc() {
    local base_dir="$1"
    local timestamp="$2"
    local s3_key="mysql/inc/${timestamp}.tar.gz"
    local restore_base="${RESTORE_DIR}/mysql"

    log_info "[MySQL] 下载并解密增量备份: ${timestamp}"
    local archive
    archive="$(download_and_decrypt "${s3_key}")"

    tar -xzf "${archive}" -C "${restore_base}"
    rm -f "${archive}"

    local inc_dir="${restore_base}/${timestamp}"
    log_info "[MySQL] xtrabackup --prepare（应用增量 ${timestamp}）"
    xtrabackup --prepare --apply-log-only \
        --target-dir="${base_dir}" \
        --incremental-dir="${inc_dir}"
    rm -rf "${inc_dir}"

    log_ok "[MySQL] 增量 ${timestamp} 已应用到 ${base_dir}"
}

restore_mysql_finalize() {
    local base_dir="$1"
    log_info "[MySQL] xtrabackup --prepare（最终 redo log 回放）"
    xtrabackup --prepare --target-dir="${base_dir}"
    log_ok "[MySQL] 备份已可用，恢复步骤："
    log_ok "  1. 停止 MySQL：systemctl stop mysql"
    log_ok "  2. 清空数据目录：rm -rf /var/lib/mysql/*"
    log_ok "  3. 复制备份：xtrabackup --copy-back --target-dir=${base_dir}"
    log_ok "  4. 修改权限：chown -R mysql:mysql /var/lib/mysql"
    log_ok "  5. 启动 MySQL：systemctl start mysql"
}

# ──────────────────────────────────────────────────────
# PostgreSQL 恢复
# ──────────────────────────────────────────────────────
restore_postgres() {
    local db="$1"
    local timestamp="$2"
    local s3_key="postgres/inc/${db}/${timestamp}.pgdump"
    local restore_target="${PG_RESTORE_TARGET:-${db}_restored}"

    log_info "[PostgreSQL] 下载并解密: ${db} @ ${timestamp}"
    local dump_file
    dump_file="$(download_and_decrypt "${s3_key}")"

    log_info "[PostgreSQL] 恢复数据库: ${restore_target}"
    PGPASSWORD="${PG_PASSWORD:-}" psql \
        --host="${PG_HOST}" \
        --port="${PG_PORT}" \
        --username="${PG_USER:-postgres}" \
        -c "CREATE DATABASE \"${restore_target}\";" 2>/dev/null || true

    PGPASSWORD="${PG_PASSWORD:-}" pg_restore \
        --host="${PG_HOST}" \
        --port="${PG_PORT}" \
        --username="${PG_USER:-postgres}" \
        --dbname="${restore_target}" \
        --no-password \
        --verbose \
        "${dump_file}"

    rm -f "${dump_file}"
    log_ok "[PostgreSQL] 恢复完成: ${restore_target}"
}

# ──────────────────────────────────────────────────────
# 主入口
# ──────────────────────────────────────────────────────
case "${1:-}" in
    mysql)
        case "${2:-}" in
            full)
                base_dir="$(restore_mysql_full "${3}")"
                restore_mysql_finalize "${base_dir}"
                ;;
            inc)
                if [ -z "${4:-}" ]; then
                    log_error "用法: restore.sh mysql inc <base_timestamp> <inc_timestamp>"
                    exit 1
                fi
                base_dir="$(restore_mysql_full "${3}")"
                restore_mysql_inc "${base_dir}" "${4}"
                restore_mysql_finalize "${base_dir}"
                ;;
            *)
                echo "用法: restore.sh mysql {full|inc} <timestamp> [inc_timestamp]"
                exit 1
                ;;
        esac
        ;;
    postgres)
        restore_postgres "${2}" "${3}"
        ;;
    marzban)
        marzban_archive="$(download_and_decrypt "marzban/marzban_${2}.sqlite3.gz")"
        log_ok "[Marzban] 解密文件: ${marzban_archive}"
        ;;
    *)
        echo "用法: restore.sh {mysql|postgres|marzban} ..."
        echo "  restore.sh mysql full <timestamp>"
        echo "  restore.sh mysql inc  <full_timestamp> <inc_timestamp>"
        echo "  restore.sh postgres   <db_name> <timestamp>"
        echo "  restore.sh marzban    <timestamp>"
        exit 1
        ;;
esac