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

if [[ "$(id -u)" != 0 ]]; then
  echo "Run this script as root through sudo on the PostgreSQL host." >&2
  exit 1
fi

operator=${1:-${SUDO_USER:-}}
if [[ -z "$operator" || "$operator" == root ]]; then
  echo "Usage: sudo $0 OPERATOR_USER [OUTPUT_DIRECTORY]" >&2
  exit 1
fi

operator_entry=$(getent passwd "$operator") || {
  echo "Operator user does not exist: $operator" >&2
  exit 1
}
operator_home=$(cut -d: -f6 <<<"$operator_entry")
operator_group=$(id -gn "$operator")
output_dir=${2:-"$operator_home/.config/who_need_help"}

postgres_version=${POSTGRES_CLUSTER_VERSION:-18}
postgres_cluster=${POSTGRES_CLUSTER_NAME:-main}
socket_dir=${POSTGRES_SOCKET_DIR:-/var/run/postgresql}

production_role=wnh_production
production_database=who_need_help_production
staging_role=wnh_staging
staging_database=who_need_help_staging
hba_marker="# BEGIN Who Need Help managed local socket authentication"

for command in awk cat chgrp chmod chown cut date getent grep id install mktemp \
  openssl pg_ctlcluster pg_lsclusters psql rm runuser sed stat tr; do
  command -v "$command" >/dev/null 2>&1 || {
    echo "Required command is unavailable: $command" >&2
    exit 1
  }
done

cluster_status=$(
  pg_lsclusters --no-header |
    awk -v version="$postgres_version" -v cluster="$postgres_cluster" \
      '$1 == version && $2 == cluster {print $4}'
)
if [[ "$cluster_status" != online ]]; then
  echo "PostgreSQL cluster $postgres_version/$postgres_cluster is not online." >&2
  exit 1
fi

postgres_psql=(
  runuser -u postgres --
  psql --no-psqlrc --set ON_ERROR_STOP=1 --quiet --tuples-only --no-align
)

hba_file=$("${postgres_psql[@]}" --dbname postgres --command 'SHOW hba_file')
configured_socket_dirs=$(
  "${postgres_psql[@]}" --dbname postgres --command 'SHOW unix_socket_directories'
)
postgres_port=$("${postgres_psql[@]}" --dbname postgres --command 'SHOW port')

