#!/usr/bin/env bash
#
# Incrementally sync the production Postgres database into the local
# `prop_dev`, pulling only rows whose `updated_at` is newer than the
# corresponding row already present locally. Use this when you already
# have a recent prod-flavoured DB and just want to top it up.
#
# For a first-time bootstrap (empty local DB) or a clean slate, run
# `scripts/restore_prod_to_local.sh` first — that dump+restore is still
# faster than syncing a fresh 50M+ row table one COPY at a time.
#
# Usage:
#
#   scripts/sync_prod_to_local.sh                # sync everything
#   scripts/sync_prod_to_local.sh contacts       # one table
#   scripts/sync_prod_to_local.sh --dry-run      # count rows, no write
#
# Before running: stop `mix phx.server` so Ecto doesn't hold connections.
# Requires direnv (or `source .envrc`) so PROP_PROD_DB_URL is set.

set -euo pipefail

: "${PROP_PROD_DB_URL:?not set — run via direnv or 'source .envrc' first}"

LOCAL_DB="${LOCAL_DB:-prop_dev}"
LOCAL_USER="${LOCAL_USER:-postgres}"
LOCAL_HOST="${LOCAL_HOST:-localhost}"
LOCAL_TMP="${LOCAL_TMP:-/tmp}"

# Prod runs Postgres 18; pin to the Homebrew @18 binary so the client's
# wire protocol stays in step with the server (same reason as the full
# restore script).
PG_BIN="${PG_BIN:-/opt/homebrew/opt/postgresql@18/bin}"
PSQL="${PG_BIN}/psql"

# Tables to skip. `schema_migrations` is owned by `mix ecto.migrate`;
# `oban_jobs` / `oban_peers` / `oban_producers` are runtime state, not
# data we want carried across envs. Add more here if they cause trouble.
SKIP_TABLES="${SKIP_TABLES:-schema_migrations oban_jobs oban_peers oban_producers oban_met_dim_datapoints}"

log() { printf '%s sync_prod_to_local: %s\n' "$(date -u +'%Y-%m-%dT%H:%M:%SZ')" "$*"; }
die() { log "ERROR: $*"; exit 1; }

[[ -x "$PSQL" ]] || die "$PSQL missing — brew install postgresql@18"

# --- Arg parsing -------------------------------------------------------

dry_run=0
only_table=""
for arg in "$@"; do
    case "$arg" in
        --dry-run) dry_run=1 ;;
        --help|-h)
            sed -n '1,/^set -euo/p' "$0" | sed '/^set/d; s/^# \{0,1\}//'
            exit 0
            ;;
        --*)
            die "unknown flag: $arg"
            ;;
        *)
            if [[ -n "$only_table" ]]; then
                die "at most one table name allowed (got '$only_table' and '$arg')"
            fi
            only_table="$arg"
            ;;
    esac
done

# --- psql command arrays ---------------------------------------------

local_psql=("$PSQL" -h "$LOCAL_HOST" -U "$LOCAL_USER" -d "$LOCAL_DB" -v ON_ERROR_STOP=1 --quiet)
prod_psql=("$PSQL" "$PROP_PROD_DB_URL" -v ON_ERROR_STOP=1 --quiet)

# --- Table discovery --------------------------------------------------

discover_sql="
SELECT c.relname
FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE n.nspname = 'public'
  -- 'r' = regular, 'p' = partitioned parent. Keep the parent, drop
  -- the leaf partitions (relispartition=true) — an INSERT into the
  -- parent routes data to the correct child and lets us use the
  -- parent's cross-partition unique index as the ON CONFLICT target.
  AND c.relkind IN ('r', 'p')
  AND NOT c.relispartition
  AND EXISTS (
      SELECT 1 FROM pg_attribute a
      WHERE a.attrelid = c.oid AND a.attname = 'updated_at' AND NOT a.attisdropped
  )
ORDER BY c.relname
"

mapfile -t tables < <("${local_psql[@]}" -tAc "$discover_sql")

# Filter out skip list.
declare -A skip_set=()
for t in $SKIP_TABLES; do skip_set["$t"]=1; done

filtered=()
for t in "${tables[@]}"; do
    [[ -n "${skip_set[$t]:-}" ]] && continue
    if [[ -n "$only_table" && "$t" != "$only_table" ]]; then continue; fi
    filtered+=("$t")
done

tables=("${filtered[@]}")

