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

# set -x

# Defaults
PGDATA_DEFAULT="/var/lib/postgresql/data"
PGPORT_DEFAULT=5342
LOCALE_DEFAULT="en_US.UTF-8"
ENCODING_DEFAULT="UTF8"
AUTH_DEFAULT="scram-sha-256"   # alternatives: md5 | password | trust
TEMP_PORT_DEFAULT="55432"
SOCKET_DIR_DEFAULT="/tmp"
ROLEFILE_PATH_DEFAULT="."

# Args
PGDATA="$PGDATA_DEFAULT"
PGPORT="$PGPORT_DEFAULT"
DB_LOCALE="$LOCALE_DEFAULT"
DB_ENCODING="$ENCODING_DEFAULT"
AUTH_METHOD="$AUTH_DEFAULT"
TEMP_PORT="$TEMP_PORT_DEFAULT"
SOCKET_DIR="$SOCKET_DIR_DEFAULT"
ROLEFILE_PATH="$ROLEFILE_PATH_DEFAULT"

POSTGRES_PASSWORD=""
ADMIN_USER="fwd_admin"
ADMIN_PASS=""
STANDARD_USER="fwd_user"
STANDARD_PASS=""
force=false
start_cluster=false

function show_usage()
{
   cat <<EOF
Usage: $0 [options]

Creates a PostgreSQL cluster with upstream initdb, compatible with the official
postgres Docker image. Safe to run on Ubuntu (host) or in an Ubuntu-based container.

Options:
  --pgdata DIR                 Data directory (default: $PGDATA_DEFAULT)
  --pgport PORT                Conection port (default: $PGPORT_DEFAULT)
  --locale LOCALE              Cluster locale (default: $LOCALE_DEFAULT)
  --encoding ENC               Cluster encoding (default: $ENCODING_DEFAULT)
  --auth METHOD                pg_hba auth method: scram-sha-256 | md5 | password | trust
                               (default: $AUTH_DEFAULT)

  --postgres-password PASS     Initial password for postgres superuser
  --admin-user NAME            Admin role name (default: $ADMIN_USER)
  --admin-pass PASS            Admin role password
  --user-user NAME             Standard role name (default: $STANDARD_USER)
  --user-pass PASS             Standard role password

  --temp-port PORT             Temporary port for bootstrap (default: $TEMP_PORT_DEFAULT)
  --socket-dir DIR             Unix socket dir for bootstrap (default: $SOCKET_DIR_DEFAULT)
  --rolefile-path FILE         Path to create_roles.sql file (default: $ROLEFILE_PATH_DEFAULT)

  --force                      Remove the cluster if it already exists.
  --start                      Start the cluster after creating it.

Examples:
  $0 --pgdata /data/pg17 --locale en_US.UTF-8 --encoding UTF8 \\
     --pgport 5433
     --auth scram-sha-256 \\
     --postgres-password adminpw \\
     --admin-user app_admin --admin-pass S3cret! \\
     --user-user app_user --user-pass userpw

Notes:
- Uses initdb (NOT pg_createcluster), so the layout matches postgres Docker images.
- Ensures correct ownership and permissions on PGDATA.
EOF
}

# Parse args
while [[ $# -gt 0 ]]; do
   case "$1" in
      --pgdata)             PGDATA="$2"; shift 2 ;;
      --pgport)             PGPORT="$2"; shift 2 ;;
      --locale)             DB_LOCALE="$2"; shift 2 ;;
      --encoding)           DB_ENCODING="$2"; shift 2 ;;
      --auth)               AUTH_METHOD="$2"; shift 2 ;;
      --postgres-password)  POSTGRES_PASSWORD="$2"; shift 2 ;;
      --admin-user)         ADMIN_USER="$2"; shift 2 ;;
      --admin-pass)         ADMIN_PASS="$2"; shift 2 ;;
      --user-user)          STANDARD_USER="$2"; shift 2 ;;
      --user-pass)          STANDARD_PASS="$2"; shift 2 ;;
      --temp-port)          TEMP_PORT="$2"; shift 2 ;;
      --socket-dir)         SOCKET_DIR="$2"; shift 2 ;;
      --rolefile-path)      ROLEFILE_PATH="$2"; shift 2 ;;
      --force)              force=true; shift 1 ;;
      --start)              start_cluster=true; shift 1 ;;
      -h|--help|-?)         show_usage; exit 0 ;;
      *) echo "ERROR: Unknown option: $1"; show_usage; exit 1 ;;
   esac
done

function force_drop()
{
   echo "INFO: Checking if cluster is running..."
   if ${pg_path}/pg_ctl -D "$PGDATA" status >/dev/null 2>&1; then
     echo "INFO: Stopping cluster..."
     run_as_postgres ${pg_path}/pg_ctl -D "$PGDATA" stop -m fast
   else
     echo "INFO: Cluster not running."
   fi

   # Delete data directory
   if [[ -d "$PGDATA" ]]; then
      echo "INFO: Removing data directory $PGDATA"
      sudo rm -rf "$PGDATA"
   fi
}

