#!/bin/bash

cd $(dirname "$0")

readonly PG_MAJOR_VERSION='17'
readonly BRAND='orchid'
readonly POSTGRES_ADMIN='orchid_postgres_admin'
readonly BIN_DIR="/opt/$BRAND/bin"
# Each Postgres major gets its own prefix: pg_upgrade needs both installed,
# and their libraries share filenames.
readonly PG_BIN_DIR="/opt/$BRAND/pg/pg$PG_MAJOR_VERSION/bin"
readonly LIB_DIR=''

get_property_value () {
    local property="$1"

    # Get the property line. Ignore any preceding whitespace.
    local property_line=$(grep "^\s*$property" "$PROPERTIES_FILE")

    # Get the property value. Ignore arbitrary whitespace on either side of the
    # first detected property assignment character (':' or '=').
    local property_value=$(echo "$property_line" | sed -r "s/^\s*$property\s*[:=]\s*(.*$)/\1/")

    # Return the property without a newline and trailing whitespace.
    echo -n "$property_value" | sed -e 's/\s*$//'
}

assert_root_user () {
    if [[ $EUID -ne 0 ]]; then
        echo 'Must run as root.'
        exit 1
    fi
}>&2

error_out () {
    echo "Error: $@"
    exit 1
}>&2

# Logs to stderr, because both service managers discard this script's stdout.
log () {
    echo "$@"
}>&2

show_usage () {
    cat << EOF
usage: $(basename $0) [SCRIPT OPTION] <DB OPERATION>

SCRIPT OPTIONS

  -c, --config <path/to/properties/directory>
          Path to directory containing properties file.
          - If option is not specified, the following directories are checked, in order:
            1) $PWD
            2) /etc/opt

DB OPERATIONS

  cli
  start
  restart
  stop
  status
EOF
    exit
}

using_sysvinit () {
    [[ $(cat /proc/1/comm) != 'systemd' ]]
    return $?
}

readonly UPGRADE_WAIT_POLL_SECONDS=10

# The cluster's major version, or nothing if there is no cluster.
get_cluster_major_version () {
    local version_file="$POSTGRES_STORAGE_LOCATION/PG_VERSION"
    [[ -f "$version_file" ]] || return 0

    # initdb writes PG_VERSION once, so it is reliable with no server running, even after an unclean shutdown.
    echo "$(< "$version_file")"
}

# Path of orchid_server's upgrade journal. It exists from the start of an upgrade until the upgrade completes,
# and its existence alone means this script must wait.
get_upgrade_journal_path () {
    # dirname and basename ignore a trailing `/`, so the journal never lands inside the storage directory.
    local cluster_dir="$POSTGRES_STORAGE_LOCATION"
    echo "$(dirname "$cluster_dir")/.$(basename "$cluster_dir").upgrade_journal.sqlite"
}

upgrade_is_in_progress () {
    [[ -e "$POSTGRES_UPGRADE_JOURNAL" ]]
}

cluster_needs_upgrade () {
    local cluster_version="$(get_cluster_major_version)"

    # No cluster means nothing to upgrade; setup_smart_search_db creates one.
    [[ -n "$cluster_version" ]] && [[ "$cluster_version" != "$PG_MAJOR_VERSION" ]]
}

# No cluster, but orchid_server's upgrade staging directories remain: an interrupted upgrade, with the data
# inside one of them. setup_smart_search_db would otherwise run initdb here,
# and only initdb's refusal of a non-empty directory would keep it off the data.
#
# Keyed on the staging directories rather than a non-empty location,
# because a location that is a mount point holds lost+found even on a fresh install.
# These names must match new_cluster_dir() and old_cluster_dir() in postgres_cluster_layout.cpp.
cluster_is_missing_but_staging_remains () {
    [[ -z "$(get_cluster_major_version)" ]] || return 1

    [[ -d "$POSTGRES_STORAGE_LOCATION/cluster.old" ]] || [[ -d "$POSTGRES_STORAGE_LOCATION/cluster.new" ]]
}

