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

usage() {
  cat >&2 <<'EOF'
Usage: rehearse-restore.sh [--keep-on-failure] <frameflow-backup-directory>

Restores a verified FrameFlow backup into isolated, randomly named Docker
resources. No host ports are published. Resources are removed automatically
unless --keep-on-failure is supplied and the rehearsal fails.
EOF
  exit "${1:-2}"
}

fail() {
  echo "Restore rehearsal failed: $1" >&2
  exit 1
}

keep_on_failure=0
while [[ $# -gt 0 ]]; do
  case "$1" in
    --keep-on-failure)
      keep_on_failure=1
      shift
      ;;
    --help|-h)
      usage 0
      ;;
    --)
      shift
      break
      ;;
    -*)
      usage
      ;;
    *)
      break
      ;;
  esac
done

[[ $# -eq 1 ]] || usage

script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
backup_input="$1"
timeout_seconds="${FRAMEFLOW_RESTORE_TIMEOUT_SECONDS:-120}"
metrics_dir="${FRAMEFLOW_RESTORE_METRICS_DIR:-${FRAMEFLOW_BACKUP_METRICS_DIR:-/var/lib/frameflow/metrics}}"
started_at_seconds="$(date +%s)"
metrics_tmp=""

[[ "${timeout_seconds}" =~ ^[0-9]+$ ]] || fail "FRAMEFLOW_RESTORE_TIMEOUT_SECONDS must be an integer"
(( timeout_seconds >= 5 && timeout_seconds <= 3600 )) \
  || fail "FRAMEFLOW_RESTORE_TIMEOUT_SECONDS must be between 5 and 3600"

# Verification deliberately precedes every Docker call.
"${script_dir}/verify-backup.sh" "${backup_input}"
backup_dir="$(cd "${backup_input}" && pwd -P)"

command -v docker >/dev/null 2>&1 || fail "docker is required"
docker version >/dev/null 2>&1 || fail "the Docker engine is unavailable"

random_hex() {
  local bytes="$1"
  local value
  value="$(LC_ALL=C od -An -N "${bytes}" -tx1 /dev/urandom | tr -d '[:space:]')"
  [[ "${value}" =~ ^[0-9a-f]+$ ]] || fail "could not generate a random resource identifier"
  printf '%s' "${value}"
}

timestamp="$(date -u +%Y%m%d%H%M%S)"
suffix="$(random_hex 4)"
run_id="${timestamp}-${suffix}"
prefix="frameflow-restore-${run_id}"
label_flag="com.frameflow.restore-rehearsal"
label_id="com.frameflow.restore-rehearsal.id"
expected_labels="true|${run_id}"
labels=(--label "${label_flag}=true" --label "${label_id}=${run_id}")

network_name="${prefix}-network"
postgres_volume="${prefix}-postgres"
redis_volume="${prefix}-redis"
minio_volume="${prefix}-minio"
postgres_container="${prefix}-postgres"
redis_container="${prefix}-redis"
minio_container="${prefix}-minio"
redis_extract_container="${prefix}-redis-extract"
minio_extract_container="${prefix}-minio-extract"
minio_check_container="${prefix}-minio-check"

postgres_image="postgres:16-alpine"
redis_image="redis:7.4-alpine"
alpine_image="alpine:3.20"
minio_image="minio/minio@sha256:14cea493d9a34af32f524e538b8346cf79f3321eff8e708c1e2960462bd8936e"
mc_image="minio/mc@sha256:a7fe349ef4bd8521fb8497f55c6042871b2ae640607cf99d9bede5e9bdf11727"

postgres_db="frameflow"
postgres_user="frameflow_restore"
postgres_password="$(random_hex 24)"
redis_password="$(random_hex 24)"
minio_access_key="restore${suffix}"
minio_secret_key="$(random_hex 32)"

resources_cleaned=0

resource_labels() {
  local kind="$1"
  local name="$2"
  case "${kind}" in
    container)
      docker container inspect --format "{{ index .Config.Labels \"${label_flag}\" }}|{{ index .Config.Labels \"${label_id}\" }}" "${name}" 2>/dev/null
      ;;
    network|volume)
      docker "${kind}" inspect --format "{{ index .Labels \"${label_flag}\" }}|{{ index .Labels \"${label_id}\" }}" "${name}" 2>/dev/null
      ;;
    *)
      return 2
      ;;
  esac
}