# ---- Helper: run command as the postgres system user when available ----
function run_as_postgres()
{
   if command -v gosu >/dev/null && [ "$(id -u)" = "0" ]; then
      gosu postgres "$@"
   elif [ "$(id -un)" = "postgres" ]; then
      "$@"
   else
      sudo -u postgres "$@"
   fi
}

# ---- Checks ----
function need()
{
   command -v "$1" >/dev/null 2>&1 || { echo "ERROR: Missing required command: $1" >&2; exit 1; }
}

# Setup full path, since sudo gets a PATH without PG
pg_path="/usr/lib/postgresql/${PG_VERSION}/bin"
need ${pg_path}/initdb
need ${pg_path}/pg_ctl
need ${pg_path}/psql

# ---- Prepare PGDATA (ownership and permissions) ----
PGDATA="$(readlink -f "$PGDATA")"
CLUSTER_NAME="${PGDATA##*/}"
PG_BASE="${PGDATA%/*}"

# ---- Check if cluster exists. ----
if run_as_postgres test -s "$PGDATA/PG_VERSION"; then
   if [ "$force" == true ]; then
      echo "WARNING: Cluster $CLUSTER_NAME already exists at $PGDATA, but --force specified. Dropping cluster."
      force_drop
   else
      echo "WARNING: Cluster $CLUSTER_NAME already exists at $PGDATA; exiting."
      exit 1
   fi
fi

sudo mkdir -p "$PGDATA"
# Ensure ownership by postgres user (UID/GID may be 999:999 inside official postgres image)
if id postgres >/dev/null 2>&1; then
   sudo chown -R postgres:postgres "$PGDATA" || true
fi
sudo chmod 700 "$PGDATA" || true

echo "INFO: Initializing cluster in $PGDATA (locale=$DB_LOCALE, encoding=$DB_ENCODING)..."

PWFILE=""
if [[ -n "$POSTGRES_PASSWORD" ]]; then
   PWFILE="$(mktemp)"
   printf '%s' "$POSTGRES_PASSWORD" > "$PWFILE"
   sudo chown postgres:postgres "$PWFILE"
fi

initdb_args=( -D "$PGDATA" --locale="$DB_LOCALE" --encoding="$DB_ENCODING" )
# Set bootstrap superuser name explicitly (default is 'postgres' anyway)
initdb_args+=( -U postgres )
if [[ -n "$PWFILE" ]]; then
   initdb_args+=( --pwfile="$PWFILE" )
fi

run_as_postgres ${pg_path}/initdb "${initdb_args[@]}"

[[ -n "$PWFILE" ]] && run_as_postgres rm -f "$PWFILE"

# ---- Configure password_encryption if SCRAM selected ----
POSTGRESQL_CONF="${PGDATA}/postgresql.conf"
echo "INFO: Configuring $POSTGRESQL_CONF with auth method: $AUTH_METHOD"
if [[ "$AUTH_METHOD" == "scram-sha-256" ]]; then
   #if ! grep -q '^[# ]*password_encryption[[:space:]]*=' "$POSTGRESQL_CONF" 2>/dev/null; then
   #   echo "password_encryption = scram-sha-256" | run_as_postgres tee -a "$POSTGRESQL_CONF" > /dev/null
   #else
      #run_as_postgres sed -i -e "s|^[# ]*password_encryption[[:space:]]*=.*|password_encryption = scram-sha-256|g" "$POSTGRESQL_CONF"
      run_as_postgres sed -i -e 's/^\s*#\s*\(password_encryption\s*=\s*\).*$/\1scram-sha-256/g' "$POSTGRESQL_CONF"
   #fi
fi
run_as_postgres sed -i -e 's/^\s*#\(\s*log_destination\)\( *=\)\(.*\)$/\1\2\3/g' \
                       -e 's/^\s*#\(\s*logging_collector\)\( *=\)\(.*\)$/\1\2\3/g' \
                       -e 's/^\s*#\(\s*log_directory\)\( *=\)\(.*\)$/\1\2\3/g' \
                       -e 's/^\s*#\(\s*log_filename\)\( *=\)\(.*\)$/\1\2\3/g' \
                       -e 's/^\s*#\(\s*log_file_mode\)\( *=\)\(.*\)$/\1\2\3/g' \
                       -e 's/^\s*#\(\s*log_rotation_age\)\( *=\)\(.*\)$/\1\2\3/g' \
                       -e 's/^\s*#\(\s*log_rotation_size\)\( *=\)\(.*\)$/\1\2\3/g' \
                       -e "s/^\s*#\(\s*port\)\( *=\)\(.*\)$/\1\2 ${PGPORT}/g" \
                "$POSTGRESQL_CONF"

