#!/bin/bash
set -e

BOOT_START=$(date +%s)

echo "[entrypoint] mem0 Self-Hosted on HuggingFace Spaces"
echo "==============================================================================="

# ── Configuration ───────────────────────────────────────────────────────────
# Affine-style: PGDATA lives DIRECTLY on the HF persistent volume (/data).
# No restore-rsync, no backup loop — both hammered the FUSE mount until it
# died ("Transport endpoint is not connected"). FUSE quirks handled below:
# empty dirs are dropped and .git/ is injected into tracked dirs, so dirs
# Postgres scans as object stores must be purged at every boot.
MEM0_WORK="/opt/data"                    # Fast local FS: history, sessions, logs
PGDATA="/data/pgdata"                    # Persistent HF bucket (FUSE)
MEM0_UID=1000
MEM0_GID=1000

# ── Singleton guard ──────────────────────────────────────────────────────────
# Hanya satu instance yang menjalankan steps. Instance duplikat (dari HF rapid-
# restart) cukup sleep infinity — tidak exit (exit = HF restart lagi).
# MEM0_ENTRYPOINT_OWNER di-export sebelum exec gosu sehingga instance setelah
# privilege-drop tahu dia adalah kelanjutan sah, bukan duplikat.
ENTRYPOINT_PID_FILE="/tmp/mem0-entrypoint.pid"
ENTRYPOINT_LOG="/tmp/mem0-entrypoint.log"

# Always clear stale PID file at container startup (new container = new PID namespace)
rm -f "$ENTRYPOINT_PID_FILE" "$ENTRYPOINT_LOG" 2>/dev/null || true

if [ -z "${MEM0_ENTRYPOINT_OWNER:-}" ]; then
  # Instance baru tanpa token → cek apakah ada instance lain yang sudah running
  if [ -f "$ENTRYPOINT_PID_FILE" ]; then
    OLD_PID=$(cat "$ENTRYPOINT_PID_FILE" 2>/dev/null || echo "")
    if [ -n "$OLD_PID" ] && kill -0 "$OLD_PID" 2>/dev/null; then
      echo "[entrypoint] Instance $OLD_PID already running — standing by silently."
      sleep infinity
      exit 0
    else
      echo "[entrypoint] Stale PID $OLD_PID (dead) — taking over..."
      rm -f "$ENTRYPOINT_PID_FILE" "$ENTRYPOINT_LOG" 2>/dev/null || true
    fi
  fi
  echo $$ > "$ENTRYPOINT_PID_FILE"
  export MEM0_ENTRYPOINT_OWNER=$$
  # Buat log world-writable sekarang (sebagai root) agar mem0 bisa write setelah gosu
  touch "$ENTRYPOINT_LOG" && chmod 666 "$ENTRYPOINT_LOG"
fi

exec > >(tee -a "$ENTRYPOINT_LOG") 2>&1

# ── Stage 1: Prepare persistent storage (as root) ───────────────────────────
if [ "$(id -u)" = "0" ]; then
  echo ""
  echo "=== 1. Preparing persistent storage ==="

  mkdir -p "$MEM0_WORK"/{history,sessions,logs}
  mkdir -p "$PGDATA"
  chown ${MEM0_UID}:${MEM0_GID} "$MEM0_WORK"/* 2>/dev/null || true
  chown -R ${MEM0_UID}:${MEM0_GID} "$PGDATA"

  # Init if empty
  FRESH_INIT=0
  if [ ! -f "$PGDATA/PG_VERSION" ]; then
    FRESH_INIT=1
    echo "[entrypoint] Initializing PostgreSQL at $PGDATA"
    chmod 700 "$PGDATA"
    gosu mem0 /usr/lib/postgresql/15/bin/initdb -D "$PGDATA" --auth=trust --encoding=UTF8 --locale=C.UTF-8

    # Configure PostgreSQL for local connections
    cat >> "$PGDATA/postgresql.conf" <<EOF
listen_addresses = 'localhost'
port = 5432
max_connections = 100
shared_buffers = 128MB
effective_cache_size = 512MB
maintenance_work_mem = 64MB
checkpoint_completion_target = 0.9
wal_buffers = 16MB
default_statistics_target = 100
random_page_cost = 1.1
effective_io_concurrency = 200
work_mem = 4MB
min_wal_size = 1GB
max_wal_size = 4GB
EOF

    # Allow local connections without password (incl. replication for future use)
    cat > "$PGDATA/pg_hba.conf" <<EOF
local   all             all                                     trust
host    all             all             127.0.0.1/32            trust
host    all             all             ::1/128                 trust
local   replication     all                                     trust
host    replication     all             127.0.0.1/32            trust
host    replication     all             ::1/128                 trust
EOF

    chown -R ${MEM0_UID}:${MEM0_GID} "$PGDATA"
    echo "[TIMER] PostgreSQL init complete"
  else
    echo "PostgreSQL data exists at $PGDATA"
  fi

  # FUSE repair, every boot (mirrors affine bootstrap.sh):
  # - HF bucket drops empty dirs and injects .git/ into tracked dirs.
  # - Postgres scans these dirs as object stores; any stray entry → PANIC.
  #   They must exist and be COMPLETELY EMPTY (no .gitkeep).
  for dir in pg_replslot pg_twophase pg_dynshmem pg_notify; do
    mkdir -p "$PGDATA/$dir"
    find "$PGDATA/$dir" -mindepth 1 -delete 2>/dev/null || true
  done
  # These tolerate a .gitkeep marker (Postgres ignores unknown files here)
  for dir in pg_commit_ts pg_serial pg_stat pg_stat_tmp \
             pg_logical/snapshots pg_logical/mappings; do
    mkdir -p "$PGDATA/$dir"
    touch "$PGDATA/$dir/.gitkeep"
    find "$PGDATA/$dir" -mindepth 1 -name '.git' -type d -exec rm -rf {} + 2>/dev/null || true
  done
  # pg_tblspc must be a symlink to tmpfs — FUSE drops real empty dirs
  mkdir -p /tmp/pg_tblspc
  rm -rf "$PGDATA/pg_tblspc"
  ln -sfn /tmp/pg_tblspc "$PGDATA/pg_tblspc"

  # Defensive re-assert: only UID 1000 persists on the FUSE idmapped mount.
  # Scoped chown on dirs only — recursive chown over FUSE has killed the mount.
  chown ${MEM0_UID}:${MEM0_GID} "$PGDATA" "$PGDATA"/*/ 2>/dev/null || true
  chmod 700 "$PGDATA"

  # Fix socket dir so mem0 can create the unix socket lock file
  mkdir -p /var/run/postgresql
  chown ${MEM0_UID}:${MEM0_GID} /var/run/postgresql
  export FRESH_INIT
