#!/usr/bin/env bash

# 使用方式：
#   ./merge_db_restore /absolute/path/to/merge_setup.ini DB_NAME_POSTFIX <AccountDB/GameDB/WorldDB>
#
# 流程：
#   1. source load_merge_db_setup
#   2. load_merge_db_setup <merge_setup.ini>
#   3. 依參數 3 決定要還原哪些 DB：
#        - 無參數 3：AccountDB + GameDB + 全部 WorldDB
#        - AccountDB：只還原 AccountDB
#        - GameDB：只還原 GameDB
#        - WorldDB：只還原所有 WorldDB
#   4. 對每一個要還原的 DB：
#        - 檢查備份 DB "<DB_NAME><POSTFIX>" 是否存在
#        - 若存在：
#            a. 踢掉原 DB 的連線並 DROP DATABASE IF EXISTS "<DB_NAME>"
#            b. CREATE DATABASE "<DB_NAME>" TEMPLATE "<DB_NAME><POSTFIX>";
#
# 依賴：
#   - load_merge_db_setup 內已實作：
#       load_merge_db_setup
#       get_accountdb_value
#       get_gamedb_value
#       get_worlddb_value
#       world_order (依 setup 順序的 WorldID array)
#
#   - 每個 DB 的設定提供：
#       AccountDBIP / AccountDBName / AccountDBUser / AccountDBPW / (AccountDBPort)
#       GameDBIP / GameDB / GameDBUser / GameDBPassword / (GameDBPort)
#       WorldDBIP / WorldDBName / WorldDBUser / WorldDBPassword / (WorldDBPort)

set -u

if [[ $# -lt 2 ]]; then
    echo "Usage: $0 /absolute/path/to/merge_setup.ini DB_NAME_POSTFIX <AccountDB|GameDB|WorldDB>" >&2
    exit 1
fi

SETUP_FILE="$1"
DB_POSTFIX="$2"
TARGET_KIND="${3-}"    # 可能不存在，必須用 ${3-}

# 強制要求 setup 檔案用絕對路徑
if [[ "${SETUP_FILE:0:1}" != "/" ]]; then
    echo "請使用 setup 檔案的絕對路徑（目前：$SETUP_FILE）" >&2
    exit 1
fi

if [[ ! -f "$SETUP_FILE" ]]; then
    echo "找不到 setup 檔案：$SETUP_FILE" >&2
    exit 1
fi

# 驗證第三個參數（若有提供）
if [[ -n "$TARGET_KIND" ]]; then
    case "$TARGET_KIND" in
        AccountDB|GameDB|WorldDB)
            # 合法
            ;;
        *)
            echo "第三個參數若提供，只能是：AccountDB / GameDB / WorldDB" >&2
            echo "實際輸入：$TARGET_KIND" >&2
            exit 1
            ;;
    esac
fi

# 取得目前腳本所在目錄，假設 load_merge_db_setup 跟這支在同一個目錄
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# 引入 load_merge_db_setup
# shellcheck source=/dev/null
source "$SCRIPT_DIR/load_merge_db_setup"

# 載入設定
if ! load_merge_db_setup "$SETUP_FILE"; then
    echo "載入設定檔失敗：$SETUP_FILE" >&2
    exit 1
fi