remove_labeled_resource() {
  local kind="$1"
  local name="$2"
  local actual

  if ! actual="$(resource_labels "${kind}" "${name}")"; then
    return 0
  fi
  if [[ "${actual}" != "${expected_labels}" ]]; then
    echo "Refusing to remove ${kind} ${name}: rehearsal label mismatch" >&2
    return 1
  fi

  case "${kind}" in
    container) docker rm -f "${name}" >/dev/null ;;
    network) docker network rm "${name}" >/dev/null ;;
    volume) docker volume rm "${name}" >/dev/null ;;
  esac
}

cleanup_resources() {
  local result=0
  local name

  for name in \
    "${minio_check_container}" \
    "${minio_extract_container}" \
    "${redis_extract_container}" \
    "${minio_container}" \
    "${redis_container}" \
    "${postgres_container}"; do
    if ! remove_labeled_resource container "${name}"; then
      result=1
    fi
  done

  if ! remove_labeled_resource network "${network_name}"; then
    result=1
  fi
  for name in "${minio_volume}" "${redis_volume}" "${postgres_volume}"; do
    if ! remove_labeled_resource volume "${name}"; then
      result=1
    fi
  done
  return "${result}"
}

cleanup_on_exit() {
  local status=$?
  trap - EXIT

  if [[ -n "${metrics_tmp}" && -f "${metrics_tmp}" ]]; then
    rm -f "${metrics_tmp}"
  fi

  if (( status != 0 && keep_on_failure == 1 )); then
    cat >&2 <<EOF
Restore rehearsal resources retained for diagnosis:
  network: ${network_name}
  volumes: ${postgres_volume}, ${redis_volume}, ${minio_volume}
  containers: ${postgres_container}, ${redis_container}, ${minio_container}
Remove them only after confirming both labels ${label_flag}=true and ${label_id}=${run_id}.
EOF
    exit "${status}"
  fi

  if (( resources_cleaned == 0 )); then
    if ! cleanup_resources && (( status == 0 )); then
      status=1
    fi
  fi
  exit "${status}"
}

trap cleanup_on_exit EXIT
trap 'exit 130' INT
trap 'exit 143' TERM

wait_for_postgres() {
  local deadline=$((SECONDS + timeout_seconds))
  until docker exec "${postgres_container}" pg_isready -U "${postgres_user}" -d "${postgres_db}" >/dev/null 2>&1; do
    (( SECONDS < deadline )) || fail "PostgreSQL did not become ready within ${timeout_seconds} seconds"
    sleep 1
  done
}

wait_for_redis() {
  local deadline=$((SECONDS + timeout_seconds))
  until docker exec "${redis_container}" redis-cli --no-auth-warning -a "${redis_password}" ping 2>/dev/null | grep -q '^PONG$'; do
    (( SECONDS < deadline )) || fail "Redis did not become ready within ${timeout_seconds} seconds"
    sleep 1
  done
}

wait_for_minio() {
  local deadline=$((SECONDS + timeout_seconds))
  until docker exec "${minio_container}" curl -fsS http://127.0.0.1:9000/minio/health/live >/dev/null 2>&1; do
    (( SECONDS < deadline )) || fail "MinIO did not become ready within ${timeout_seconds} seconds"
    sleep 1
  done
}

echo "Creating isolated restore resources for rehearsal ${run_id}..."
docker network create --internal "${labels[@]}" "${network_name}" >/dev/null
docker volume create "${labels[@]}" "${postgres_volume}" >/dev/null
docker volume create "${labels[@]}" "${redis_volume}" >/dev/null
docker volume create "${labels[@]}" "${minio_volume}" >/dev/null