fi

# ── Stage 2: First-boot DB setup (fresh cluster only) ──────────────────────
if [ "$(id -u)" = "0" ] && [ "${FRESH_INIT:-0}" = "1" ]; then
  echo ""
  echo "=== 2. First-boot database setup ==="
  
  mkdir -p /var/log/postgresql
  chown -R ${MEM0_UID}:${MEM0_GID} /var/log/postgresql
  
  # FUSE repair before first pg_ctl start (same logic as supervisord wrapper)
  for dir in pg_replslot pg_twophase pg_dynshmem pg_notify; do
    mkdir -p "$PGDATA/$dir"
    find "$PGDATA/$dir" -mindepth 1 -delete 2>/dev/null || true
  done
  for dir in pg_commit_ts pg_serial pg_stat pg_stat_tmp pg_logical/snapshots pg_logical/mappings; do
    mkdir -p "$PGDATA/$dir"
    touch "$PGDATA/$dir/.gitkeep"
    find "$PGDATA/$dir" -mindepth 1 -name .git -type d -exec rm -rf {} + 2>/dev/null || true
  done
  mkdir -p /tmp/pg_tblspc
  rm -rf "$PGDATA/pg_tblspc"
  ln -sfn /tmp/pg_tblspc "$PGDATA/pg_tblspc"

  # Start postgres temporarily to run init-db + migrations, then stop
  echo "Running database initialization script..."
  PGLOG="/tmp/pg-init.log"
  touch "$PGLOG" && chown ${MEM0_UID}:${MEM0_GID} "$PGLOG"
  set +e
  gosu mem0 /usr/lib/postgresql/15/bin/pg_ctl -D "$PGDATA" -l "$PGLOG" -w start
  PG_START_RC=$?
  set -e
  echo "[pg_ctl start exit: $PG_START_RC]"
  echo "=== pg startup log ==="
  cat "$PGLOG"
  echo "=== end pg log ==="
  if [ "$PG_START_RC" -ne 0 ]; then
    echo "[ERROR] PostgreSQL failed to start; aborting"
    exit 1
  fi
  gosu mem0 /init-db.sh

  # Run app-DB migrations (users/settings/etc. in mem0_app) BEFORE first boot.
  # Fresh volumes have no tables; without this, uvicorn crashes at import on
  # any DB access and supervisord puts mem0-api into FATAL.
  echo "=== Running database migrations ==="
  gosu mem0 sh -c 'cd /opt/mem0/api && python -m alembic upgrade head' \
    || echo "[WARN] alembic failed (non-fatal; API will retry tables lazily)"

  gosu mem0 /usr/lib/postgresql/15/bin/pg_ctl -D "$PGDATA" -w stop
fi

echo ""
echo "=== 3. Starting services via supervisord ==="

# Stream program logs to container stdout so tracebacks are visible in HF logs.
# (supervisord itself cannot reopen /dev/fd/1 -> EACCES on every dispatcher)
touch /var/log/supervisor/mem0-api.log /var/log/supervisor/mem0-api.err.log \
      /var/log/supervisor/mem0-dashboard.log /var/log/supervisor/mem0-dashboard.err.log \
      /var/log/supervisor/postgresql.log /var/log/supervisor/postgresql.err.log \
      /var/log/supervisor/nginx.log /var/log/supervisor/nginx.err.log
chmod 666 /var/log/supervisor/*.log
tail -q -F /var/log/supervisor/*.log &

# supervisord runs as root (user=root in supervisord.conf) so it can recreate
# FUSE-volatile dirs before each postgres spawn via the pg wrapper command.
exec supervisord -c /etc/supervisor/conf.d/supervisord.conf