#!/usr/bin/env bashio
# shellcheck shell=bash
set -e

CONFIG_HOME="/config"
PGDATA="${PGDATA:-/config/database}"
export PGDATA
PG_MAJOR_VERSION="${PG_MAJOR:-15}"

mkdir -p "$PGDATA"
chown -R postgres:postgres "$PGDATA"
chmod 700 "$PGDATA"
chmod -R 755 "$CONFIG_HOME"

RESTART_NEEDED=false

cd /config || true

bashio::log.info "Starting PostgreSQL..."

if [ "$(bashio::info.arch)" = "armv7" ]; then
    bashio::log.warning "ARMv7 detected: Starting without vectors.so"
    /usr/local/bin/immich-docker-entrypoint.sh postgres & true
    exit 0
else
    /usr/local/bin/immich-docker-entrypoint.sh postgres -c config_file=/etc/postgresql/postgresql.conf & true
fi

bashio::log.info "Waiting for PostgreSQL to start..."

(
wait_for_postgres() {
    local tries=0
    while ! pg_isready -h "$DB_HOSTNAME" -p "$DB_PORT" -U "$DB_USERNAME" >/dev/null 2>&1; do
        tries=$((tries+1))
        if [ "$tries" -ge 60 ]; then
            bashio::log.error "Postgres did not start after 2 minutes, aborting."
            exit 1
        fi
        echo "PostgreSQL is starting up... ($tries/60)"
        sleep 2
    done
}

DB_PORT=5432
DB_HOSTNAME=localhost
DB_PASSWORD="$(bashio::config 'POSTGRES_PASSWORD')"
DB_PASSWORD="$(jq -rn --arg x "$DB_PASSWORD" '$x|@uri')"
DB_USERNAME=postgres
if bashio::config.has_value "POSTGRES_USER"; then
    DB_USERNAME="$(bashio::config "POSTGRES_USER")"
fi
export DB_PORT DB_HOSTNAME DB_USERNAME DB_PASSWORD

wait_for_postgres

# --------- Version detection directly from cluster & databases ----------

get_pgdata_version() {
    if [ -f "$PGDATA/PG_VERSION" ]; then
        cat "$PGDATA/PG_VERSION"
    else
        bashio::log.error "FATAL: $PGDATA/PG_VERSION not found; cannot determine cluster version."
        exit 1
    fi
}

get_available_extension_version() {
    local extname="$1"
    psql "postgres://$DB_USERNAME:$DB_PASSWORD@$DB_HOSTNAME:$DB_PORT/postgres" -v ON_ERROR_STOP=1 -tAc \
        "SELECT default_version FROM pg_available_extensions WHERE name = '$extname';" 2>/dev/null | xargs
}

is_extension_available() {
    local extname="$1"
    local result
    result=$(psql "postgres://$DB_USERNAME:$DB_PASSWORD@$DB_HOSTNAME:$DB_PORT/postgres" -v ON_ERROR_STOP=1 -tAc \
        "SELECT 1 FROM pg_available_extensions WHERE name = '$extname';" 2>/dev/null | xargs)
    [[ "$result" == "1" ]]
}

get_user_databases() {
    psql "postgres://$DB_USERNAME:$DB_PASSWORD@$DB_HOSTNAME:$DB_PORT/postgres" -v ON_ERROR_STOP=1 -tAc \
        "SELECT datname FROM pg_database WHERE datistemplate = false AND datallowconn = true;"
}

get_installed_extension_version() {
    local extname="$1"
    local dbname="$2"
    psql "postgres://$DB_USERNAME:$DB_PASSWORD@$DB_HOSTNAME:$DB_PORT/$dbname" -v ON_ERROR_STOP=1 -tAc \
        "SELECT extversion FROM pg_extension WHERE extname = '$extname';" 2>/dev/null | xargs
}

compare_versions() {
    # Returns 0 if $1 < $2, 1 otherwise
    local v1="$1"
    local v2="$2"
    if [ "$v1" = "$v2" ]; then return 1; fi
    if [ "$(printf '%s\n' "$v1" "$v2" | sort -V | head -n1)" = "$v1" ]; then
        return 0
    fi
    return 1
}

# --------- Main logic ----------

upgrade_postgres_if_needed() {
    CLUSTER_VERSION=$(get_pgdata_version)
    IMAGE_VERSION="$PG_MAJOR_VERSION"

    if [ "$CLUSTER_VERSION" != "$IMAGE_VERSION" ]; then
        bashio::log.warning "Postgres data directory version is $CLUSTER_VERSION but image wants $IMAGE_VERSION. Running upgrade..."

        export DATA_DIR="$PGDATA"
        export BINARIES_DIR="/usr/lib/postgresql/$IMAGE_VERSION/bin"
        export BACKUP_DIR="/config/backups"
        export PSQL_VERSION="$IMAGE_VERSION"

        apt-get update &>/dev/null
        apt-get install -y procps rsync postgresql-$IMAGE_VERSION postgresql-$CLUSTER_VERSION &>/dev/null

        TMP_SCRIPT=$(mktemp)
        wget https://raw.githubusercontent.com/linkyard/postgres-upgrade/refs/heads/main/upgrade-postgres.sh -O "$TMP_SCRIPT"
        chmod +x "$TMP_SCRIPT"
        "$TMP_SCRIPT"
        UPGRADE_STATUS=$?
        if [ "$UPGRADE_STATUS" -ne 0 ]; then
            bashio::log.error "Postgres major version upgrade failed. Aborting startup."
            exit 1
        fi

        bashio::log.info "PostgreSQL major version upgrade complete."
        RESTART_NEEDED=true

        # Wait for server to come back online after upgrade
        wait_for_postgres
    else
        bashio::log.info "PostgreSQL data directory version ($CLUSTER_VERSION) matches image version ($IMAGE_VERSION)."
    fi
}

upgrade_extension_if_needed() {
    local extname="$1"
    if ! is_extension_available "$extname"; then
        bashio::log.info "$extname extension not available on this Postgres instance."
        return
    fi
    local available_version
    available_version=$(get_available_extension_version "$extname")
    if [ -z "$available_version" ]; then
        bashio::log.info "Could not determine available version for $extname."
        return
    fi
    for db in $(get_user_databases); do
        local installed_version
        installed_version=$(get_installed_extension_version "$extname" "$db")
        if [ -n "$installed_version" ]; then
            compare_versions "$installed_version" "$available_version"
            if [ $? -eq 0 ]; then
                bashio::log.info "Upgrading $extname in $db from $installed_version to $available_version"
                psql "postgres://$DB_USERNAME:$DB_PASSWORD@$DB_HOSTNAME:$DB_PORT/$db" -v ON_ERROR_STOP=1 -c "ALTER EXTENSION $extname UPDATE;" \
                    && RESTART_NEEDED=true \
                    || { bashio::log.error "Failed to upgrade $extname in $db. Aborting startup."; exit 1; }
            else
                bashio::log.info "$extname in $db already at latest version ($installed_version)"
            fi
        fi
    done
}

# --------- Run logic ----------

upgrade_postgres_if_needed

upgrade_extension_if_needed "vectors"
upgrade_extension_if_needed "vchord"

if [ "$RESTART_NEEDED" = true ]; then
    bashio::log.warning "A critical update (Postgres or extension) occurred. Restarting the addon for changes to take effect."
    bashio::addon.restart
    exit 0
fi

bashio::log.info "All initialization/version check steps completed successfully!"
) & true