#!/usr/bin/env bash
# =====================================================================
# 02_setup_streaming_replication.sh
# ---------------------------------------------------------------------
# 用 docker 起一主一从 PG 16 流复制集群。
#
# 主库：pg_master，端口 5432
# 从库：pg_replica，端口 5433
#
# 默认账号：postgres / s3cret
# 复制账号：repl_user / r3pl_pwd
#
# 用法：
#   ./02_setup_streaming_replication.sh up        # 起集群
#   ./02_setup_streaming_replication.sh status    # 查复制状态
#   ./02_setup_streaming_replication.sh write     # 主库写入测试数据
#   ./02_setup_streaming_replication.sh read      # 从库验证可读
#   ./02_setup_streaming_replication.sh failover  # 演示 promote 从库
#   ./02_setup_streaming_replication.sh down      # 销毁
# =====================================================================

set -euo pipefail

PG_IMAGE="postgres:16"
NETWORK="pg_repl_net"
MASTER="pg_master"
REPLICA="pg_replica"
PGPASS="s3cret"
REPL_USER="repl_user"
REPL_PASS="r3pl_pwd"

cmd_up() {
    docker network inspect "$NETWORK" >/dev/null 2>&1 || docker network create "$NETWORK"

    echo "[1/5] 启动主库 $MASTER..."
    docker run -d --name "$MASTER" --network "$NETWORK" \
        -e POSTGRES_PASSWORD="$PGPASS" \
        -e POSTGRES_DB=learn_pg \
        -p 5432:5432 \
        "$PG_IMAGE" \
        -c wal_level=logical \
        -c max_wal_senders=10 \
        -c max_replication_slots=10 \
        -c hot_standby=on \
        -c synchronous_commit=on \
        -c wal_keep_size=128MB

    echo "[2/5] 等待主库就绪..."
    for i in {1..30}; do
        if docker exec "$MASTER" pg_isready -U postgres >/dev/null 2>&1; then break; fi
        sleep 1
    done

    echo "[3/5] 创建复制账号 + 修改 pg_hba.conf..."
    docker exec "$MASTER" psql -U postgres -c \
        "CREATE ROLE $REPL_USER WITH REPLICATION LOGIN PASSWORD '$REPL_PASS';"
    docker exec "$MASTER" bash -c \
        "echo 'host replication $REPL_USER 0.0.0.0/0 md5' >> /var/lib/postgresql/data/pg_hba.conf"
    docker exec "$MASTER" psql -U postgres -c "SELECT pg_reload_conf();"

    echo "[4/5] 在主库创建物理复制槽 + PUBLICATION..."
    docker exec "$MASTER" psql -U postgres -d learn_pg -c \
        "SELECT pg_create_physical_replication_slot('ch15_standby1_slot');"
    docker exec "$MASTER" psql -U postgres -d learn_pg -c \
        "CREATE TABLE IF NOT EXISTS ch15_demo (id BIGSERIAL PRIMARY KEY, msg TEXT, ts TIMESTAMPTZ DEFAULT now());"

    echo "[5/5] 启动从库 $REPLICA（pg_basebackup 自动克隆主库）..."
    docker run -d --name "$REPLICA" --network "$NETWORK" \
        -e POSTGRES_PASSWORD="$PGPASS" \
        -e PGPASSWORD="$REPL_PASS" \
        -p 5433:5432 \
        --entrypoint /bin/bash \
        "$PG_IMAGE" \
        -c "
            set -e
            rm -rf /var/lib/postgresql/data/*
            echo '正在做基础备份...'
            pg_basebackup -h $MASTER -p 5432 -U $REPL_USER \
                -D /var/lib/postgresql/data \
                -Fp -Xs -P -R -S ch15_standby1_slot
            echo '基础备份完成，启动从库...'
            chown -R postgres:postgres /var/lib/postgresql/data
            chmod 0700 /var/lib/postgresql/data
            su postgres -c 'postgres -D /var/lib/postgresql/data'
        "

    sleep 5
    echo
    echo "✅ 集群已起。下面查复制状态："
    cmd_status
}

cmd_status() {
    echo "--- 主库 pg_stat_replication ---"
    docker exec "$MASTER" psql -U postgres -d learn_pg -c "
        SELECT application_name, client_addr, state, sync_state,
               pg_wal_lsn_diff(pg_current_wal_lsn(), replay_lsn) AS lag_bytes
        FROM pg_stat_replication;
    "
    echo "--- 主库 pg_replication_slots ---"
    docker exec "$MASTER" psql -U postgres -d learn_pg -c "
        SELECT slot_name, slot_type, active, restart_lsn FROM pg_replication_slots;
    "
    echo "--- 从库角色 ---"
    docker exec "$REPLICA" psql -U postgres -d learn_pg -c \
        "SELECT pg_is_in_recovery() AS in_recovery, pg_last_wal_replay_lsn() AS replayed;" || true
}

cmd_write() {
    echo "--- 主库写入 5 行数据 ---"
    docker exec "$MASTER" psql -U postgres -d learn_pg -c \
        "INSERT INTO ch15_demo(msg) SELECT 'hello-' || g FROM generate_series(1,5) g RETURNING *;"
}

cmd_read() {
    echo "--- 从库查询数据（应能看到主库写入的行） ---"
    sleep 1
    docker exec "$REPLICA" psql -U postgres -d learn_pg -c \
        "SELECT * FROM ch15_demo ORDER BY id DESC LIMIT 10;"
    echo "--- 试试在从库写（应失败：read-only） ---"
    docker exec "$REPLICA" psql -U postgres -d learn_pg -c \
        "INSERT INTO ch15_demo(msg) VALUES('从库写入')" || true
}

cmd_failover() {
    echo "--- Step 1: 停掉主库 ---"
    docker stop "$MASTER"

    echo "--- Step 2: 在从库 pg_promote() ---"
    docker exec "$REPLICA" psql -U postgres -d learn_pg -c \
        "SELECT pg_promote(wait => true, wait_seconds => 30);"

    echo "--- Step 3: 验证从库变成主库 ---"
    docker exec "$REPLICA" psql -U postgres -d learn_pg -c \
        "SELECT pg_is_in_recovery() AS still_replica;"

    echo "--- Step 4: 在新主库写入 ---"
    docker exec "$REPLICA" psql -U postgres -d learn_pg -c \
        "INSERT INTO ch15_demo(msg) VALUES('写入新主库') RETURNING *;"
}

cmd_down() {
    docker rm -f "$MASTER" "$REPLICA" 2>/dev/null || true
    docker network rm "$NETWORK" 2>/dev/null || true
    echo "✅ 已清理。"
}

case "${1:-}" in
    up)        cmd_up ;;
    status)    cmd_status ;;
    write)     cmd_write ;;
    read)      cmd_read ;;
    failover)  cmd_failover ;;
    down)      cmd_down ;;
    *)
        echo "Usage: $0 {up|status|write|read|failover|down}"
        exit 1
        ;;
esac