echo "Restoring PostgreSQL into an empty rehearsal volume..."
docker run -d \
  --name "${postgres_container}" \
  --network "${network_name}" \
  "${labels[@]}" \
  -e POSTGRES_DB="${postgres_db}" \
  -e POSTGRES_USER="${postgres_user}" \
  -e POSTGRES_PASSWORD="${postgres_password}" \
  -v "${postgres_volume}:/var/lib/postgresql/data" \
  "${postgres_image}" >/dev/null
wait_for_postgres
docker exec -i "${postgres_container}" \
  pg_restore -U "${postgres_user}" -d "${postgres_db}" \
  --exit-on-error --no-owner --no-privileges \
  < "${backup_dir}/postgres.dump"

core_tables_ok="$(docker exec "${postgres_container}" psql -v ON_ERROR_STOP=1 -U "${postgres_user}" -d "${postgres_db}" -Atc \
  "select case when to_regclass('drizzle.__drizzle_migrations') is not null
    and to_regclass('public.users') is not null
    and to_regclass('public.workspaces') is not null
    and to_regclass('public.projects') is not null
    and to_regclass('public.episodes') is not null
    and to_regclass('public.assets') is not null
    and to_regclass('public.asset_versions') is not null
    and to_regclass('public.shots') is not null
    and to_regclass('public.shot_versions') is not null
    and to_regclass('public.generation_jobs') is not null
    and to_regclass('public.renders') is not null
    and to_regclass('public.provider_verifications') is not null
    then 'ok' else 'missing' end")"
[[ "${core_tables_ok}" == "ok" ]] || fail "PostgreSQL restore is missing required tables"

docker exec "${postgres_container}" psql -v ON_ERROR_STOP=1 -U "${postgres_user}" -d "${postgres_db}" -Atc \
  "select count(*) from users;
   select count(*) from workspaces;
   select count(*) from projects;
   select count(*) from episodes;
   select count(*) from assets;
   select count(*) from shots;
   select count(*) from generation_jobs;
   select count(*) from renders;
   select count(*) from provider_verifications;" >/dev/null

migration_count="$(docker exec "${postgres_container}" psql -v ON_ERROR_STOP=1 -U "${postgres_user}" -d "${postgres_db}" -Atc \
  'select count(*) from drizzle.__drizzle_migrations')"
table_count="$(docker exec "${postgres_container}" psql -v ON_ERROR_STOP=1 -U "${postgres_user}" -d "${postgres_db}" -Atc \
  "select count(*) from pg_tables where schemaname in ('public', 'drizzle')")"
[[ "${migration_count}" =~ ^[0-9]+$ && "${table_count}" =~ ^[0-9]+$ ]] \
  || fail "PostgreSQL verification returned invalid counts"
(( migration_count > 0 )) || fail "PostgreSQL migration history is empty"

echo "Restoring Redis persistence into an empty rehearsal volume..."
docker run --rm \
  --name "${redis_extract_container}" \
  "${labels[@]}" \
  -v "${redis_volume}:/target" \
  --mount "type=bind,src=${backup_dir},dst=/backup,readonly" \
  --entrypoint /bin/sh \
  "${alpine_image}" -ec \
  'test -z "$(find /target -mindepth 1 -maxdepth 1 -print -quit)" && tar -xzf /backup/redis-data.tar.gz -C /target'
docker run -d \
  --name "${redis_container}" \
  --network "${network_name}" \
  "${labels[@]}" \
  -v "${redis_volume}:/data" \
  "${redis_image}" redis-server --appendonly yes --requirepass "${redis_password}" >/dev/null
wait_for_redis
redis_key_count="$(docker exec "${redis_container}" redis-cli --no-auth-warning -a "${redis_password}" --raw DBSIZE)"
redis_persistence="$(docker exec "${redis_container}" redis-cli --no-auth-warning -a "${redis_password}" --raw INFO persistence | tr -d '\r')"
[[ "${redis_key_count}" =~ ^[0-9]+$ ]] || fail "Redis verification returned an invalid key count"
grep -q '^loading:0$' <<< "${redis_persistence}" || fail "Redis is still loading restored persistence"
grep -q '^aof_last_write_status:ok$' <<< "${redis_persistence}" || fail "Redis AOF persistence is not healthy"