# orchid_server performs Postgres major upgrades. Until the cluster is one these binaries may open,
# leave it alone: wait (systemd start) or exit, depending on the operation.
await_upgraded_cluster () {
    local op="$1"

    if ! upgrade_is_in_progress && ! cluster_needs_upgrade && ! cluster_is_missing_but_staging_remains; then
        return 0
    fi

    local cluster_version="$(get_cluster_major_version)"
    local situation
    if upgrade_is_in_progress; then
        situation='a Postgres upgrade is in progress'
    elif cluster_is_missing_but_staging_remains; then
        situation="$POSTGRES_STORAGE_LOCATION holds no cluster, but staging directories from an interrupted upgrade hold its data"
    elif [[ ! $cluster_version =~ ^[0-9]+$ ]]; then
        situation="the cluster's PG_VERSION holds '$cluster_version', which is not a Postgres major version"
    elif (( 10#$cluster_version > PG_MAJOR_VERSION )); then
        # Reached after a package rollback. Nothing will resolve it, but waiting is still the safe response.
        situation="the cluster is Postgres $cluster_version, which is newer than the Postgres $PG_MAJOR_VERSION this version includes"
    else
        situation="the cluster is Postgres $cluster_version, which will be upgraded to Postgres $PG_MAJOR_VERSION"
    fi

    case $op in
        start|restart)
            if [[ -z $INVOKED_BY_SYSTEMD ]]; then
                # Our caller expects this script to return, so exit instead of waiting.
                # sysvinit neither retries nor gates the rest of the boot on our exit status,
                # so orchid_server still starts and can perform the upgrade.
                # monit re-runs the init script on each check, and that is the wait.
                log "Not starting Postgres, because $situation"
                exit 1
            fi

            # Under systemd this process is the service's main process for its whole life,
            # so waiting here keeps the unit active instead of failing and restarting throughout the upgrade.
            # A stop while waiting is a normal outcome, so it exits 0; a non-zero exit would mark the unit failed.
            trap 'log "Stopped while waiting for the Postgres upgrade to finish"; exit 0' SIGINT SIGTERM

            log "Waiting to start Postgres, because $situation"

            local started=$SECONDS
            while upgrade_is_in_progress || cluster_needs_upgrade || cluster_is_missing_but_staging_remains; do
                # Backgrounded so the traps above run immediately rather than after the sleep.
                sleep $UPGRADE_WAIT_POLL_SECONDS & wait $!
                log "Still waiting for the Postgres upgrade to finish | waited: $((SECONDS - started)) seconds"
            done

            trap - SIGINT SIGTERM
            log "Finished waiting for the Postgres upgrade | cluster version: $(get_cluster_major_version)"
            ;;
        stop)
            # Succeed, because nothing is wrong. orchid_server stops any cluster it started for the upgrade,
            # and failing the service's ExecStop would only mark the unit failed.
            log "Not stopping Postgres, because $situation"
            exit 0
            ;;
        status)
            # 3 matches the "not running" exit status of pg_ctl status.
            log "Postgres is not running, because $situation"
            exit 3
            ;;
        *)
            log "Cannot connect to Postgres, because $situation"
            exit 1
            ;;
    esac
}

run_as_pg_admin () {
    local ld_library_path
    [[ -n "$LIB_DIR" ]] && ld_library_path="LD_LIBRARY_PATH=$LIB_DIR"
    eval "sudo -u $POSTGRES_ADMIN "$ld_library_path $@""
    return $?
}

get_pg_port () {
    local default_port=5448
    local port=$(get_property_value smart_search.postgres.port)
    echo ${port:-$default_port}
}

get_pg_logging_enabled () {
    get_property_value smart_search.postgres.logging
}

