#!/usr/bin/env bash
# ============================================================
# 第 14 章 · 演示 3：完整 PITR 流程脚本
# ------------------------------------------------------------
# 这是一个「教学用」端到端 PITR 演练脚本，会在「同一台主机」上
# 启动一个独立 PG 实例，做：
#   1) 配置 archive_mode + archive_command
#   2) 跑业务（写入数据 → 模拟事故 DROP TABLE）
#   3) 用 pg_basebackup 拉一次基础备份
#   4) 继续业务，产生新 WAL
#   5) 模拟事故时间点 T_BAD
#   6) 用 PITR 恢复到 T_BAD - 1s
#
# ⚠️ 重要：脚本需要的环境
#   - 必须以 postgres OS 用户（或有写权限的用户）运行
#   - 需要 initdb / pg_ctl / pg_basebackup 在 PATH 中
#   - 需要写入临时目录的权限（默认 /tmp/pitr_demo）
#   - 不会影响你的主 PG 实例（用独立端口 5499）
#
# 用法：
#   bash 03_basebackup_pitr.sh
#
# 如果你不想真实跑（比如没有 root/initdb 权限），
# 可以只阅读注释，理解 PITR 全流程。
# ============================================================
set -euo pipefail

WORK_DIR=${WORK_DIR:-/tmp/pitr_demo}
PGDATA="$WORK_DIR/data"
ARCHIVE="$WORK_DIR/archive"
BACKUP="$WORK_DIR/backup"
PORT=${PITR_PORT:-5499}

bar() { echo -e "\n============================================================\n$1\n============================================================"; }

cleanup() {
  bar "清理：停掉演示实例并删除目录"
  pg_ctl stop -D "$PGDATA" -m immediate >/dev/null 2>&1 || true
  rm -rf "$WORK_DIR"
}

trap 'echo "[!] 出错，请查看 $PGDATA/log/ 下日志"' ERR

# ------------------------------------------------------------
# Step 0：彻底清理上一次留下的环境
# ------------------------------------------------------------
cleanup
mkdir -p "$ARCHIVE" "$BACKUP"

# ------------------------------------------------------------
# Step 1：initdb 创建一个新实例
# ------------------------------------------------------------
bar "Step 1：initdb 创建新实例（端口 $PORT）"
initdb -D "$PGDATA" --auth-local=trust --auth-host=trust -U postgres -E UTF8 --locale=C

# 改 postgresql.conf：开归档、用本地 socket、监听独立端口
cat >> "$PGDATA/postgresql.conf" <<EOF

# === PITR 演示配置 ===
port = $PORT
unix_socket_directories = '$WORK_DIR'
listen_addresses = ''

wal_level = replica
archive_mode = on
archive_command = 'test ! -f $ARCHIVE/%f && cp %p $ARCHIVE/%f'
archive_timeout = 30
wal_keep_size = 64MB

logging_collector = on
log_directory = 'log'
log_filename = 'postgres-%Y-%m-%d.log'
log_min_messages = info
EOF

mkdir -p "$PGDATA/log"

# ------------------------------------------------------------
# Step 2：启动实例，建库 + 跑些初始数据
# ------------------------------------------------------------
bar "Step 2：启动实例，建库写初始数据"
pg_ctl start -D "$PGDATA" -l "$PGDATA/log/start.log" -w

PSQL="psql -h $WORK_DIR -p $PORT -U postgres"
$PSQL -d postgres -c "CREATE DATABASE pitr_demo;"
$PSQL -d pitr_demo -c "
  CREATE TABLE money_log (
    id BIGSERIAL PRIMARY KEY,
    msg TEXT,
    amount NUMERIC,
    created_at TIMESTAMPTZ DEFAULT now()
  );
  INSERT INTO money_log(msg, amount)
  SELECT '初始数据 ' || g, 100*g
  FROM generate_series(1,5) g;
"
$PSQL -d pitr_demo -c "SELECT count(*) AS init_rows FROM money_log;"