echo "Restoring MinIO data into an empty rehearsal volume..."
docker run --rm \
  --name "${minio_extract_container}" \
  "${labels[@]}" \
  -v "${minio_volume}:/target" \
  --mount "type=bind,src=${backup_dir},dst=/backup,readonly" \
  --entrypoint /bin/sh \
  "${alpine_image}" -ec \
  'test -z "$(find /target -mindepth 1 -maxdepth 1 -print -quit)" && tar -xzf /backup/minio-data.tar.gz -C /target'
docker run -d \
  --name "${minio_container}" \
  --network "${network_name}" \
  "${labels[@]}" \
  -e MINIO_ROOT_USER="${minio_access_key}" \
  -e MINIO_ROOT_PASSWORD="${minio_secret_key}" \
  -v "${minio_volume}:/data" \
  "${minio_image}" server /data --console-address :9001 >/dev/null
wait_for_minio
minio_object_count="$(docker run --rm \
  --name "${minio_check_container}" \
  --network "${network_name}" \
  "${labels[@]}" \
  -e MC_HOST_RESTORE="http://${minio_access_key}:${minio_secret_key}@${minio_container}:9000" \
  --entrypoint /bin/sh \
  "${mc_image}" -ec \
  'mc admin info RESTORE >/dev/null && mc ls RESTORE >/dev/null && mc ls --recursive RESTORE | wc -l | tr -d "[:space:]"')"
[[ "${minio_object_count}" =~ ^[0-9]+$ ]] || fail "MinIO verification returned an invalid object count"

echo "Cleaning isolated restore resources..."
cleanup_resources || fail "one or more rehearsal resources failed label verification or cleanup"
resources_cleaned=1

completed_at_seconds="$(date +%s)"
mkdir -p "${metrics_dir}"
metrics_tmp="$(mktemp "${metrics_dir%/}/frameflow_restore.prom.tmp.XXXXXX")"
cat > "${metrics_tmp}" <<EOF
# HELP frameflow_restore_rehearsal_last_success_timestamp_seconds Unix timestamp of the last successful isolated restore rehearsal.
# TYPE frameflow_restore_rehearsal_last_success_timestamp_seconds gauge
frameflow_restore_rehearsal_last_success_timestamp_seconds ${completed_at_seconds}
# HELP frameflow_restore_rehearsal_last_duration_seconds Duration of the last successful isolated restore rehearsal.
# TYPE frameflow_restore_rehearsal_last_duration_seconds gauge
frameflow_restore_rehearsal_last_duration_seconds $((completed_at_seconds - started_at_seconds))
# HELP frameflow_restore_rehearsal_postgres_table_count Number of restored PostgreSQL application and migration tables.
# TYPE frameflow_restore_rehearsal_postgres_table_count gauge
frameflow_restore_rehearsal_postgres_table_count ${table_count}
# HELP frameflow_restore_rehearsal_redis_key_count Number of Redis keys observed after restore.
# TYPE frameflow_restore_rehearsal_redis_key_count gauge
frameflow_restore_rehearsal_redis_key_count ${redis_key_count}
# HELP frameflow_restore_rehearsal_minio_object_count Number of MinIO objects listed after restore.
# TYPE frameflow_restore_rehearsal_minio_object_count gauge
frameflow_restore_rehearsal_minio_object_count ${minio_object_count}
EOF
chmod 0644 "${metrics_tmp}"
mv -f "${metrics_tmp}" "${metrics_dir%/}/frameflow_restore.prom"
metrics_tmp=""

echo "Restore rehearsal passed: PostgreSQL ${table_count} tables, Redis ${redis_key_count} keys, MinIO ${minio_object_count} objects."
echo "Restore rehearsal metrics updated: ${metrics_dir%/}/frameflow_restore.prom"