get_pg_runtime_options () {
    local logging_enabled=$(get_pg_logging_enabled)
    local logging_runtime_options=""

    # Check if postgres logging is enabled.
    if [[ "$logging_enabled" == "true" ]]; then
        # Prior to 24.9.0, when the SS db was initialized logging was enabled in postgresql.conf with unlimited log storage. 
        # The settings may still be present within postgresql.conf so we need to override them here.
        #
        # These settings create one log per day, and we retain one week of logs.
        logging_runtime_options="-c logging_collector=on -c log_filename=postgresql-%A.log -c log_rotation_age=1440 -c log_rotation_size=0 -c log_truncate_on_rotation=on"
    else
        logging_runtime_options="-c logging_collector=off"
    fi

    # Pin the listener to literal loopback IPs so startup never depends on name resolution.
    # A missing /etc/hosts can otherwise resolve "localhost" to a bogus address that postgres
    # cannot bind, and it refuses to start. Applied here rather than in postgresql.conf so it
    # also holds for pre-existing data dirs.
    local listen_runtime_options="-c listen_addresses=127.0.0.1,::1"

    # Options specified in smart_search.postgres.runtime_options will override any of the above settings
    local default_special_runtime_options="-c max_locks_per_transaction=1825 -c parallel_tuple_cost=0 -c timescaledb.max_tuples_decompressed_per_dml_transaction=0"
    local special_runtime_options=$(get_property_value smart_search.postgres.runtime_options)

    echo "$listen_runtime_options $logging_runtime_options ${special_runtime_options:-$default_special_runtime_options}"
}

wipe_log_dir_if_logging_disabled () {
    local logging_enabled=$(get_pg_logging_enabled)

    if [[ "$logging_enabled" != "true" ]]; then
        rm -rf "$POSTGRES_STORAGE_LOCATION/log"
    fi
}

get_default_properties_file_path () {
    local props_file="${BRAND}_server.properties"
    local path="$PWD/$props_file"
    [[ ! -f "$path" ]] && path="/etc/opt/$props_file"
    echo "$path"
}

pg_cmd () {
    _run_pg_in_foreground_until_finished () {
        local postmaster_pid="$pg_path/postmaster.pid"

        run_as_pg_admin "$PG_BIN_DIR/postgres -D '$pg_path' $runtime_pg_opts > /dev/null" &

        # If for some odd reason a user ran this script manually with
        # INVOKED_BY_SYSTEMD set, handle CTRL+C so postgres actually shuts down.
        trap 'kill -INT $(head -1 $postmaster_pid)' SIGINT

        # Wait on the PID of the backgrounded postgres process so we can return the status code
        # on exit. If we return a non-zero code, systemd should automatically restart this service.
        wait $!

        return $?
    }

    local op=$1
    local pg_path="${2:-$POSTGRES_STORAGE_LOCATION}"
    local run_in_foreground='false'
    local extra_pg_ctl_opts=''

    # Check if postgres is already running
    local running='false'
    run_as_pg_admin "$PG_BIN_DIR/pg_ctl status -D '$pg_path' -o \"-p $POSTGRES_PORT\"" &> /dev/null
    [[ $? -eq 0 ]] && running='true'

    case $op in
        start|restart)
            if [[ $op == 'start' ]]; then
                # noop if we're trying to start and postgres is already running
                [[ $running == 'true' ]] && return 0

                # For systemd services only, we run postgres in the foreground. This
                # allows systemd to track/manage the process without using a pidfile.
                [[ -n $INVOKED_BY_SYSTEMD ]] && run_in_foreground='true'
            fi

            local runtime_pg_opts="-p $POSTGRES_PORT $POSTGRES_RUNTIME_OPTIONS"

            # If we're on sysvinit build, we need to create a PID file.
            if using_sysvinit; then
                local pidfile_path="$POSTGRES_ADMIN_DIR/smart_search_postgres.pid"
                runtime_pg_opts+=" -c external_pid_file='$pidfile_path'"
            fi

            extra_pg_ctl_opts="-o \"${runtime_pg_opts}\""
            ;;
        stop)
            # noop if we're trying to stop and postgres is already stopped
            [[ $running != 'true' ]] && return 0
            ;;
        status)
            ;;
        *)
            # This should be unreachable.
            error_out "${FUNCNAME[0]} was called with an invalid option: $op."
            ;;
    esac

    if [[ $run_in_foreground == 'true' ]]; then
        _run_pg_in_foreground_until_finished
    else
        run_as_pg_admin "$PG_BIN_DIR/pg_ctl $op -D '$pg_path' $extra_pg_ctl_opts"
    fi

    return $?
}

