#!/usr/bin/env bash
# Compare + repair row data on logical-replication subscribers vs publisher.
#
# replicas:sync only ships DDL. If a subscriber missed initial copy (permissions,
# stalled subscription, import before replication caught up), rows stay missing
# until repaired. This script makes publisher the source of truth for core tables.
#
# Usage (on publisher, same deploy/replicas.env as sync-replica-schema.sh):
#   bash deploy/sync-replica-data.sh --check
#   bash deploy/sync-replica-data.sh --repair
#   bash deploy/sync-replica-data.sh --repair-if-needed
set -euo pipefail

API_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
CONFIG_FILE="${REPLICAS_ENV_FILE:-$API_DIR/deploy/replicas.env}"
MODE="${1:---repair-if-needed}"

# Tables that must match publisher for public API (order: parents before children).
SYNC_TABLES=(setting_templates tenants settings)
APP_DB_USER="${REPLICA_APP_DB_USER:-montime}"

case "$MODE" in
  --check|--repair|--repair-if-needed) ;;
  *)
    echo "Usage: $0 [--check|--repair|--repair-if-needed]" >&2
    exit 1
    ;;
esac

command -v psql >/dev/null || { echo "psql not found" >&2; exit 1; }
command -v pg_dump >/dev/null || { echo "pg_dump not found" >&2; exit 1; }

if [[ -f "$CONFIG_FILE" ]]; then
  # shellcheck disable=SC1090
  set -a && source "$CONFIG_FILE" && set +a
fi

if [[ -z "${REPLICA_DSNS:-}" ]]; then
  echo "REPLICA_DSNS is empty. Create $CONFIG_FILE from deploy/replicas.env.example" >&2
  exit 1
fi

if [[ -z "${DATABASE_URL:-}" && -f "$API_DIR/.env" ]]; then
  # shellcheck disable=SC1091
  set -a && source "$API_DIR/.env" && set +a
fi

if [[ -z "${DATABASE_URL:-}" ]]; then
  echo "DATABASE_URL is empty (publisher)." >&2
  exit 1
fi

mapfile -t REPLICAS < <(
  printf '%s\n' "$REPLICA_DSNS" | tr ',' '\n' | sed 's/^[[:space:]]*//;s/[[:space:]]*$//' | grep -v '^$'
)

run_psql() {
  psql "$@" </dev/null
}

libpq_dsn() {
  local u="$1"
  echo "${u%%\?*}"
}

label() {
  local dsn="$1"
  if [[ "$dsn" == *"://"* ]]; then
    sed -E 's#^[a-z]+://([^@]*@)?##; s#\?.*$##' <<< "$dsn"
    return
  fi
  local host port db
  host="$(sed -nE 's/.*host=([^ ]+).*/\1/p' <<< "$dsn")"
  port="$(sed -nE 's/.*port=([^ ]+).*/\1/p' <<< "$dsn")"
  db="$(sed -nE 's/.*dbname=([^ ]+).*/\1/p' <<< "$dsn")"
  echo "${host:-?}:${port:-5432}/${db:-?}"
}

table_count() {
  local dsn="$1" table="$2"
  run_psql "$(libpq_dsn "$dsn")" -v ON_ERROR_STOP=1 -At -c \
    "SELECT count(*)::bigint FROM public.\"$table\""
}

subscription_names() {
  local dsn="$1"
  run_psql "$(libpq_dsn "$dsn")" -v ON_ERROR_STOP=1 -At -c \
    "SELECT subname FROM pg_subscription ORDER BY subname"
}

grant_app_user() {
  local dsn="$1"
  run_psql "$(libpq_dsn "$dsn")" -v ON_ERROR_STOP=1 -q -c "
    GRANT USAGE ON SCHEMA public TO ${APP_DB_USER};
    GRANT SELECT ON ALL TABLES IN SCHEMA public TO ${APP_DB_USER};
    GRANT SELECT ON ALL SEQUENCES IN SCHEMA public TO ${APP_DB_USER};
    ALTER DEFAULT PRIVILEGES IN SCHEMA public
      GRANT SELECT ON TABLES TO ${APP_DB_USER};
  " 2>/dev/null || true
}

copy_table_from_publisher() {
  local replica_dsn="$1" table="$2"
  local pub_dsn
  pub_dsn="$(libpq_dsn "$DATABASE_URL")"
  pg_dump "$pub_dsn" \
    -t "public.\"$table\"" \
    --data-only \
    --no-owner \
    --no-privileges \
    --disable-triggers \
    | run_psql "$(libpq_dsn "$replica_dsn")" -v ON_ERROR_STOP=1 -q
}