[[ "$hba_file" == /* && -f "$hba_file" ]] || {
  echo "PostgreSQL reported an unusable hba_file path." >&2
  exit 1
}
[[ "$socket_dir" == /* && -d "$socket_dir" ]] || {
  echo "PostgreSQL socket directory is unavailable: $socket_dir" >&2
  exit 1
}
if ! tr ',' '\n' <<<"$configured_socket_dirs" |
  sed -e "s/^[[:space:]']*//" -e "s/[[:space:]']*$//" |
  grep -Fx "$socket_dir" >/dev/null; then
  echo "POSTGRES_SOCKET_DIR is not listed in unix_socket_directories." >&2
  exit 1
fi

socket_path="$socket_dir/.s.PGSQL.$postgres_port"
[[ -S "$socket_path" ]] || {
  echo "PostgreSQL Unix socket is unavailable: $socket_path" >&2
  exit 1
}
socket_mode=$(stat -c '%a' "$socket_path")
case "${socket_mode: -1}" in
  6 | 7) ;;
  *)
    echo "PostgreSQL socket is not writable by the unprivileged container user." >&2
    exit 1
    ;;
esac

if grep -F "$hba_marker" "$hba_file" >/dev/null; then
  echo "Who Need Help HBA rules already exist; refusing an ambiguous reprovision." >&2
  exit 1
fi

existing_objects=$(
  "${postgres_psql[@]}" --dbname postgres --command "
    SELECT 'role:' || rolname FROM pg_roles
      WHERE rolname IN ('$production_role', '$staging_role')
    UNION ALL
    SELECT 'database:' || datname FROM pg_database
      WHERE datname IN ('$production_database', '$staging_database')
    ORDER BY 1;
  "
)
if [[ -n "$existing_objects" ]]; then
  echo "Project-scoped PostgreSQL roles or databases already exist:" >&2
  while IFS= read -r object; do
    printf '  %s\n' "$object" >&2
  done <<<"$existing_objects"
  echo "Inspect them before deciding whether to reuse, rotate, or remove them." >&2
  exit 1
fi

production_fragment="$output_dir/database-production.env"
staging_fragment="$output_dir/database-staging.env"
if [[ -e "$production_fragment" || -e "$staging_fragment" ]]; then
  echo "Database credential fragments already exist; refusing to overwrite them." >&2
  exit 1
fi

umask 077
production_password=$(openssl rand -hex 32)
staging_password=$(openssl rand -hex 32)
work_dir=$(mktemp -d)
chgrp postgres "$work_dir"
chmod 750 "$work_dir"
hba_candidate="$work_dir/pg_hba.conf"
hba_backup="$work_dir/pg_hba.conf.original"
sql_file="$work_dir/provision.sql"
backup_dir=/var/backups/who_need_help
backup_file="$backup_dir/pg_hba.conf.before-who-need-help-$(date -u +%Y%m%dT%H%M%SZ)"
provision_started=false
provision_finished=false

cleanup() {
  status=$?

  if [[ "$status" != 0 && "$provision_started" == true && "$provision_finished" == false ]]; then
    "${postgres_psql[@]}" --dbname postgres --command \
      "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname IN ('$production_database', '$staging_database') AND pid <> pg_backend_pid();" \
      >/dev/null 2>&1 || true
    "${postgres_psql[@]}" --dbname postgres --command \
      "DROP DATABASE IF EXISTS $staging_database;" >/dev/null 2>&1 || true
    "${postgres_psql[@]}" --dbname postgres --command \
      "DROP DATABASE IF EXISTS $production_database;" >/dev/null 2>&1 || true
    "${postgres_psql[@]}" --dbname postgres --command \
      "DROP ROLE IF EXISTS $staging_role; DROP ROLE IF EXISTS $production_role;" \
      >/dev/null 2>&1 || true

    if [[ -f "$hba_backup" ]]; then
      install -m 640 -o postgres -g postgres "$hba_backup" "$hba_file" || true
      pg_ctlcluster "$postgres_version" "$postgres_cluster" reload || true
    fi

    rm -f "$production_fragment" "$staging_fragment"
  fi

  rm -rf "$work_dir"
  unset production_password staging_password
  exit "$status"
}
trap cleanup EXIT HUP INT TERM

install -m 600 -o root -g root "$hba_file" "$hba_backup"
{
  printf '%s\n' "$hba_marker"
  printf 'local %s %s scram-sha-256\n' "$production_database" "$production_role"
  printf 'local %s %s scram-sha-256\n' "$staging_database" "$staging_role"
  printf '%s\n' '# END Who Need Help managed local socket authentication'
  cat "$hba_backup"
} >"$hba_candidate"

install -d -m 700 -o root -g root "$backup_dir"
install -m 600 -o root -g root "$hba_backup" "$backup_file"
install -m 640 -o postgres -g postgres "$hba_candidate" "$hba_file"
provision_started=true

hba_errors=$(
  "${postgres_psql[@]}" --dbname postgres --command \
    "SELECT count(*) FROM pg_hba_file_rules WHERE error IS NOT NULL;"
)
if [[ "$hba_errors" != 0 ]]; then
  echo "PostgreSQL rejected the candidate pg_hba.conf; restoring the backup." >&2
  exit 1
fi
pg_ctlcluster "$postgres_version" "$postgres_cluster" reload

cat >"$sql_file" <<SQL
SET password_encryption = 'scram-sha-256';
CREATE ROLE $production_role LOGIN PASSWORD '$production_password'
  NOSUPERUSER NOCREATEDB NOCREATEROLE NOREPLICATION NOBYPASSRLS;
CREATE ROLE $staging_role LOGIN PASSWORD '$staging_password'
  NOSUPERUSER NOCREATEDB NOCREATEROLE NOREPLICATION NOBYPASSRLS;
CREATE DATABASE $production_database OWNER $production_role TEMPLATE template0;
CREATE DATABASE $staging_database OWNER $staging_role TEMPLATE template0;
REVOKE CONNECT ON DATABASE $production_database FROM PUBLIC;
REVOKE CONNECT ON DATABASE $staging_database FROM PUBLIC;
GRANT CONNECT ON DATABASE $production_database TO $production_role;
GRANT CONNECT ON DATABASE $staging_database TO $staging_role;
\connect $production_database
CREATE EXTENSION IF NOT EXISTS citext;
CREATE EXTENSION IF NOT EXISTS postgis;
\connect $staging_database
CREATE EXTENSION IF NOT EXISTS citext;
CREATE EXTENSION IF NOT EXISTS postgis;
SQL
chmod 600 "$sql_file"
chown postgres:postgres "$sql_file"
"${postgres_psql[@]}" --dbname postgres --file "$sql_file"

PGPASSWORD="$production_password" psql --no-psqlrc --set ON_ERROR_STOP=1 \
  --host "$socket_dir" --username "$production_role" \
  --dbname "$production_database" --quiet --tuples-only --no-align \
  --command 'SELECT current_user, current_database(), PostGIS_Version()' >/dev/null
PGPASSWORD="$staging_password" psql --no-psqlrc --set ON_ERROR_STOP=1 \
  --host "$socket_dir" --username "$staging_role" \
  --dbname "$staging_database" --quiet --tuples-only --no-align \
  --command 'SELECT current_user, current_database(), PostGIS_Version()' >/dev/null

install -d -m 700 -o "$operator" -g "$operator_group" "$output_dir"
production_tmp="$work_dir/database-production.env"
staging_tmp="$work_dir/database-staging.env"
printf '%s\n' \
  'PRODUCTION_DATABASE_MODE=external' \
  "PRODUCTION_DATABASE_URL=ecto://$production_role:$production_password@localhost/$production_database" \
  "PRODUCTION_DATABASE_SOCKET_DIR=$socket_dir" >"$production_tmp"
printf '%s\n' \
  'PRODUCTION_DATABASE_MODE=external' \
  "PRODUCTION_DATABASE_URL=ecto://$staging_role:$staging_password@localhost/$staging_database" \
  "PRODUCTION_DATABASE_SOCKET_DIR=$socket_dir" >"$staging_tmp"
install -m 600 -o "$operator" -g "$operator_group" "$production_tmp" "$production_fragment"
install -m 600 -o "$operator" -g "$operator_group" "$staging_tmp" "$staging_fragment"

provision_finished=true

echo "Provisioned isolated production and staging PostgreSQL roles and databases."
echo "Verified SCRAM authentication, citext, and PostGIS through $socket_dir."
echo "Credential fragments (mode 0600):"
echo "  $production_fragment"
echo "  $staging_fragment"
echo "Original HBA backup (mode 0600): $backup_file"