# ------------------------------------------------------------
# Step 3：基础备份
# ------------------------------------------------------------
bar "Step 3：pg_basebackup 拉一次基础备份"
rm -rf "$BACKUP"/*
pg_basebackup -h "$WORK_DIR" -p "$PORT" -U postgres \
              -D "$BACKUP" -Ft -X stream -P -c fast --label='pitr_demo_base'
echo "备份产物："
ls -lh "$BACKUP"

# ------------------------------------------------------------
# Step 4：继续业务，写一些「重要数据」
# ------------------------------------------------------------
bar "Step 4：备份后继续写入「重要数据」（10 行）"
$PSQL -d pitr_demo -c "
  INSERT INTO money_log(msg, amount)
  SELECT '重要交易 ' || g, 1000*g
  FROM generate_series(1,10) g;
"
$PSQL -d pitr_demo -c "SELECT count(*) AS after_backup_rows FROM money_log;"

# 强制切一段 WAL 触发归档
$PSQL -d pitr_demo -c "SELECT pg_switch_wal();"
sleep 2

# ------------------------------------------------------------
# Step 5：记录「事故时间点」 T_BAD
# ------------------------------------------------------------
sleep 1
T_BAD=$($PSQL -d pitr_demo -tA -c "SELECT now()")
echo "📌 记录事故前时间 T_BAD = $T_BAD"
sleep 2

# ------------------------------------------------------------
# Step 6：模拟事故 —— DROP TABLE
# ------------------------------------------------------------
bar "Step 6：模拟事故：DROP TABLE money_log"
$PSQL -d pitr_demo -c "DROP TABLE money_log;"
$PSQL -d pitr_demo -c "\dt"

# 强切 WAL 把 DROP 这条记录归档
$PSQL -d pitr_demo -c "SELECT pg_switch_wal();"
sleep 2

# ------------------------------------------------------------
# Step 7：停库，准备 PITR
# ------------------------------------------------------------
bar "Step 7：停掉实例，开始 PITR 恢复到 $T_BAD"
pg_ctl stop -D "$PGDATA" -m fast -w

# 备份原数据目录便于对比
mv "$PGDATA" "${PGDATA}.broken"

# 解压基础备份到新 PGDATA
mkdir -p "$PGDATA/pg_wal"
tar -xf "$BACKUP/base.tar"   -C "$PGDATA"
tar -xf "$BACKUP/pg_wal.tar" -C "$PGDATA/pg_wal/"

# 配置恢复
cat >> "$PGDATA/postgresql.auto.conf" <<EOF

# === PITR 恢复配置 ===
restore_command          = 'cp $ARCHIVE/%f %p'
recovery_target_time     = '$T_BAD'
recovery_target_action   = 'promote'
EOF

# 创建 recovery.signal（PG 12+ 用法，替代 recovery.conf）
touch "$PGDATA/recovery.signal"

mkdir -p "$PGDATA/log"

# ------------------------------------------------------------
# Step 8：启动并观察 recovery
# ------------------------------------------------------------
bar "Step 8：启动 PG，自动重放 WAL 直到 T_BAD"
pg_ctl start -D "$PGDATA" -l "$PGDATA/log/start.log" -w

echo "---- 启动日志（最后 30 行） ----"
tail -n 30 "$PGDATA/log"/postgres-*.log || true

# 等 promote 完成
sleep 3

bar "Step 9：验证恢复结果"
echo "money_log 应该已经回来，行数应该是 5（初始）+ 10（重要） = 15："
$PSQL -d pitr_demo -c "SELECT count(*) AS recovered_rows FROM money_log;"
$PSQL -d pitr_demo -c "SELECT id, msg, amount FROM money_log ORDER BY id LIMIT 5;"

bar "🎉 PITR 演示完成"
echo "  - 事故前时间 T_BAD : $T_BAD"
echo "  - WAL 归档目录     : $ARCHIVE （内有 $(ls -1 $ARCHIVE | wc -l) 个 WAL）"
echo "  - 旧坏 PGDATA     : ${PGDATA}.broken （已保留供对比）"
echo
echo "若要清理本次演示："
echo "    pg_ctl stop -D $PGDATA -m fast"
echo "    rm -rf $WORK_DIR"