#!/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

# Prisma DATABASE_URL often has @ in the password (09128081392@Montime).
# libpq URI parsing then treats host as "Montime@127.0.0.1". Convert to keyword form.
uri_to_conninfo() {
  local u="${1%%\?*}"
  if [[ "$u" != postgresql://* && "$u" != postgres://* ]]; then
    echo "$u"
    return 0
  fi
  local rest="${u#*://}"
  if [[ "$rest" =~ ^([^:]+):(.+)@([^:/]+):([0-9]+)/([^/?]+) ]]; then
    printf 'host=%s port=%s dbname=%s user=%s password=%s' \
      "${BASH_REMATCH[3]}" "${BASH_REMATCH[4]}" "${BASH_REMATCH[5]}" \
      "${BASH_REMATCH[1]}" "${BASH_REMATCH[2]}"
    return 0
  fi
  echo "$u"
}

resolve_publisher_dsn() {
  if [[ -n "${PUBLISHER_DSN:-}" ]]; then
    echo "$PUBLISHER_DSN"
    return 0
  fi
  if [[ -n "${PGHOST:-}" && -n "${PGUSER:-}" ]]; then
    echo "host=${PGHOST} port=${PGPORT:-5432} dbname=${PGDATABASE:-superadmin} user=${PGUSER} password=${PGPASSWORD:-}"
    return 0
  fi
  if [[ -n "${DATABASE_URL:-}" ]]; then
    echo "$DATABASE_URL"
    return 0
  fi
  echo "Set DATABASE_URL or PUBLISHER_DSN (publisher Postgres)." >&2
  return 1
}

PUBLISHER="$(resolve_publisher_dsn)" || exit 1

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

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

libpq_dsn() {
  uri_to_conninfo "$1"
}

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:-?}"
}

if ! run_psql "$(libpq_dsn "$PUBLISHER")" -v ON_ERROR_STOP=1 -At -c "SELECT 1" >/dev/null 2>&1; then
  echo "Cannot connect to publisher ($(label "$PUBLISHER"))." >&2
  echo "Add PUBLISHER_DSN to deploy/replicas.env — URL-encode @ in passwords as %40." >&2
  exit 1
fi

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
}

grant_replica_writer() {
  local dsn="$1"
  local user=""
  if [[ "$dsn" == *"://"* ]]; then
    user="$(sed -E 's#^[a-z]+://([^:/]+).*#\1#' <<< "$dsn")"
  else
    user="$(sed -nE 's/.*user=([^ ]+).*/\1/p' <<< "$dsn")"
  fi
  [[ -n "$user" ]] || return 0
  echo "  grant write to $user"
  run_psql "$(libpq_dsn "$dsn")" -v ON_ERROR_STOP=1 -q -c "
    GRANT USAGE ON SCHEMA public TO ${user};
    GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON ALL TABLES IN SCHEMA public TO ${user};
    GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA public TO ${user};
  " || true
}

copy_table_from_publisher() {
  local replica_dsn="$1" table="$2"
  local dump
  dump="$(mktemp)"
  pg_dump "$(libpq_dsn "$PUBLISHER")" \
    -t "$table" \
    --data-only \
    --no-owner \
    --no-privileges \
    --column-inserts \
    --disable-triggers \
    >"$dump"
  local bytes
  bytes="$(wc -c <"$dump" | tr -d ' ')"
  echo "  dumped $table (${bytes} bytes)"
  if [[ "${bytes:-0}" -lt 20 ]]; then
    echo "  ! pg_dump produced empty output for $table" >&2
    rm -f "$dump"
    return 1
  fi
  run_psql "$(libpq_dsn "$replica_dsn")" -v ON_ERROR_STOP=1 -f "$dump"
  rm -f "$dump"
}

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 "$PUBLISHER" "$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

  grant_replica_writer "$replica_dsn"

  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 "$PUBLISHER" "$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 "$PUBLISHER" "$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:-$API_DIR/deploy/.replica-data.lock}"
# --repair is operator-initiated; never block on flock (avoids /tmp permission fights with pm2 user).
if [[ "$MODE" == "--repair-if-needed" ]] && command -v flock >/dev/null; then
  mkdir -p "$(dirname "$LOCK_FILE")"
  touch "$LOCK_FILE" 2>/dev/null || true
  chmod 666 "$LOCK_FILE" 2>/dev/null || true
  if exec 9>>"$LOCK_FILE" 2>/dev/null; then
    if ! flock -n 9; then
      echo "replica data sync already running — skip"
      exit 0
    fi
  else
    echo "  ! lock unavailable ($LOCK_FILE) — continuing without flock" >&2
  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