#!/usr/bin/env bash
# PostgreSQL 논리 복원 (pg_restore). custom format(.dump) 전용.
# 사용 예:
#   bash scripts/pg-restore.sh /home/xavis/workspace/backup/ai_platform/latest/ai_web_platform.dump
#   bash scripts/pg-restore.sh --all --globals --confirm /home/xavis/workspace/backup/ai_platform/latest/
#   bash scripts/pg-restore.sh --test --confirm /home/xavis/workspace/backup/ai_platform/latest/ai_web_platform.dump
set -euo pipefail

usage() {
  cat <<'EOF'
Usage: bash scripts/pg-restore.sh [options] <path-to.dump | backup-dir>

Options:
  --all        백업 디렉터리 안의 *.dump 전체 복원 (00_globals.sql 있으면 --globals 권장)
  --test       복원 대상 DB를 {dbname}_restore_test 로 생성 (단일 dump 또는 --all)
  --clean      pg_restore --clean --if-exists (기존 객체 삭제 후 복원, --test 미사용 시 위험)
  --confirm    확인 없이 실행 (cron·자동화용; --test 없이 --clean 시 필수)
  --globals    같은 디렉터리의 00_globals.sql 을 먼저 적용 (슈퍼유저 env 필요)
  -h, --help   도움말

환경 변수 (.env):
  DB_*                     연결 정보
  PG_BACKUP_SUPERUSER      --globals 시 사용 (기본 postgres)
  PG_BACKUP_SUPERUSER_PASSWORD
EOF
}

REPO_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
# shellcheck source=scripts/lib/load-env.sh
source "$REPO_ROOT/scripts/lib/load-env.sh"

MODE="prod"
CLEAN=0
CONFIRM=0
GLOBALS=0
RESTORE_ALL=0
TARGET_PATH=""

while [[ $# -gt 0 ]]; do
  case "$1" in
    --test) MODE="test"; shift ;;
    --clean) CLEAN=1; shift ;;
    --confirm) CONFIRM=1; shift ;;
    --globals) GLOBALS=1; shift ;;
    --all) RESTORE_ALL=1; shift ;;
    -h|--help) usage; exit 0 ;;
    -*) echo "Unknown option: $1" >&2; usage >&2; exit 1 ;;
    *)
      if [[ -n "$TARGET_PATH" ]]; then
        echo "Unexpected argument: $1" >&2
        exit 1
      fi
      TARGET_PATH="$1"
      shift
      ;;
  esac
done

if [[ -z "$TARGET_PATH" ]]; then
  usage >&2
  exit 1
fi

load_project_env "$REPO_ROOT/.env"

: "${DB_HOST:?DB_HOST is required in .env}"
: "${DB_DATABASE:?DB_DATABASE is required in .env}"
: "${DB_USERNAME:?DB_USERNAME is required in .env}"
: "${DB_PASSWORD:?DB_PASSWORD is required in .env}"

DB_PORT="${DB_PORT:-5432}"

export PGHOST="$DB_HOST"
export PGPORT="$DB_PORT"
export PGUSER="$DB_USERNAME"
export PGPASSWORD="$DB_PASSWORD"

if ! command -v pg_restore >/dev/null 2>&1; then
  echo "pg-restore: pg_restore not found. Install postgresql-client." >&2
  exit 1
fi

log_ts() { date '+%Y-%m-%dT%H:%M:%S%z'; }

resolve_superuser_creds() {
  SUPERUSER="${PG_BACKUP_SUPERUSER:-postgres}"
  if [[ -n "${PG_BACKUP_SUPERUSER_PASSWORD:-}" ]]; then
    SU_USER="$SUPERUSER"
    SU_PASS="$PG_BACKUP_SUPERUSER_PASSWORD"
  else
    SU_USER="$DB_USERNAME"
    SU_PASS="$DB_PASSWORD"
  fi
}

apply_globals() {
  local globals_file="$1"
  if [[ ! -f "$globals_file" ]]; then
    echo "pg-restore: globals file not found: $globals_file" >&2
    exit 1
  fi
  : "${PG_BACKUP_SUPERUSER_PASSWORD:?PG_BACKUP_SUPERUSER_PASSWORD required for --globals}"
  resolve_superuser_creds
  export PGUSER="$SU_USER"
  export PGPASSWORD="$SU_PASS"
  echo "[$(log_ts)] applying globals: $globals_file"
  psql -h "$PGHOST" -p "$PGPORT" -U "$PGUSER" -v ON_ERROR_STOP=1 -f "$globals_file"
  export PGUSER="$DB_USERNAME"
  export PGPASSWORD="$DB_PASSWORD"
}