########################################
# 共用：如 DB 有連線，才重啟遠端 PostgreSQL
########################################
# 參數：HOST SRC_DB USER PASS [PORT]
# 邏輯：
#   1. 用 pg_stat_activity 查 datname = SRC_DB 的連線數
#   2. 若連線數 > 0，則：
#        - 查遠端 psql 版本（主版號）
#        - sudo systemctl restart postgresql@${PG_VERSION}-main.service
#   3. 若連線數 = 0，則略過 restart
pg_check_and_restart_if_in_use() {
    local host="${1-}"
    local src_db="${2-}"
    local user="${3-}"
    local pass="${4-}"
    local port="${5-}"

    if [[ -z "$host" || -z "$src_db" || -z "$user" ]]; then
        echo "[NG]  pg_check_and_restart_if_in_use 參數不完整 (host='$host', db='$src_db', user='$user')" >&2
        return 1
    fi

    if [[ -z "$port" ]]; then
        port=5432
    fi

    echo "[INFO] 檢查 ${host}:${port} DB='${src_db}' 是否有連線佔用..."

    local conn_cnt
    conn_cnt=$(
        PGPASSWORD="$pass" PGCONNECT_TIMEOUT=3 \
        psql -h "$host" -p "$port" -U "$user" -d postgres \
             -tA -c "SELECT count(*) FROM pg_stat_activity WHERE datname = '${src_db}' AND pid <> pg_backend_pid();" \
             2>/dev/null || echo "ERROR"
    )

    conn_cnt=$(echo "$conn_cnt" | tr -d '[:space:]')

    if [[ "$conn_cnt" == "ERROR" || -z "$conn_cnt" ]]; then
        echo "[WARN] 無法取得 ${host}:${port} DB='${src_db}' 連線數，略過 restart 判斷" >&2
        return 0
    fi

    if ! [[ "$conn_cnt" =~ ^[0-9]+$ ]]; then
        echo "[WARN] 取得的連線數不是數字：'${conn_cnt}'，略過 restart 判斷" >&2
        return 0
    fi

    if (( conn_cnt == 0 )); then
        echo "[INFO] ${host}:${port} DB='${src_db}' 無活動連線，略過 PostgreSQL restart"
        return 0
    fi

    echo "[INFO] ${host}:${port} DB='${src_db}' 目前有 ${conn_cnt} 個連線，執行遠端 PostgreSQL restart..."

    local PG_VERSION
    PG_VERSION=$(ssh -o StrictHostKeyChecking=no "$host" "psql --version | awk '{print \$3}' | cut -d. -f1" 2>/dev/null || echo "")

    if [[ -z "$PG_VERSION" ]]; then
        echo "[NG]  無法在 ${host} 取得 psql 版本，無法決定 systemd 服務名稱，略過 restart" >&2
        return 1
    fi

    if ssh -o StrictHostKeyChecking=no "$host" "sudo systemctl restart postgresql@${PG_VERSION}-main.service"; then
        echo "[OK]  已在 ${host} 重啟 postgresql@${PG_VERSION}-main.service"
        ssh -o StrictHostKeyChecking=no "$host" "sudo systemctl status postgresql@${PG_VERSION}-main.service --no-pager" >/dev/null 2>&1 || true
        sleep 3
        return 0
    else
        echo "[NG]  在 ${host} 重啟 postgresql@${PG_VERSION}-main.service 失敗" >&2
        return 1
    fi
}

########################################
# 共用：在指定主機上用備份 DB 還原原 DB
########################################
# 參數：HOST SRC_DB USER PASS [PORT] POSTFIX
# 動作：
#   - 檢查備份 DB "<SRC_DB><POSTFIX>" 是否存在
#   - 若存在：
#       SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '<SRC_DB>';
#       DROP DATABASE IF EXISTS "<SRC_DB>";
#       CREATE DATABASE "<SRC_DB>" TEMPLATE "<SRC_DB><POSTFIX>";
#
# 若備份 DB 不存在，回傳失敗，不動原 DB。
pg_restore_db() {
    local host="${1-}"
    local src_db="${2-}"
    local user="${3-}"
    local pass="${4-}"
    local port="${5-}"
    local postfix="${6-}"

    if [[ -z "$host" || -z "$src_db" || -z "$user" ]]; then
        echo "[NG]  參數不完整，host/src_db/user 不可為空 (host='$host', db='$src_db')" >&2
        return 1
    fi

    if [[ -z "$port" ]]; then
        port=5432
    fi

    local backup_db="${src_db}${postfix}"

    echo ">>> ${host}:${port} - restore DB '${src_db}' from backup '${backup_db}'"

    # 檢查備份 DB 是否存在
    local exists
    exists=$(
        PGPASSWORD="$pass" PGCONNECT_TIMEOUT=3 \
        psql -h "$host" -p "$port" -U "$user" -d postgres \
             -tA -c "SELECT 1 FROM pg_database WHERE datname = '${backup_db}'" \
             2>/dev/null || true
    )

    if [[ "$exists" != "1" ]]; then
        echo "[NG]    ${host}:${port} - 備份 DB '${backup_db}' 不存在，無法還原 '${src_db}'" >&2
        return 1
    fi

    # 終止原 DB 連線 + DROP + CREATE from TEMPLATE
    if PGPASSWORD="$pass" PGCONNECT_TIMEOUT=5 \
        psql -h "$host" -p "$port" -U "$user" -d postgres -v ON_ERROR_STOP=1 <<SQL
SELECT pg_terminate_backend(pid)
  FROM pg_stat_activity
 WHERE datname = '${src_db}';

DROP DATABASE IF EXISTS "${src_db}";

CREATE DATABASE "${src_db}" TEMPLATE "${backup_db}";
SQL
    then
        echo "[OK]    ${host}:${port} - 已由 '${backup_db}' 還原 DB '${src_db}'"
        return 0
    else
        echo "[NG]    ${host}:${port} - 還原 DB '${src_db}' (from '${backup_db}') 失敗" >&2
        return 1
    fi
}

