#!/usr/bin/env bash
set -euo pipefail
trap 'echo "⚠️ An error occurred. Consider running rollback or checking backups."' ERR

COMPOSE=./docker-compose.yaml
SERVICE=postgres
DATA_DIR=./database
PG_VERSION_FILE="$DATA_DIR/PG_VERSION"

echo "🧪 Validating docker-compose config..."
docker compose -f "$COMPOSE" config > /dev/null || {
  echo "❌ docker-compose config failed. Restore aborted."
  exit 1
}

if [ ! -d "$DATA_DIR" ]; then
  echo "❌ Expected data directory '${DATA_DIR}' does not exist. Aborting."
  exit 1
fi

# echo "🔍 Checking if Postgres service is already running..."
# if ! docker compose ps --services --filter "status=running" | grep -q "^${SERVICE}$"; then
#   echo "⚠️  '${SERVICE}' service is not running. Skipping auto-upgrade step."
#   echo "🔄 Attempting to start '${SERVICE}' service to detect version..."
#   docker compose up -d $SERVICE

#   echo "⏳ Waiting for PostgreSQL to become ready..."
#   for i in $(seq 1 60); do
#     if docker compose exec -T $SERVICE pg_isready -U postgres > /dev/null 2>&1; then
#       break
#     fi
#     echo "⏳ Still waiting... (${i}s)"
#     sleep 1
#   done

#   if ! docker compose exec -T $SERVICE pg_isready -U postgres > /dev/null 2>&1; then
#     echo "❌ PostgreSQL did not become ready in time. Aborting."
#     echo "💡 Postgres is not running. Revert to the old version on you docker-compose.yaml file and start start the service!"
#     echo "1. Run: docker compose up -d --force-recreate $SERVICE"
#     echo "2. Run: docker compose --profile postgres-rollback run --rm postgres-auto-rollback"
#     exit 1
#   fi
# fi

# echo "⏳ Waiting for PostgreSQL to become ready before dumping SQL..."
# for i in $(seq 1 120); do
#   if docker compose exec -T $SERVICE pg_isready -U postgres > /dev/null 2>&1; then
#     break
#   fi
#   echo "⏳ Still waiting... (${i}s)"
#   sleep 1
# done

# if ! docker compose exec -T $SERVICE pg_isready -U postgres > /dev/null 2>&1; then
#   echo "❌ PostgreSQL did not become ready in time. Aborting."
#   exit 1
# fi

echo "📡 Detecting running PostgreSQL version..."
OLD_VERSION=$(cat "$PG_VERSION_FILE")
echo "🔍 Detected running PostgreSQL version: $OLD_VERSION"
OLD_MAJOR=$(echo "$OLD_VERSION" | cut -d. -f1)
echo "🔍 Detected running PostgreSQL major version: $OLD_MAJOR"
OLD_IMG="${OLD_VERSION}-alpine"


echo "🆕 Detecting target version from docker-compose.yaml..."
NEW_IMG=$(docker compose -f $COMPOSE config | grep "image:" | grep "$SERVICE" | awk '{print $2}')

# Ensure NEW_IMG was detected
if [[ -z "$NEW_IMG" ]]; then
  echo "❌ Failed to detect target Postgres image from $COMPOSE. Aborting."
  exit 1
fi

NEW_VERSION=$(echo "$NEW_IMG" | sed -E 's/^postgres://; s/-alpine.*$//')
NEW_MAJOR=$(echo "$NEW_VERSION" | cut -d. -f1)

echo "🔁 From $OLD_VERSION (major $OLD_MAJOR) → $NEW_VERSION (major $NEW_MAJOR)"

if [[ "$NEW_VERSION" == *beta* ]] || [[ "$NEW_VERSION" == *rc* ]] || [[ "$NEW_VERSION" == *bookworm* ]]; then
  echo "❌ Target version $NEW_VERSION appears to be a pre-release (beta/rc/bookworm). Skipping upgrade."
  echo "💡 Please upgrade to a stable version of Postgres."
  exit 1
fi

# Early exit if no upgrade needed
if [ "$OLD_MAJOR" -eq "$NEW_MAJOR" ]; then
  echo "✅ Already running target major version. Skipping upgrade."
  exit 0
