#!/bin/bash

cd $(dirname "$0")

readonly BRAND='orchid'
readonly POSTGRES_ADMIN='orchid_postgres_admin'
readonly BIN_DIR="/opt/$BRAND/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

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 $?
}

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

    # 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"
    local special_runtime_options=$(get_property_value smart_search.postgres.runtime_options)

    echo "$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 "$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 "$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 "$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

        # Make sure any processes owned by $POSTGRES_ADMIN are stopped before proceeding.
        killall -u "$POSTGRES_ADMIN" -QUIT

        # Wipe the database. After the issue is resolved, this will allow subsequent
        # attempts to configure the database correctly.
        rm -rf "$POSTGRES_STORAGE_LOCATION"

        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
        { 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.
        run_as_pg_admin "$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 to provide a 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_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 "$BIN_DIR/psql -h localhost -p $POSTGRES_PORT -d smart_search"
else
    pg_cmd $pg_op
fi