ensure_database() {
  local db_name="$1"
  resolve_superuser_creds
  export PGUSER="$SU_USER"
  export PGPASSWORD="$SU_PASS"
  psql -h "$PGHOST" -p "$PGPORT" -U "$PGUSER" -v ON_ERROR_STOP=1 -tc \
    "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '$db_name' AND pid <> pg_backend_pid();" \
    >/dev/null 2>&1 || true
  if [[ "$MODE" == "test" ]]; then
    psql -h "$PGHOST" -p "$PGPORT" -U "$PGUSER" -v ON_ERROR_STOP=1 -c "DROP DATABASE IF EXISTS \"$db_name\";"
    psql -h "$PGHOST" -p "$PGPORT" -U "$PGUSER" -v ON_ERROR_STOP=1 -c \
      "CREATE DATABASE \"$db_name\" OWNER \"$DB_USERNAME\";"
  fi
  export PGUSER="$DB_USERNAME"
  export PGPASSWORD="$DB_PASSWORD"
}

restore_one_dump() {
  local dump_path="$1"
  local source_db
  source_db="$(basename "$dump_path" .dump)"
  local target_db="$source_db"
  if [[ "$MODE" == "test" ]]; then
    target_db="${source_db}_restore_test"
    ensure_database "$target_db"
  fi

  local restore_args=(
    -h "$PGHOST"
    -p "$PGPORT"
    -U "$PGUSER"
    -d "$target_db"
    --no-owner
    --role="$DB_USERNAME"
    -v
  )
  if [[ "$CLEAN" -eq 1 ]]; then
    restore_args+=(--clean --if-exists)
  fi

  echo "[$(log_ts)] pg_restore → $target_db ($dump_path)"
  pg_restore "${restore_args[@]}" "$dump_path"
  echo "[$(log_ts)] pg-restore done: $target_db"
}

confirm_restore() {
  local message="$1"
  if [[ "$MODE" == "prod" && "$CLEAN" -eq 1 && "$CONFIRM" -ne 1 ]]; then
    echo "pg-restore: --clean on production DB requires --confirm" >&2
    exit 1
  fi
  if [[ "$CONFIRM" -eq 1 ]]; then
    return 0
  fi
  echo "$message"
  read -r -p "Continue? [y/N] " ans
  [[ "$ans" == "y" || "$ans" == "Y" ]] || exit 0
}

if [[ -d "$TARGET_PATH" ]]; then
  BACKUP_DIR="${TARGET_PATH%/}"
  if [[ "$RESTORE_ALL" -ne 1 ]]; then
    echo "pg-restore: path is a directory. Use --all to restore every *.dump inside." >&2
    exit 1
  fi
  confirm_restore "About to restore ALL dumps in: $BACKUP_DIR on $DB_HOST:$PGPORT"
  if [[ "$GLOBALS" -eq 1 ]]; then
    apply_globals "$BACKUP_DIR/00_globals.sql"
  fi
  shopt -s nullglob
  dumps=( "$BACKUP_DIR"/*.dump )
  shopt -u nullglob
  if [[ ${#dumps[@]} -eq 0 ]]; then
    echo "pg-restore: no .dump files in $BACKUP_DIR" >&2
    exit 1
  fi
  for dump in "${dumps[@]}"; do
    restore_one_dump "$dump"
  done
  exit 0
fi

if [[ ! -f "$TARGET_PATH" ]]; then
  echo "pg-restore: path not found: $TARGET_PATH" >&2
  exit 1
fi

if [[ "$RESTORE_ALL" -eq 1 ]]; then
  echo "pg-restore: --all requires a backup directory path." >&2
  exit 1
fi

TARGET_DB="$DB_DATABASE"
if [[ "$MODE" == "test" ]]; then
  TARGET_DB="${DB_DATABASE}_restore_test"
fi

confirm_restore "About to restore into database: $TARGET_DB on $DB_HOST:$PGPORT
Dump: $TARGET_PATH"

if [[ "$GLOBALS" -eq 1 ]]; then
  apply_globals "$(dirname "$TARGET_PATH")/00_globals.sql"
fi

if [[ "$MODE" == "test" ]]; then
  ensure_database "$TARGET_DB"
fi

restore_one_dump "$TARGET_PATH"

if [[ "$MODE" == "test" ]]; then
  echo "Verify with: psql -h $DB_HOST -U $DB_USERNAME -d $TARGET_DB -c '\\dt'"
  echo "Drop test DB: psql ... -c 'DROP DATABASE \"$TARGET_DB\";'"
fi