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

log() {
    printf '[postgres-ha-bootstrap] %s\n' "$*" >&2
}

die() {
    log "ERROR: $*"
    exit 1
}

require_var() {
    local name="$1"
    [[ -n "${!name:-}" ]] || die "Required environment variable $name is not set"
}

require_file() {
    local path="$1"
    [[ -r "$path" ]] || die "Required file is not readable: $path"
}

as_postgres() {
    runuser -u postgres -- "$@"
}

# PostgreSQL replication slot names are intentionally restricted here
# so values can safely be used in SQL and configuration.
validate_slot_name() {
    [[ "$1" =~ ^[a-z0-9_]+$ ]] ||
        die "Invalid replication slot name: $1"
}

install_secret_for_postgres() {
    local source="$1"
    local destination="$2"

    require_file "$source"

    install \
        -o postgres \
        -g postgres \
        -m 0600 \
        "$source" \
        "$destination"
}

pgdata_is_initialized() {
    [[ -s "$PGDATA/PG_VERSION" ]] &&
    [[ -s "$PGDATA/global/pg_control" ]]
}

pgdata_is_empty() {
    [[ ! -d "$PGDATA" ]] && return 0

    local first
    first="$(find "$PGDATA" -mindepth 1 -maxdepth 1 -print -quit 2>/dev/null || true)"
    [[ -z "$first" ]]
}

marker_path() {
    printf '%s/%s\n' \
        "$PGDATA" \
        "${POSTGRES_BOOTSTRAP_MARKER:-.loxthelion-ha-initialized}"
}

write_managed_config() {
    local config="$PGDATA/postgres-ha.conf"

    cat > "$config" <<EOF_CONFIG
# Managed by postgres-ha-bootstrap.
# Do not place node-role-specific settings in this file.

listen_addresses = '*'
wal_level = '${POSTGRES_WAL_LEVEL:-replica}'
max_wal_senders = ${POSTGRES_MAX_WAL_SENDERS:-10}
max_replication_slots = ${POSTGRES_MAX_REPLICATION_SLOTS:-10}
hot_standby = ${POSTGRES_HOT_STANDBY:-on}
password_encryption = 'scram-sha-256'
EOF_CONFIG

    chown postgres:postgres "$config"
    chmod 0600 "$config"

    if ! grep -Eq \
        "^[[:space:]]*include_if_exists[[:space:]]*=[[:space:]]*'postgres-ha\.conf'" \
        "$PGDATA/postgresql.conf"; then

        cat >> "$PGDATA/postgresql.conf" <<'EOF_CONFIG'

# Loxthelion PostgreSQL HA configuration
include_if_exists = 'postgres-ha.conf'
EOF_CONFIG
    fi
}

write_managed_hba() {
    local hba="$PGDATA/pg_hba.conf"
    local tmp

    tmp="$(mktemp)"

    # Remove a previous block generated by this script.
    awk '
        /^# BEGIN LOXTHELION POSTGRES HA$/ { skip=1; next }
        /^# END LOXTHELION POSTGRES HA$/   { skip=0; next }
        !skip { print }
    ' "$hba" > "$tmp"

    cat >> "$tmp" <<EOF_HBA

# BEGIN LOXTHELION POSTGRES HA
host    replication    ${POSTGRES_REPLICATION_USER}    ${POSTGRES_HA_CIDR}    scram-sha-256
host    postgres       ${POSTGRES_REPLICATION_USER}    ${POSTGRES_HA_CIDR}    scram-sha-256
# END LOXTHELION POSTGRES HA
EOF_HBA

    cat "$tmp" > "$hba"
    rm -f "$tmp"

    chown postgres:postgres "$hba"
    chmod 0600 "$hba"
}

create_persistent_pgpass() {
    require_var POSTGRES_PRIMARY_HOST
    require_var POSTGRES_PRIMARY_PORT
    require_var POSTGRES_REPLICATION_USER
    require_var POSTGRES_REPLICATION_PASSWORD_FILE

    require_file "$POSTGRES_REPLICATION_PASSWORD_FILE"

    local password
    password="$(cat "$POSTGRES_REPLICATION_PASSWORD_FILE")"

    mkdir -p /var/lib/postgresql
    chown postgres:postgres /var/lib/postgresql

    printf '%s:%s:*:%s:%s\n' \
        "$POSTGRES_PRIMARY_HOST" \
        "$POSTGRES_PRIMARY_PORT" \
        "$POSTGRES_REPLICATION_USER" \
        "$password" \
        > /var/lib/postgresql/.pgpass

    chown postgres:postgres /var/lib/postgresql/.pgpass
    chmod 0600 /var/lib/postgresql/.pgpass
}