configure_pg_admin () {
    local pg_admin_dir="/var/lib/${BRAND}_server/postgres"
    local pg_admin_passwd_info="$(getent passwd $POSTGRES_ADMIN)"
    local pg_admin_default_shell='/usr/sbin/nologin'

    # We need to disable logins for the postgres admin account. Certain services (e.g. globalprotect) can
    # automatically spawn processes for accounts with interactive shells (e.g. /bin/bash), which we do not
    # want to happen with this postgres account. We achieve this by using nologin as the default shell.
    #
    # Note that prior to v23.3.0, bash was used as the default shell for the postgres account.

    if [[ -z "$pg_admin_passwd_info" ]]; then
        # If the Smart Search admin account doesn't exist, but the group does, remove it. This state was once possible if
        # 'USERGROUPS_ENAB no' was set in /etc/login.defs and the orchid package was uninstalled, then reinstalled again.
        [[ $(getent group $POSTGRES_ADMIN) ]] && groupdel $POSTGRES_ADMIN
        useradd -d "$pg_admin_dir" -r -m -k /dev/null -U -s $pg_admin_default_shell $POSTGRES_ADMIN
    else
        [[ "$pg_admin_passwd_info" != *$pg_admin_default_shell* ]] && usermod -s $pg_admin_default_shell $POSTGRES_ADMIN
        [[ ! -d "$pg_admin_dir" ]] && mkdir -p "$pg_admin_dir" && chown -R $POSTGRES_ADMIN: "$pg_admin_dir"
    fi

    echo "$pg_admin_dir"
}