if [[ ${#tables[@]} -eq 0 ]]; then
    die "no candidate tables (only_table='$only_table', skip='$SKIP_TABLES')"
fi

log "syncing ${#tables[@]} tables (dry_run=$dry_run)"

# --- Per-table sync ---------------------------------------------------

work=$(mktemp -d "${LOCAL_TMP}/sync_XXXXXX")
trap 'rm -rf "$work"' EXIT

total_rows=0

for tbl in "${tables[@]}"; do
    # Intersect local + prod column lists so schema drift (a new column
    # added on prod but not yet migrated locally, or vice-versa) turns
    # into "transfer the columns that exist on both sides" instead of
    # a hard COPY/INSERT failure.
    local_cols_raw=$(
        "${local_psql[@]}" -tAc "
            SELECT column_name
            FROM information_schema.columns
            WHERE table_schema='public' AND table_name='$tbl'
            ORDER BY column_name
        "
    )
    prod_cols_raw=$(
        "${prod_psql[@]}" -tAc "
            SELECT column_name
            FROM information_schema.columns
            WHERE table_schema='public' AND table_name='$tbl'
            ORDER BY column_name
        "
    )

    common_cols=$(comm -12 <(echo "$local_cols_raw") <(echo "$prod_cols_raw"))

    if [[ -z "$common_cols" ]]; then
        log "  $tbl: no shared columns between local and prod, skipping"
        continue
    fi

    cols=$(echo "$common_cols" | awk 'NF { printf "%s\"%s\"", sep, $0; sep="," }')

    # Resolve the ON CONFLICT target. Prefer the primary key; if the
    # table doesn't have one (partitioned tables often don't), fall
    # back to the shortest non-partial unique index. If neither
    # exists, skip — we can't do a safe upsert without a uniqueness
    # guarantee.
    pk_cols_raw=$(
        "${local_psql[@]}" -tAc "
            WITH candidates AS (
                -- Primary key, ranked first.
                SELECT 0 AS priority, i.indexrelid,
                       string_agg(a.attname, E'\n' ORDER BY array_position(i.indkey::int[], a.attnum)) AS col_list
                FROM pg_index i
                JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)
                WHERE i.indrelid = '$tbl'::regclass AND i.indisprimary
                GROUP BY i.indexrelid
                UNION ALL
                -- Any non-partial unique index, ranked by fewest cols.
                SELECT 1 + array_length(i.indkey::int[], 1) AS priority, i.indexrelid,
                       string_agg(a.attname, E'\n' ORDER BY array_position(i.indkey::int[], a.attnum)) AS col_list
                FROM pg_index i
                JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)
                WHERE i.indrelid = '$tbl'::regclass
                  AND i.indisunique AND NOT i.indisprimary
                  AND i.indpred IS NULL
                GROUP BY i.indexrelid, i.indkey
            )
            SELECT col_list FROM candidates ORDER BY priority LIMIT 1
        "
    )

    if [[ -z "$pk_cols_raw" ]]; then
        log "  $tbl: no primary key or unique index, skipping"
        continue
    fi

    missing_pk=0
    while IFS= read -r pk; do
        [[ -z "$pk" ]] && continue
        if ! echo "$common_cols" | grep -qx -- "$pk"; then
            missing_pk=1
        fi
    done <<<"$pk_cols_raw"

    if [[ "$missing_pk" -eq 1 ]]; then
        log "  $tbl: primary-key column missing on prod, skipping"
        continue
    fi

    pk_cols=$(echo "$pk_cols_raw" | awk 'NF { printf "%s\"%s\"", sep, $0; sep="," }')

    # ON CONFLICT DO UPDATE SET list: every common non-PK column.
    declare -A pk_set=()
    while IFS= read -r pk; do
        [[ -n "$pk" ]] && pk_set["$pk"]=1
    done <<<"$pk_cols_raw"

    update_set=""
    while IFS= read -r col; do
        [[ -z "$col" ]] && continue
        [[ -n "${pk_set[$col]:-}" ]] && continue
        if [[ -n "$update_set" ]]; then
            update_set+=","
        fi
        update_set+="\"${col}\"=EXCLUDED.\"${col}\""
    done <<<"$common_cols"
    unset pk_set

    if [[ -z "$update_set" ]]; then
        log "  $tbl: only PK columns shared, skipping (nothing to update)"
        continue
    fi

    local_max=$(
        "${local_psql[@]}" -tAc "
            SELECT COALESCE(MAX(updated_at), '1970-01-01 00:00:00+00')::text FROM $tbl
        "
    )

    # Fetch changed rows from prod into a local temp file first (so
    # upsert failures don't leave a half-streamed COPY hanging).
    chunk="${work}/${tbl}.copy"
    "${prod_psql[@]}" -c "\\COPY (SELECT $cols FROM $tbl WHERE updated_at > '$local_max') TO '$chunk'"

    rows=$(wc -l < "$chunk" | tr -d ' ')

    if [[ "$rows" -eq 0 ]]; then
        log "  $tbl: up to date (max=$local_max)"
        continue
    fi

    log "  $tbl: $rows new/changed rows since $local_max"

    if [[ "$dry_run" -eq 1 ]]; then
        continue
    fi

    # Load into a temp table with the same shape, then upsert. Wrapping
    # in SET session_replication_role = replica defers FK checks for
    # the session so cross-table ordering doesn't block us — the
    # script is dev-only and prod's RI is assumed correct.
    "${local_psql[@]}" <<SQL
BEGIN;
SET LOCAL session_replication_role = replica;
CREATE TEMP TABLE "_sync_${tbl}" (LIKE "$tbl" INCLUDING DEFAULTS) ON COMMIT DROP;
\\COPY "_sync_${tbl}" ($cols) FROM '$chunk';
INSERT INTO "$tbl" ($cols)
SELECT $cols FROM "_sync_${tbl}"
ON CONFLICT ($pk_cols) DO UPDATE SET $update_set;
COMMIT;
SQL

    total_rows=$((total_rows + rows))
done

log "done — synced $total_rows rows across ${#tables[@]} tables"