bootstrap_primary() {
    require_var POSTGRES_DB
    require_var POSTGRES_USER
    require_var POSTGRES_REPLICATION_USER
    require_var POSTGRES_HA_CIDR
    require_var POSTGRES_PASSWORD_FILE
    require_var POSTGRES_REPLICATION_PASSWORD_FILE

    require_file "$POSTGRES_PASSWORD_FILE"
    require_file "$POSTGRES_REPLICATION_PASSWORD_FILE"

    pgdata_is_empty ||
        die "primary bootstrap requires an empty PGDATA"

    log "Initializing fresh PostgreSQL primary"

    mkdir -p "$PGDATA"
    chown -R postgres:postgres /var/lib/postgresql
    chmod 0700 "$PGDATA"

    local pwfile
    pwfile="/tmp/postgres-bootstrap-password"

    install_secret_for_postgres \
        "$POSTGRES_PASSWORD_FILE" \
        "$pwfile"

    as_postgres initdb \
        -D "$PGDATA" \
        --username="$POSTGRES_USER" \
        --pwfile="$pwfile" \
        --auth-host=scram-sha-256 \
        --auth-local=trust \
        --data-checksums

    rm -f "$pwfile"

    write_managed_config
    write_managed_hba

    log "Starting PostgreSQL temporarily for initial role/database setup"

    as_postgres pg_ctl \
        -D "$PGDATA" \
        -w \
        start

    trap 'as_postgres pg_ctl -D "$PGDATA" -m fast -w stop >/dev/null 2>&1 || true' RETURN

    local repl_password
    repl_password="$(cat "$POSTGRES_REPLICATION_PASSWORD_FILE")"

    # psql variables keep the password out of SQL construction itself.
    as_postgres psql \
        --username="$POSTGRES_USER" \
        --dbname=postgres \
        -v ON_ERROR_STOP=1 \
        --set=repl_user="$POSTGRES_REPLICATION_USER" \
        --set=repl_password="$repl_password" <<'SQL'
SELECT format(
    'CREATE ROLE %I WITH LOGIN REPLICATION PASSWORD %L',
    :'repl_user',
    :'repl_password'
)
WHERE NOT EXISTS (
    SELECT 1
    FROM pg_roles
    WHERE rolname = :'repl_user'
)
\gexec
SQL

    if ! as_postgres psql \
        --username="$POSTGRES_USER" \
        --dbname=postgres \
        --tuples-only \
        --no-align \
        --set=ON_ERROR_STOP=1 \
        --set=db="$POSTGRES_DB" <<'SQL' | grep -qx 1
SELECT 1
FROM pg_database
WHERE datname = :'db';
SQL
    then
        as_postgres createdb \
            --username="$POSTGRES_USER" \
            --owner="$POSTGRES_USER" \
            "$POSTGRES_DB"
    fi

    as_postgres pg_ctl \
        -D "$PGDATA" \
        -m fast \
        -w \
        stop

    trap - RETURN

    touch "$(marker_path)"
    chown postgres:postgres "$(marker_path)"
    chmod 0600 "$(marker_path)"

    log "Fresh primary bootstrap complete"
}

ensure_remote_replication_slot() {
    validate_slot_name "$POSTGRES_REPLICATION_SLOT"

    local slot_type

    slot_type="$(
        as_postgres env \
            PGPASSFILE=/var/lib/postgresql/.pgpass \
            psql \
            --host="$POSTGRES_PRIMARY_HOST" \
            --port="$POSTGRES_PRIMARY_PORT" \
            --username="$POSTGRES_REPLICATION_USER" \
            --dbname=postgres \
            --tuples-only \
            --no-align \
            --set=ON_ERROR_STOP=1 \
            --command="
                SELECT slot_type
                FROM pg_replication_slots
                WHERE slot_name = '$POSTGRES_REPLICATION_SLOT';
            " |
        tr -d '[:space:]'
    )"

    case "$slot_type" in
        physical)
            log "Replication slot already exists: $POSTGRES_REPLICATION_SLOT"
            ;;

        "")
            log "Creating physical replication slot: $POSTGRES_REPLICATION_SLOT"

            as_postgres env \
                PGPASSFILE=/var/lib/postgresql/.pgpass \
                psql \
                --host="$POSTGRES_PRIMARY_HOST" \
                --port="$POSTGRES_PRIMARY_PORT" \
                --username="$POSTGRES_REPLICATION_USER" \
                --dbname=postgres \
                --set=ON_ERROR_STOP=1 \
                --command="
                    SELECT pg_create_physical_replication_slot(
                        '$POSTGRES_REPLICATION_SLOT'
                    );
                "
            ;;

        *)
            die "Slot $POSTGRES_REPLICATION_SLOT exists but is not physical"
            ;;
    esac
}

