#!/bin/bash
set -e

PG_VERSION=15
PG_BIN=/usr/lib/postgresql/$PG_VERSION/bin
PG_DATA=/data/pgdata
PG_LOG=/data/pgdata/pg.log

: "${POSTGRES_USER:=root}"
: "${POSTGRES_PASSWORD:=123456}"
: "${POSTGRES_DB:=new-api}"
: "${SQL_DSN:=postgresql://root:123456@127.0.0.1:5432/new-api?sslmode=disable}"

export LANG=en_US.UTF-8
export LC_ALL=en_US.UTF-8
export SQL_DSN

mkdir -p "$PG_DATA"
chown -R postgres:postgres "$PG_DATA"

if [ ! -f "$PG_DATA/PG_VERSION" ]; then
    echo "Initializing PostgreSQL data directory..."

    su postgres -c "$PG_BIN/initdb -D '$PG_DATA' --locale=en_US.UTF-8"

    if [ ! -f "$PG_DATA/postgresql.conf" ]; then
        echo "FATAL: initdb did not create postgresql.conf" >&2
        exit 1
    fi
fi

echo "Writing PostgreSQL configuration..."

cat > "$PG_DATA/pg_hba.conf" << EOF
local   all   all                 trust
host    all   all   127.0.0.1/32  scram-sha-256
host    all   all   ::1/128       scram-sha-256
EOF

# Listen only on localhost
if ! grep -q "^listen_addresses" "$PG_DATA/postgresql.conf" 2>/dev/null; then
    echo "listen_addresses = '127.0.0.1,::1'" >> "$PG_DATA/postgresql.conf"
fi

# Optimize for Hugging Face Spaces / slow persistent storage
if ! grep -q "^fsync = off" "$PG_DATA/postgresql.conf" 2>/dev/null; then
    echo "Applying PostgreSQL storage optimizations..."

    cat >> "$PG_DATA/postgresql.conf" << EOF

# Hugging Face Spaces storage optimizations
fsync = off
synchronous_commit = off
full_page_writes = off
EOF
fi

# Remove stale PID file from previous run
if [ -f "$PG_DATA/postmaster.pid" ]; then
    OLD_PID=$(head -1 "$PG_DATA/postmaster.pid")

    if ! kill -0 "$OLD_PID" 2>/dev/null; then
        echo "Removing stale PostgreSQL PID file, old PID: $OLD_PID"
        rm -f "$PG_DATA/postmaster.pid"
    fi
fi

# Fix permissions on every start
chown -R postgres:postgres "$PG_DATA"
chmod 0700 "$PG_DATA"

# Recreate missing critical directories
# HF persistent storage may sometimes lose empty dirs after restart
for dir in \
    base \
    global \
    pg_commit_ts \
    pg_dynshmem \
    pg_logical \
    pg_logical/snapshots \
    pg_logical/mappings \
    pg_multixact \
    pg_multixact/members \
    pg_multixact/offsets \
    pg_notify \
    pg_replslot \
    pg_serial \
    pg_snapshots \
    pg_stat \
    pg_stat_tmp \
    pg_subtrans \
    pg_tblspc \
    pg_twophase \
    pg_wal \
    pg_wal/archive_status; do

    if [ ! -d "$PG_DATA/$dir" ]; then
        echo "Recreating missing directory: $dir"
        mkdir -p "$PG_DATA/$dir"
    fi
done

chown -R postgres:postgres "$PG_DATA"

rm -f "$PG_LOG"

echo "Starting PostgreSQL..."

su postgres -c "$PG_BIN/pg_ctl -D '$PG_DATA' -l '$PG_LOG' start" || {
    echo "=== PostgreSQL log ==="
    cat "$PG_LOG" 2>/dev/null || echo "(no PostgreSQL log file)"
    exit 1
}

echo "Waiting for PostgreSQL to be ready..."

for i in $(seq 1 120); do
    if su postgres -c "$PG_BIN/pg_isready -q -h 127.0.0.1 -p 5432"; then
        echo "PostgreSQL is ready."
        break
    fi

    if [ "$i" -eq 120 ]; then
        echo "PostgreSQL startup timeout."
        echo "=== PostgreSQL log ==="
        cat "$PG_LOG" 2>/dev/null || echo "(no PostgreSQL log file)"
        exit 1
    fi

    sleep 1
done

echo "Creating or fixing database user and database..."

# Create role if missing
if ! su postgres -c "$PG_BIN/psql -tAc \"SELECT 1 FROM pg_roles WHERE rolname='$POSTGRES_USER'\"" | grep -q 1; then
    echo "Creating PostgreSQL user: $POSTGRES_USER"
    su postgres -c "$PG_BIN/psql -c \"CREATE USER \\\"$POSTGRES_USER\\\" WITH PASSWORD '$POSTGRES_PASSWORD'\""
fi

# Always reset password
# This is important when /data/pgdata already exists with an old password
echo "Resetting PostgreSQL password for user: $POSTGRES_USER"
su postgres -c "$PG_BIN/psql -c \"ALTER USER \\\"$POSTGRES_USER\\\" WITH PASSWORD '$POSTGRES_PASSWORD'\""

# Create database if missing
if ! su postgres -c "$PG_BIN/psql -tAc \"SELECT 1 FROM pg_database WHERE datname='$POSTGRES_DB'\"" | grep -q 1; then
    echo "Creating PostgreSQL database: $POSTGRES_DB"
    su postgres -c "$PG_BIN/psql -c \"CREATE DATABASE \\\"$POSTGRES_DB\\\" OWNER \\\"$POSTGRES_USER\\\"\""
fi

# Ensure owner is correct
echo "Fixing database owner..."
su postgres -c "$PG_BIN/psql -c \"ALTER DATABASE \\\"$POSTGRES_DB\\\" OWNER TO \\\"$POSTGRES_USER\\\"\""

echo "Testing PostgreSQL login..."

PGPASSWORD="$POSTGRES_PASSWORD" "$PG_BIN/psql" \
    "postgresql://$POSTGRES_USER@127.0.0.1:5432/$POSTGRES_DB?sslmode=disable" \
    -c "SELECT current_user, current_database();" || {
        echo "FATAL: PostgreSQL login test failed."
        echo "=== PostgreSQL log ==="
        cat "$PG_LOG" 2>/dev/null || echo "(no PostgreSQL log file)"
        exit 1
    }

cd /data

echo "Starting NewAPI..."
echo "SQL_DSN=$SQL_DSN"

/new-api &
NEWAPI_PID=$!

cleanup() {
    echo "Shutting down..."

    if kill -0 "$NEWAPI_PID" 2>/dev/null; then
        kill "$NEWAPI_PID" 2>/dev/null || true
        wait "$NEWAPI_PID" 2>/dev/null || true
    fi

    su postgres -c "$PG_BIN/pg_ctl -D '$PG_DATA' stop -m fast" || true

    exit 0
}

trap cleanup SIGTERM SIGINT

wait "$NEWAPI_PID"
EXIT_CODE=$?

echo "NewAPI exited with code $EXIT_CODE"

su postgres -c "$PG_BIN/pg_ctl -D '$PG_DATA' stop -m fast" || true

exit "$EXIT_CODE"