fi

# Paths
BACKUP_DIR=${DATA_DIR}_backup_${OLD_IMG}_$(date +%Y%m%d_%H%M%S)
OLD_DATA_DIR=./database_old
UPGRADE_DIR=./database_tmp_upgrade

# 1. Stop services
echo "🛑 Stopping services..."
docker compose -f $COMPOSE down

# 2. Backup database directory
echo "🔐 Creating backup at ${BACKUP_DIR}..."
cp -a "$DATA_DIR" "$BACKUP_DIR"


echo "📦 Dumping full SQL backup using temporary PostgreSQL container..."
DUMP_FILE="backup_dump_${OLD_IMG}_$(date +%Y%m%d_%H%M%S).sql"
TMP_CONTAINER_NAME="pg-dump-${OLD_MAJOR}"

# Run temporary postgres container with existing data dir
docker run -d --rm \
  --name "$TMP_CONTAINER_NAME" \
  -v "$DATA_DIR:/var/lib/postgresql/data" \
  -e POSTGRES_USER=postgres \
  postgres:${OLD_IMG}

echo "⏳ Waiting for pg_dump container to become ready..."
for i in $(seq 1 30); do
  if docker exec "$TMP_CONTAINER_NAME" pg_isready -U postgres > /dev/null 2>&1; then
    break
  fi
  echo "⏳ Still waiting... (${i}s)"
  sleep 1
done

if ! docker exec "$TMP_CONTAINER_NAME" pg_isready -U postgres > /dev/null 2>&1; then
  echo "❌ Temporary container for SQL dump did not become ready. Aborting."
  docker rm -f "$TMP_CONTAINER_NAME" > /dev/null 2>&1 || true
  exit 1
fi

docker exec "$TMP_CONTAINER_NAME" pg_dumpall -U postgres > "$DUMP_FILE"

echo "🧹 Cleaning up older SQL dump files..."
ALL_DUMPS=( $(ls -t backup_dump_*.sql 2>/dev/null || true) )

if [ "${#ALL_DUMPS[@]}" -gt 1 ]; then
  LATEST_DUMP="${ALL_DUMPS[0]}"
  TO_DELETE=( "${ALL_DUMPS[@]:1}" )

  for dump in "${TO_DELETE[@]}"; do
    echo "🗑️ Removing old dump: $dump"
    rm -f "$dump"
  done

  echo "✅ Only latest dump '${LATEST_DUMP}' preserved."
else
  echo "ℹ️ Only one dump file found. No cleanup needed."
fi

docker rm -f "$TMP_CONTAINER_NAME" > /dev/null 2>&1 || true

# 3. Create upgrade target folder
echo "📁 Creating upgrade workspace ${UPGRADE_DIR}..."
mkdir -p "$UPGRADE_DIR"

# 4. Perform pg_upgrade
echo "🔧 Running pg_upgrade via tianon image..."
docker run --rm \
  -v "${BACKUP_DIR}:/var/lib/postgresql/${OLD_MAJOR}/data" \
  -v "${UPGRADE_DIR}:/var/lib/postgresql/${NEW_MAJOR}/data" \
  tianon/postgres-upgrade:${OLD_MAJOR}-to-${NEW_MAJOR} --copy

# 5. Promote new data
echo "🔁 Swapping data directories..."
rm -rf "$DATA_DIR"
mv "$UPGRADE_DIR" "$DATA_DIR"

# 6. Restore pg_hba.conf before startup
echo "🔄 Restoring pg_hba.conf if it existed..."
cp "${BACKUP_DIR}/pg_hba.conf" "${DATA_DIR}/pg_hba.conf" || echo "✅ No custom pg_hba.conf to restore."

# 7. Update image in docker-compose.yaml
echo "📝 Updating docker-compose to use image ${NEW_IMG}..."
sed -i.bak -E "s#postgres:[^ ]*${OLD_MAJOR}[^ ]*#postgres:${NEW_IMG}#" "$COMPOSE"

# 8. Start container
echo "🚀 Starting upgraded container..."
docker compose -f $COMPOSE up -d $SERVICE