bootstrap_standby() {
    require_var POSTGRES_PRIMARY_HOST
    require_var POSTGRES_PRIMARY_PORT
    require_var POSTGRES_REPLICATION_USER
    require_var POSTGRES_REPLICATION_SLOT
    require_var POSTGRES_APPLICATION_NAME
    require_var POSTGRES_REPLICATION_PASSWORD_FILE

    pgdata_is_empty ||
        die "standby bootstrap requires an empty PGDATA"

    validate_slot_name "$POSTGRES_REPLICATION_SLOT"

    log "Preparing fresh standby from $POSTGRES_PRIMARY_HOST"

    mkdir -p "$PGDATA"
    chmod 0700 "$PGDATA"
    chown -R postgres:postgres /var/lib/postgresql

    create_persistent_pgpass
    ensure_remote_replication_slot

    log "Running pg_basebackup"

    as_postgres env \
        PGPASSFILE=/var/lib/postgresql/.pgpass \
        pg_basebackup \
        --dbname="host=$POSTGRES_PRIMARY_HOST port=$POSTGRES_PRIMARY_PORT user=$POSTGRES_REPLICATION_USER application_name=$POSTGRES_APPLICATION_NAME passfile=/var/lib/postgresql/.pgpass" \
        --pgdata="$PGDATA" \
        --format=plain \
        --wal-method=stream \
        --slot="$POSTGRES_REPLICATION_SLOT" \
        --write-recovery-conf \
        --progress

    [[ -e "$PGDATA/standby.signal" ]] ||
        die "pg_basebackup completed without standby.signal"

    pgdata_is_initialized ||
        die "pg_basebackup did not produce a valid PostgreSQL data directory"

    touch "$(marker_path)"
    chown postgres:postgres "$(marker_path)"
    chmod 0600 "$(marker_path)"

    log "Fresh standby bootstrap complete"
}

adopt_existing() {
    pgdata_is_initialized ||
        die "POSTGRES_HA_MODE=existing requires a valid existing PostgreSQL cluster"

    log "Existing PostgreSQL cluster detected"

    write_managed_config

    # HBA management requires these only when supplied.
    if [[ -n "${POSTGRES_REPLICATION_USER:-}" &&
          -n "${POSTGRES_HA_CIDR:-}" ]]; then
        write_managed_hba
    fi

    touch "$(marker_path)"
    chown postgres:postgres "$(marker_path)"
    chmod 0600 "$(marker_path)"

    log "Existing cluster adopted; no initdb or basebackup performed"
}

validate_bootstrapped_cluster() {
    pgdata_is_initialized ||
        die "Bootstrap marker exists but PGDATA is not a valid PostgreSQL cluster"

    local version
    version="$(cat "$PGDATA/PG_VERSION")"

    [[ "$version" == "18" ]] ||
        die "Expected PostgreSQL major version 18, found $version"

    as_postgres pg_controldata "$PGDATA" >/dev/null ||
        die "pg_controldata validation failed"

    local pgdata_mode
    pgdata_mode="$(stat -c '%a' "$PGDATA")"

    [[ "$pgdata_mode" == "700" || "$pgdata_mode" == "750" ]] ||
        die "PGDATA has invalid permissions: $pgdata_mode (expected 700 or 750)"

    [[ "$(stat -c '%U' "$PGDATA")" == "postgres" ]] ||
        die "PGDATA is not owned by postgres"

    log "Existing bootstrapped PostgreSQL 18 cluster validated"
}

main() {
    require_var PGDATA

    local marker
    marker="$(marker_path)"

    log "Node: ${POSTGRES_NODE_NAME:-unknown}"
    log "Mode: ${POSTGRES_HA_MODE:-unset}"
    log "PGDATA: $PGDATA"

    if [[ -e "$marker" ]]; then
        validate_bootstrapped_cluster
    else
        case "${POSTGRES_HA_MODE:-}" in
            primary)
                bootstrap_primary
                ;;

            standby)
                bootstrap_standby
                ;;

            existing)
                adopt_existing
                ;;

            *)
                die "POSTGRES_HA_MODE must be primary, standby, or existing"
                ;;
        esac
    fi

    [[ "$#" -gt 0 ]] ||
        die "No long-running container command supplied"

    log "Bootstrap gate passed; executing: $*"
    exec "$@"
}

main "$@"