# Ensure server will accept TCP later (listen on all; comment/uncomment as needed)
run_as_postgres sed -i -e 's/^\s*#\s*\(listen_addresses\s*=\s*\).*$/\1'\''*'\''/g' "$POSTGRESQL_CONF"

# ---- Configure pg_hba.conf ----
PG_HBA="${PGDATA}/pg_hba.conf"
echo "INFO: Configuring $PG_HBA with auth method: $AUTH_METHOD"
run_as_postgres sed -i -e "s/\(^local\s\+all\s\+all\s\+\)trust$/\1${AUTH_METHOD}/g" \
                       -e "s/\(^host\s\+all\s\+all\s\+127\.0\.0\.1\/32\s\+\)trust$/\1${AUTH_METHOD}/g" \
                       -e "s/\(^host\s\+all\s\+all\s\+\:\:1\/128\s\+\)trust$/\1${AUTH_METHOD}/g" \
                       -e "s/\(^local\s\+replication\s\+all\s\+\)trust$/\1${AUTH_METHOD}/g" \
                       -e "s/\(^host\s\+replication\s\+all\s\+127\.0\.0\.1\/32\s\+\)trust$/\1${AUTH_METHOD}/g" \
                       -e "s/\(^host\s\+replication\s\+all\s\+\:\:1\/128\s\+\)trust$/\1${AUTH_METHOD}/g" \
                   "$PG_HBA"
# Keep any existing lines, but ensure we have our standard baseline at the end
echo "host all all all $AUTH_METHOD" | run_as_postgres tee -a "$PG_HBA" > /dev/null

# Create a temporary pg_hba.conf with trust authentication
BOOTSTRAP_HBA="${PGDATA}/pg_hba.bootstrap.conf"
sudo tee "$BOOTSTRAP_HBA" > /dev/null <<'HBA'
# Bootstrap-only HBA (trust all local connections)
local   all   all                 trust
host    all   all   127.0.0.1/32  trust
host    all   all   ::1/128       trust
HBA
sudo chown -R postgres:postgres "$BOOTSTRAP_HBA"
sudo chmod 600 "$BOOTSTRAP_HBA"
sudo cp "${ROLEFILE_PATH}/create_roles.sql" /tmp
sudo chown postgres:postgres /tmp/create_roles.sql

# ---- Bootstrap with HBA locally (socket only) to create roles ----
echo "INFO: Starting temporary server for role creation..."
run_as_postgres ${pg_path}/pg_ctl -D "$PGDATA" \
  -o "-c listen_addresses='' -c unix_socket_directories='$SOCKET_DIR' -p $TEMP_PORT -c hba_file=$BOOTSTRAP_HBA" \
  -w start

# Setup 'set' variables
PSQL_ARGS=(
  --set=std_user="${STANDARD_USER}"
  --set=std_pass="${STANDARD_PASS:-}"
  --set=admin_user="${ADMIN_USER}"
  --set=admin_pass="${ADMIN_PASS:-}"
)
[[ -n "$POSTGRES_PASSWORD" ]] && PSQL_ARGS+=(--set=pg_pass="$POSTGRES_PASSWORD")

# Run external SQL file
echo "INFO: Creating roles..."
run_as_postgres psql -h "$SOCKET_DIR" -p "$TEMP_PORT" -U postgres -d postgres \
                     "${PSQL_ARGS[@]}" \
                     -v ON_ERROR_STOP=1 \
                     -f "/tmp/create_roles.sql"

# Stop temp server
echo "INFO: Stopping temporary server..."
run_as_postgres ${pg_path}/pg_ctl -D "$PGDATA" -w stop

# Final ownership/permissions setup
if id postgres >/dev/null 2>&1; then
   sudo chown -R postgres:postgres "$PGDATA" || true
fi
sudo chmod 700 "$PGDATA" || true

# Cleanup bootstrap hba and roles file
sudo rm -f "$BOOTSTRAP_HBA" || true
sudo rm -f "/tmp/create_roles.sql" || true

# Start the cluster, if requested
[ "$start_cluster" == true ] && run_as_postgres ${pg_path}/pg_ctl -D ${PGDATA} -w start || true

if ${pg_path}/pg_ctl -D "$PGDATA" status >/dev/null 2>&1; then
   echo ✅ Cluster is running at: $PGDATA
else
# Show some useful information
   cat <<SUM

✅ Cluster is ready to run at: $PGDATA

Next steps:
- To run with the official postgres image:
    docker run -e POSTGRES_PASSWORD='${POSTGRES_PASSWORD:-<set>}' \\
               -v <volume_name>:"$PGDATA" \\
               -p $PGPORT:$PGPORT \\
               postgres:${PG_VERSION}

- Or start locally (outside Docker):
    as postgres:  ${pg_path}/pg_ctl -D "$PGDATA" -w start

Notes:
- listen_addresses='*' has been set. Adjust pg_hba.conf to restrict host access before exposing ports.
- Auth method in pg_hba.conf: $AUTH_METHOD
SUM
fi