# 9. Wait until DB is accepting connections
echo "⏳ Waiting for PostgreSQL to become ready..."
until docker compose exec -T $SERVICE pg_isready -U postgres; do
  sleep 1
done

# 10. Collation and Reindexing
echo "🔧 Reindexing and refreshing collation versions..."
docker compose exec $SERVICE bash -c '
  set -e
  DBS=$(psql -U postgres -tAc "SELECT datname FROM pg_database WHERE datallowconn")
  for db in $DBS; do
    echo "➡️ Reindexing $db..."
    psql -U postgres -d "$db" -c "REINDEX DATABASE \"$db\";" || true
    psql -U postgres -d "$db" -c "REINDEX SYSTEM \"$db\";" || true

    echo "➡️ Refreshing collation version for $db..."
    if ! psql -U postgres -d "$db" -c "ALTER DATABASE \"$db\" REFRESH COLLATION VERSION;" 2>/dev/null; then
      echo "⚠️  Collation refresh failed. Forcing reset..."
      psql -U postgres -d postgres -c "UPDATE pg_database SET datcollversion = NULL WHERE datname = '\''$db'\'';" || true
      psql -U postgres -d "$db" -c "ALTER DATABASE \"$db\" REFRESH COLLATION VERSION;" || \
        echo "❌ Still failed for $db. Review manually."
    fi

    echo "➡️ Refreshing system collations in $db..."
    for coll in $(psql -U postgres -d "$db" -tAc "SELECT nspname || '\''.'\'' || quote_ident(collname) FROM pg_collation JOIN pg_namespace ON collnamespace = pg_namespace.oid WHERE collprovider = '\''c'\'';"); do
      echo "  🌀 ALTER COLLATION $coll REFRESH VERSION;"
      psql -U postgres -d "$db" -c "ALTER COLLATION $coll REFRESH VERSION;" || \
        echo "    ⚠️ Skipped $coll due to version mismatch (likely Alpine)."
    done
  done
'

# 11. Suppress collation warnings on musl (Alpine)
if docker compose exec $SERVICE ldd --version 2>&1 | grep -qi 'musl'; then
  echo "🧼 Detected musl libc (Alpine). Resetting all datcollversion values..."
  docker compose exec -T $SERVICE psql -U postgres -d postgres -c \
    "UPDATE pg_database SET datcollversion = NULL WHERE datcollversion IS NOT NULL;"
fi

# 12. Make delete_old_cluster.sh executable
DELETE_SCRIPT="./delete_old_cluster.sh"
if [[ -f "$DELETE_SCRIPT" ]]; then
  chmod +x "$DELETE_SCRIPT"
fi

# 13. Make rollback script executable
ROLLBACK_SCRIPT="./rollback_postgres_upgrade.sh"
if [[ -f "$ROLLBACK_SCRIPT" ]]; then
  chmod +x "$ROLLBACK_SCRIPT"
fi

# 14. Final message
echo "✅ Upgrade complete!"
echo "🎉 Postgres is now running ${NEW_IMG} with data in '${DATA_DIR}'."
echo "🧰 Old version is saved in '${OLD_DATA_DIR}'."
echo "💡 Next steps:"
echo "  - ✅ Run smoke tests"
echo "  - 🧹 If all OK - PLEASE MAKE SURE ON YOUR WEBSITE, YOU HAVE ALL THE DATA YOU NEED AFTER THE UPGRADE, run:"
echo "      rm -rf ./database_backup_* ./database_upgraded_*"
echo "🧹 Cleaning up older backups..."
find . -maxdepth 1 -type d -name "database_backup_*" ! -path "./${BACKUP_DIR##*/}" -exec rm -rf {} +
echo "✅ Only latest backup '${BACKUP_DIR}' preserved."

# Step 15: Restart full application
echo "🔄 Pulling latest images..."
if ! docker compose pull; then
  echo "❌ Failed to pull images. Aborting."
  exit 1
fi

echo "🔄 Starting full application stack..."
if ! docker compose up -d --force-recreate; then
  echo "❌ Failed to start application stack. Aborting."
  exit 1
fi

echo "✅ Deployment completed successfully."