echo "=== 還原 PostgreSQL DB：$SETUP_FILE，POSTFIX='${DB_POSTFIX}'，TARGET='${TARGET_KIND:-ALL}' ==="

overall_status=0

########################################
# 1. Account DB
########################################
if [[ -z "$TARGET_KIND" || "$TARGET_KIND" == "AccountDB" ]]; then
    echo "--- AccountDB ---"
    {
        host=$(get_accountdb_value "AccountDBIP")
        dbname=$(get_accountdb_value "AccountDBName")
        user=$(get_accountdb_value "AccountDBUser")
        pass=$(get_accountdb_value "AccountDBPW")
        port=$(get_accountdb_value "AccountDBPort")

        if [[ -z "${dbname-}" ]]; then
            echo "[INFO]  未設定 AccountDBName，略過"
        else
            # 有連線才重啟 PostgreSQL
            if ! pg_check_and_restart_if_in_use "$host" "$dbname" "$user" "$pass" "$port"; then
                overall_status=1
            fi

            if pg_restore_db "$host" "$dbname" "$user" "$pass" "$port" "$DB_POSTFIX"; then
                :
            else
                overall_status=1
            fi
        fi
    }
    echo
fi

########################################
# 2. Game DB
########################################
if [[ -z "$TARGET_KIND" || "$TARGET_KIND" == "GameDB" ]]; then
    echo "--- GameDB ---"
    {
        host=$(get_gamedb_value "GameDBIP")
        dbname=$(get_gamedb_value "GameDB")
        user=$(get_gamedb_value "GameDBUser")
        pass=$(get_gamedb_value "GameDBPassword")
        port=$(get_gamedb_value "GameDBPort")

        if [[ -z "${dbname-}" ]]; then
            echo "[INFO]  未設定 GameDB，略過"
        else
            if ! pg_check_and_restart_if_in_use "$host" "$dbname" "$user" "$pass" "$port"; then
                overall_status=1
            fi

            if pg_restore_db "$host" "$dbname" "$user" "$pass" "$port" "$DB_POSTFIX"; then
                :
            else
                overall_status=1
            fi
        fi
    }
    echo
fi

########################################
# 3. 所有 World DB
########################################
if [[ -z "$TARGET_KIND" || "$TARGET_KIND" == "WorldDB" ]]; then
    echo "--- WorldDB ---"
    if [[ "${#world_order[@]}" -gt 0 ]]; then
        for wid in "${world_order[@]}"; do
            host=$(get_worlddb_value "$wid" "WorldDBIP")
            dbname=$(get_worlddb_value "$wid" "WorldDBName")
            user=$(get_worlddb_value "$wid" "WorldDBUser")
            pass=$(get_worlddb_value "$wid" "WorldDBPassword")
            port=$(get_worlddb_value "$wid" "WorldDBPort")

            if [[ -z "${dbname-}" ]]; then
                echo "[INFO]  WorldID=${wid} 未設定 WorldDBName，略過"
                continue
            fi

            if ! pg_check_and_restart_if_in_use "$host" "$dbname" "$user" "$pass" "$port"; then
                overall_status=1
            fi

            if pg_restore_db "$host" "$dbname" "$user" "$pass" "$port" "$DB_POSTFIX"; then
                :
            else
                overall_status=1
            fi
        done
    else
        echo "[INFO]  沒有任何 WorldDB 設定"
    fi
fi

echo "=== 還原作業結束 ==="

exit "$overall_status"

