#!/usr/bin/env bash
set -euo pipefail

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
ENV_FILE="${1:-$ROOT_DIR/.env.production}"
BACKUP_REF="${2:-latest}"
BACKUP_ROOT="${3:-$ROOT_DIR/backups}"
COMPOSE_FILE="$ROOT_DIR/infra/docker-compose.prod.yml"
DRILL_DB="${DRILL_DB:-1}"
DRILL_S3="${DRILL_S3:-1}"
KEEP_DRILL_DB="${KEEP_DRILL_DB:-0}"
KEEP_DRILL_BUCKET="${KEEP_DRILL_BUCKET:-0}"

if [[ "${1:-}" == "--help" || "${1:-}" == "-h" ]]; then
  cat <<'USAGE'
Usage:
  scripts/ops/restore-drill.sh [env-file] [backup-dir-or-latest] [backup-root]

Examples:
  scripts/ops/restore-drill.sh
  scripts/ops/restore-drill.sh .env.production latest /var/backups/productier
  scripts/ops/restore-drill.sh .env.production backups/20260401T120000Z

Behavior:
  - Verifies selected backup.
  - Imports DB dump into temporary drill database and runs sanity queries.
  - Optionally syncs backup objects into a temporary drill bucket.
  - Cleans up temporary DB/bucket unless KEEP_DRILL_DB=1 or KEEP_DRILL_BUCKET=1.
USAGE
  exit 0
fi

if [[ ! -f "$ENV_FILE" ]]; then
  echo "[error] env file not found: $ENV_FILE" >&2
  exit 1
fi

if ! command -v docker >/dev/null 2>&1; then
  echo "[error] docker CLI is required" >&2
  exit 1
fi

node "$ROOT_DIR/scripts/check-production-env.mjs" "$ENV_FILE"

resolve_latest_backup() {
  local root="$1"
  find "$root" -mindepth 1 -maxdepth 1 -type d -printf '%f\n' \
    | grep -E '^[0-9]{8}T[0-9]{6}Z$' \
    | sort \
    | tail -n 1
}

if [[ "$BACKUP_REF" == "latest" ]]; then
  latest_name="$(resolve_latest_backup "$BACKUP_ROOT")"
  if [[ -z "$latest_name" ]]; then
    echo "[error] no backups found in $BACKUP_ROOT" >&2
    exit 1
  fi
  BACKUP_DIR="$BACKUP_ROOT/$latest_name"
elif [[ -d "$BACKUP_REF" ]]; then
  BACKUP_DIR="$BACKUP_REF"
else
  BACKUP_DIR="$BACKUP_ROOT/$BACKUP_REF"
fi

if [[ ! -d "$BACKUP_DIR" ]]; then
  echo "[error] backup directory not found: $BACKUP_DIR" >&2
  exit 1
fi

bash "$ROOT_DIR/scripts/ops/verify-backup.sh" "$BACKUP_DIR"

timestamp_compact="$(date -u +%Y%m%d%H%M%S)"
drill_db_name="productier_drill_${timestamp_compact}"
drill_bucket_name="productier-drill-${timestamp_compact}"
db_created=0
bucket_created=0

cleanup_db() {
  if [[ "$db_created" != "1" ]]; then
    return
  fi
  if [[ "$KEEP_DRILL_DB" == "1" ]]; then
    echo "[info] KEEP_DRILL_DB=1, retaining drill database: $drill_db_name"
    return
  fi
  docker compose --env-file "$ENV_FILE" -f "$COMPOSE_FILE" exec -T postgres \
    psql -U productier -d postgres -v ON_ERROR_STOP=1 \
    -c "DROP DATABASE IF EXISTS ${drill_db_name};" >/dev/null
}

cleanup_bucket() {
  if [[ "$bucket_created" != "1" ]]; then
    return
  fi
  if [[ "$KEEP_DRILL_BUCKET" == "1" ]]; then
    echo "[info] KEEP_DRILL_BUCKET=1, retaining drill bucket: $drill_bucket_name"
    return
  fi
  docker compose --env-file "$ENV_FILE" -f "$COMPOSE_FILE" run --rm --no-deps \
    --entrypoint /bin/sh \
    rustfs-init \
    -lc "set -euo pipefail; endpoint='http://rustfs:9000'; aws --endpoint-url \"\$endpoint\" s3 rb \"s3://${drill_bucket_name}\" --force >/dev/null 2>&1 || true"
}

cleanup() {
  cleanup_db
  cleanup_bucket
}

trap cleanup EXIT

if [[ "$DRILL_DB" == "1" ]]; then
  echo "[step] running DB restore drill..."
  docker compose --env-file "$ENV_FILE" -f "$COMPOSE_FILE" exec -T postgres \
    psql -U productier -d postgres -v ON_ERROR_STOP=1 \
    -c "CREATE DATABASE ${drill_db_name};" >/dev/null
  db_created=1

  gunzip -c "$BACKUP_DIR/postgres.sql.gz" \
    | docker compose --env-file "$ENV_FILE" -f "$COMPOSE_FILE" exec -T postgres \
      psql -U productier -d "$drill_db_name" -v ON_ERROR_STOP=1 >/dev/null

  table_count="$(docker compose --env-file "$ENV_FILE" -f "$COMPOSE_FILE" exec -T postgres \
    psql -U productier -d "$drill_db_name" -At -c "SELECT count(*) FROM information_schema.tables WHERE table_schema='public';")"
  if ! [[ "$table_count" =~ ^[0-9]+$ ]] || [[ "$table_count" -lt 1 ]]; then
    echo "[error] DB restore drill failed: unexpected table count ($table_count)" >&2
    exit 1
  fi
  echo "[ok] DB restore drill passed (public tables: $table_count)"
fi

if [[ "$DRILL_S3" == "1" ]]; then
  if [[ -d "$BACKUP_DIR/s3" ]]; then
    echo "[step] running object storage restore drill..."
    docker compose --env-file "$ENV_FILE" -f "$COMPOSE_FILE" run --rm --no-deps \
      -v "$BACKUP_DIR/s3:/restore:ro" \
      --entrypoint /bin/sh \
      rustfs-init \
      -lc "set -euo pipefail; endpoint='http://rustfs:9000'; aws --endpoint-url \"\$endpoint\" s3api head-bucket --bucket '${drill_bucket_name}' >/dev/null 2>&1 || aws --endpoint-url \"\$endpoint\" s3api create-bucket --bucket '${drill_bucket_name}'; aws --endpoint-url \"\$endpoint\" s3 sync /restore \"s3://${drill_bucket_name}\" --no-progress; count=\$(aws --endpoint-url \"\$endpoint\" s3 ls \"s3://${drill_bucket_name}\" --recursive | wc -l); echo \"[ok] S3 restore drill passed (objects: \$count)\""
    bucket_created=1
  else
    echo "[warn] skipping S3 restore drill: no s3/ directory in backup"
  fi
fi

echo "[ok] restore drill completed using backup: $BACKUP_DIR"