repair_replica() {
  local replica_dsn="$1"
  local lbl drifted=()
  local table pub replica

  lbl="$(label "$replica_dsn")"
  echo "== repair $lbl"

  mapfile -t subs < <(subscription_names "$replica_dsn")
  if ((${#subs[@]} == 0)); then
    echo "  ! no subscription on replica — skip data repair (create subscription first)" >&2
    return 1
  fi

  for table in "${SYNC_TABLES[@]}"; do
    pub="$(table_count "$DATABASE_URL" "$table")"
    replica="$(table_count "$replica_dsn" "$table")"
    if [[ "$pub" != "$replica" ]]; then
      drifted+=("$table")
      echo "  drift $table publisher=$pub replica=$replica"
    fi
  done

  if ((${#drifted[@]} == 0)); then
    echo "  data already in sync"
    grant_app_user "$replica_dsn"
    return 0
  fi

  # TRUNCATE tenants CASCADE wipes settings — recopy the whole core set.
  drifted=("${SYNC_TABLES[@]}")

  for sub in "${subs[@]}"; do
    echo "  disabling subscription $sub"
    run_psql "$(libpq_dsn "$replica_dsn")" -v ON_ERROR_STOP=1 -q \
      -c "ALTER SUBSCRIPTION \"${sub}\" DISABLE"
  done

  for table in "${drifted[@]}"; do
    echo "  truncating + copying $table"
    run_psql "$(libpq_dsn "$replica_dsn")" -v ON_ERROR_STOP=1 -q \
      -c "TRUNCATE TABLE public.\"$table\" CASCADE"
    copy_table_from_publisher "$replica_dsn" "$table"
    pub="$(table_count "$DATABASE_URL" "$table")"
    replica="$(table_count "$replica_dsn" "$table")"
    if [[ "$pub" != "$replica" ]]; then
      echo "  ! repair failed for $table (publisher=$pub replica=$replica)" >&2
      for sub in "${subs[@]}"; do
        run_psql "$(libpq_dsn "$replica_dsn")" -v ON_ERROR_STOP=1 -q \
          -c "ALTER SUBSCRIPTION \"${sub}\" ENABLE" || true
      done
      return 1
    fi
    echo "  ok $table=$replica"
  done

  grant_app_user "$replica_dsn"

  for sub in "${subs[@]}"; do
    echo "  enabling subscription $sub"
    run_psql "$(libpq_dsn "$replica_dsn")" -v ON_ERROR_STOP=1 -q \
      -c "ALTER SUBSCRIPTION \"${sub}\" ENABLE"
    run_psql "$(libpq_dsn "$replica_dsn")" -v ON_ERROR_STOP=1 -q \
      -c "ALTER SUBSCRIPTION \"${sub}\" REFRESH PUBLICATION" || true
  done

  echo "  repair done"
}

check_replicas() {
  local needs_repair=0
  local dsn lbl table pub replica

  for dsn in "${REPLICAS[@]}"; do
    lbl="$(label "$dsn")"
    echo "== $lbl"
    run_psql "$(libpq_dsn "$dsn")" -v ON_ERROR_STOP=1 -At -c \
      "SELECT subname || '  received=' || COALESCE(received_lsn::text, '-')
              || '  latest=' || COALESCE(latest_end_lsn::text, '-')
        FROM pg_stat_subscription" 2>/dev/null || echo "  (no pg_stat_subscription)"

    for table in "${SYNC_TABLES[@]}"; do
      pub="$(table_count "$DATABASE_URL" "$table")"
      replica="$(table_count "$dsn" "$table")"
      if [[ "$pub" == "$replica" ]]; then
        echo "  $table OK ($pub)"
      else
        echo "  $table DRIFT publisher=$pub replica=$replica"
        needs_repair=1
      fi
    done
  done

  return "$needs_repair"
}

LOCK_FILE="${REPLICA_DATA_LOCK:-/tmp/superadmin-replica-data.lock}"
if [[ "$MODE" != "--check" ]] && command -v flock >/dev/null; then
  exec 9>"$LOCK_FILE"
  if ! flock -n 9; then
    echo "replica data sync already running — skip"
    exit 0
  fi
fi

if [[ "$MODE" == "--check" ]]; then
  check_replicas
  exit $?
fi

if [[ "$MODE" == "--repair-if-needed" ]]; then
  if check_replicas; then
    echo
    echo "All replicas in sync."
    exit 0
  fi
  echo
  echo "Drift detected — repairing..."
fi

failed=0
for dsn in "${REPLICAS[@]}"; do
  repair_replica "$dsn" || failed=1
done

echo
if ((failed == 0)); then
  echo "Replica data sync complete."
  exit 0
fi

echo "Some replicas failed data repair." >&2
exit 1