setup_smart_search_db () {
    _cleanup_and_error_out () {
        # Unset the trap to make this function fail-safe.
        trap '' ERR

        # initdb refuses a non-empty target, so a half-initialized directory would fail every later attempt.
        # Removing it lets the next run start clean.
        #
        # Only a directory this run created, or found empty, is safe to remove.
        # Neither a missing postgresql.conf nor a present PG_VERSION proves an existing one disposable,
        # so it is left for the operator: a service that won't start is recoverable, a deleted cluster is not.
        if [[ $storage_location_was_populated == 'true' ]]; then
            echo "Leaving the existing contents of $POSTGRES_STORAGE_LOCATION untouched."
        else
            # Kill any postgres the failed initdb left holding files open in the directory.
            # The directory held nothing, so it cannot be an upgrade's cluster.
            killall -u "$POSTGRES_ADMIN" -QUIT
            rm -rf "$POSTGRES_STORAGE_LOCATION"
        fi

        exit 1
    }>&2

    trap 'exit 1' ERR

    # If our location has changed, make sure the current database is shutdown.
    local props_file_path_hash=$(echo "$PROPERTIES_FILE" | md5sum | awk '{print $1}')
    local pg_admin_data_symlink="$POSTGRES_ADMIN_DIR/smart_search_${props_file_path_hash}"
    if [[ -e "$pg_admin_data_symlink" ]] && [[ ! "$pg_admin_data_symlink" -ef "$POSTGRES_STORAGE_LOCATION" ]]; then
        pg_cmd 'stop' "$pg_admin_data_symlink"
    fi

    local postgres_conf="$POSTGRES_STORAGE_LOCATION/postgresql.conf"
    if [[ ! -f "$postgres_conf" ]]; then
        # Checked before mkdir can create the directory,
        # because _cleanup_and_error_out removes it only if it held nothing.
        local storage_location_was_populated='false'
        [[ -n "$(ls -A "$POSTGRES_STORAGE_LOCATION" 2> /dev/null)" ]] && storage_location_was_populated='true'

        { mkdir -p "$POSTGRES_STORAGE_LOCATION" && chown -R $POSTGRES_ADMIN: "$POSTGRES_STORAGE_LOCATION"; } || _cleanup_and_error_out

        # Initialize the database. Note that $POSTGRES_STORAGE_LOCATION has to be either previously
        # empty or non-existant, otherwise this command will fail-out.
        #
        # await_upgraded_cluster, not initdb's refusal, is what keeps this off an interrupted upgrade's data.
        run_as_pg_admin "$PG_BIN_DIR/initdb --locale=C -A trust -D '$POSTGRES_STORAGE_LOCATION'" || _cleanup_and_error_out

        # Configure the postgresql.conf file (setup logging, extensions, etc). Logging is disabled by default, but if re-enabled
        # by smart_search.postgres.logging, we want log_file_mode and log_directory to go into effect.
        sed -e "s|^#logging_collector.*$|logging_collector = off|" \
            -e "s|^#log_file_mode.*$|log_file_mode = 0644|" \
            -e "s|^#log_directory.*$|log_directory = '$POSTGRES_STORAGE_LOCATION/log'|" \
            -i "$postgres_conf" || _cleanup_and_error_out
        echo "shared_preload_libraries = 'timescaledb,vector,vchord'" >> "$postgres_conf" || _cleanup_and_error_out
        echo "timescaledb.telemetry_level = 'off'" >> "$postgres_conf" || _cleanup_and_error_out
    fi

    # Postgres requires its data directory and the files within it to be owned by the non-root
    # $POSTGRES_ADMIN account. If something external (e.g. a sysadmin operation) has left any of
    # them owned by a different account, the server won't start.
    chown -R $POSTGRES_ADMIN: "$POSTGRES_STORAGE_LOCATION"

    # We create a symlink as an easy way to track when the database location changes.
    ln -sfn "$POSTGRES_STORAGE_LOCATION" "$pg_admin_data_symlink"
}

[[ $# -eq 0 ]] && show_usage

while [[ $# -gt 0 ]]; do
    case "$1" in
        -c|--config)
            shift
            properties_dir="$1"
            [[ -z "$properties_dir" ]] && show_usage
            properties_file="$properties_dir/${BRAND}_server.properties"
            ;;
        -h|--help)
            show_usage
            ;;
        cli)
            cli='true'
            ;;
        stop)
            pg_op='stop'
            ;;
        start)
            pg_op='start'
            ;;
        restart)
            pg_op='restart'
            ;;
        status)
            pg_op='status'
            ;;
        *)
            error_out "Invalid $(basename $0) option specified: $1. (See --help)"
            ;;
    esac
    shift
done

assert_root_user

readonly PROPERTIES_FILE="${properties_file:-$(get_default_properties_file_path)}"
[[ ! -f $PROPERTIES_FILE ]] && error_out "Unable to locate properties file. (See --help)"

readonly POSTGRES_STORAGE_LOCATION="$(get_property_value smart_search.postgres.storage_location)"
[[ -z $POSTGRES_STORAGE_LOCATION ]] && error_out "Missing smart_search.postgres.storage_location in properties file."

readonly POSTGRES_UPGRADE_JOURNAL="$(get_upgrade_journal_path)"

# Must precede anything that touches the cluster.
await_upgraded_cluster "$pg_op"

readonly POSTGRES_ADMIN_DIR="$(configure_pg_admin)"
readonly POSTGRES_PORT=$(get_pg_port)
readonly POSTGRES_RUNTIME_OPTIONS=$(get_pg_runtime_options)

setup_smart_search_db
wipe_log_dir_if_logging_disabled

if [[ -n $cli ]]; then
    run_as_pg_admin "$PG_BIN_DIR/psql -h localhost -p $POSTGRES_PORT -d smart_search"
else
    pg_cmd $pg_op
fi
