#!/usr/bin/env bash
set -euo pipefail
trap 'echo "Interrupted"; exit 130' INT TERM

BASE_DIR="${BASE_DIR:-/vol1/1000/application/ClickHouse}"
ENV_FILE="${ENV_FILE:-$BASE_DIR/.env}"
PG_PASSWORD_FILE="${PG_PASSWORD_FILE:-$BASE_DIR/pg/geo.password}"
STATE_DIR="${STATE_DIR:-$BASE_DIR/state}"
LOG_DIR="${LOG_DIR:-$BASE_DIR/logs}"

CLICKHOUSE_CONTAINER="${CLICKHOUSE_CONTAINER:-geo-rankly-clickhouse}"
ARCHIVE_DB="${ARCHIVE_DB:-geo_monitoring_archive}"

PG_HOST="${PG_HOST:-10.77.0.1}"
PG_PORT="${PG_PORT:-15432}"
PG_DB="${PG_DB:-geo_monitoring}"
PG_USER="${PG_USER:-geo}"
TIMEZONE="${TIMEZONE:-Asia/Shanghai}"
FULL_IMPORT_CHUNK_DAYS="${FULL_IMPORT_CHUNK_DAYS:-31}"

usage() {
  cat <<'USAGE'
Usage:
  pg_to_clickhouse_archive.sh setup
  pg_to_clickhouse_archive.sh initial-full [--force]
  pg_to_clickhouse_archive.sh previous-day
  pg_to_clickhouse_archive.sh date YYYY-MM-DD
  pg_to_clickhouse_archive.sh status

Modes:
  setup         Create archive database and tables only.
  initial-full  Import all current PostgreSQL rows once. Refuses to run again unless --force is provided.
  previous-day Import yesterday in Asia/Shanghai.
  date          Import one explicit business day/date partition, e.g. 2026-06-30.
  status        Show row counts in ClickHouse archive tables.
USAGE
}

die() {
  echo "ERROR: $*" >&2
  exit 1
}

log() {
  printf '[%s] %s\n' "$(TZ="$TIMEZONE" date '+%F %T %Z')" "$*"
}

require_root_or_docker_access() {
  if docker inspect "$CLICKHOUSE_CONTAINER" >/dev/null 2>&1; then
    return
  fi
  die "Cannot access Docker. Run with sudo/root, for example: sudo $0 $*"
}

load_config() {
  [[ -f "$ENV_FILE" ]] || die "Missing ClickHouse env file: $ENV_FILE"
  [[ -f "$PG_PASSWORD_FILE" ]] || die "Missing PostgreSQL password file: $PG_PASSWORD_FILE"

  # shellcheck disable=SC1090
  source "$ENV_FILE"
  [[ -n "${CLICKHOUSE_USER:-}" ]] || die "CLICKHOUSE_USER is not set in $ENV_FILE"
  [[ -n "${CLICKHOUSE_PASSWORD:-}" ]] || die "CLICKHOUSE_PASSWORD is not set in $ENV_FILE"

  PG_PASSWORD="$(tr -d '\r\n' < "$PG_PASSWORD_FILE")"
  [[ -n "$PG_PASSWORD" ]] || die "PostgreSQL password file is empty: $PG_PASSWORD_FILE"

  mkdir -p "$STATE_DIR" "$LOG_DIR"
}

sql_quote() {
  sed "s/'/''/g" <<<"$1"
}

ch() {
  docker exec -i "$CLICKHOUSE_CONTAINER" clickhouse-client \
    --user "$CLICKHOUSE_USER" \
    --password "$CLICKHOUSE_PASSWORD" \
    "$@"
}

ch_query() {
  local query="$1"
  ch --query "$query"
}

ch_multi() {
  ch --multiquery
}

pg_source() {
  local table="$1"
  local pg_password_sql
  pg_password_sql="$(sql_quote "$PG_PASSWORD")"
  printf "postgresql('%s:%s', '%s', '%s', '%s', '%s')" \
    "$PG_HOST" "$PG_PORT" "$PG_DB" "$table" "$PG_USER" "$pg_password_sql"
}

psql_query() {
  local query="$1"
  PGPASSWORD="$PG_PASSWORD" PGCONNECT_TIMEOUT=15 PGTZ="$TIMEZONE" psql \
    -h "$PG_HOST" \
    -p "$PG_PORT" \
    -U "$PG_USER" \
    -d "$PG_DB" \
    -v ON_ERROR_STOP=1 \
    -At \
    -F $'\t' \
    -c "${query}"
}

psql_copy_to_stage() {
  local table="$1"
  local stage="$2"
  local pg_where="$3"
  local select_list
  select_list="$(pg_select_list "$table")"

  PGPASSWORD="$PG_PASSWORD" PGCONNECT_TIMEOUT=15 PGTZ="$TIMEZONE" psql \
    -h "$PG_HOST" \
    -p "$PG_PORT" \
    -U "$PG_USER" \
    -d "$PG_DB" \
    -v ON_ERROR_STOP=1 \
    -q \
    -c "\\copy (SELECT ${select_list} FROM public.${table} WHERE ${pg_where}) TO STDOUT WITH (FORMAT CSV, HEADER true)" \
  | ch --query "INSERT INTO ${ARCHIVE_DB}.${stage} FORMAT CSVWithNames"
}

pg_select_list() {
  case "$1" in
    question_monitor_parse_results)
      cat <<'SQL'
id,
tenant_id,
run_id,
brand_id,
question_id,
question_hash,
ai_platform_id,
business_date,
CASE WHEN brand_mentioned THEN 1 WHEN brand_mentioned IS FALSE THEN 0 ELSE NULL END AS brand_mentioned,
brand_mention_position,
CASE WHEN first_recommended THEN 1 WHEN first_recommended IS FALSE THEN 0 ELSE NULL END AS first_recommended,
sentiment_label,
matched_brand_terms,
parse_version,
parse_status,
created_at,
updated_at
SQL
      ;;
    monitoring_platform_access_snapshots)
      cat <<'SQL'
id,
tenant_id,
client_id,
ai_platform_id,
business_date,
access_status,
CASE WHEN connected THEN 1 ELSE 0 END AS connected,
CASE WHEN accessible THEN 1 ELSE 0 END AS accessible,
detected_at,
reason_code,
reason_message,
created_at,
updated_at
SQL
      ;;
    monitoring_question_config_snapshots)
      cat <<'SQL'
id,
tenant_id,
brand_id,
question_id,
keyword_id,
question_hash,
question_text_snapshot,
CASE WHEN monitor_enabled THEN 1 ELSE 0 END AS monitor_enabled,
superseded_at,
deleted_at,
projected_at,
created_at,
updated_at
SQL
      ;;
    *)
      echo "*"
      ;;
  esac
}

declare -a TABLES=(
  monitoring_collect_tasks
  question_monitor_runs
  question_monitor_parse_results
  monitoring_citation_facts
  monitoring_brand_daily
  monitoring_brand_platform_daily
  monitoring_platform_access_snapshots
  monitoring_collect_requests
  monitoring_collect_dispatch_outbox
  monitoring_question_config_snapshots
  monitoring_article_url_aliases
  monitoring_user_marked_articles
)

date_expr() {
  case "$1" in
    monitoring_collect_tasks|question_monitor_runs|question_monitor_parse_results|monitoring_citation_facts|monitoring_brand_daily|monitoring_brand_platform_daily|monitoring_platform_access_snapshots)
      echo "business_date"
      ;;
    monitoring_user_marked_articles)
      echo "toDate(marked_at, '$TIMEZONE')"
      ;;
    monitoring_collect_dispatch_outbox|monitoring_collect_requests|monitoring_question_config_snapshots|monitoring_article_url_aliases)
      echo "toDate(created_at, '$TIMEZONE')"
      ;;
    *)
      die "No date expression configured for table: $1"
      ;;
  esac
}

order_by_expr() {
  case "$1" in
    monitoring_collect_dispatch_outbox)
      echo "(toDate(created_at, '$TIMEZONE'), workspace_id, id)"
      ;;
    monitoring_user_marked_articles)
      echo "(toDate(marked_at, '$TIMEZONE'), tenant_id, id)"
      ;;
    monitoring_collect_requests|monitoring_question_config_snapshots|monitoring_article_url_aliases)
      echo "(toDate(created_at, '$TIMEZONE'), tenant_id, id)"
      ;;
    *)
      echo "(business_date, tenant_id, id)"
      ;;
  esac
}

where_for_date() {
  local table="$1"
  local day="$2"
  local next_day
  next_day="$(TZ="$TIMEZONE" date -d "$day +1 day" +%F)"

  case "$table" in
    monitoring_collect_tasks|question_monitor_runs|question_monitor_parse_results|monitoring_citation_facts|monitoring_brand_daily|monitoring_brand_platform_daily|monitoring_platform_access_snapshots)
      echo "business_date = toDate32('$day')"
      ;;
    monitoring_user_marked_articles)
      echo "marked_at >= toDateTime64('$day 00:00:00', 6, '$TIMEZONE') AND marked_at < toDateTime64('$next_day 00:00:00', 6, '$TIMEZONE')"
      ;;
    monitoring_collect_dispatch_outbox|monitoring_collect_requests|monitoring_question_config_snapshots|monitoring_article_url_aliases)
      echo "created_at >= toDateTime64('$day 00:00:00', 6, '$TIMEZONE') AND created_at < toDateTime64('$next_day 00:00:00', 6, '$TIMEZONE')"
      ;;
    *)
      die "No date filter configured for table: $table"
      ;;
  esac
}

where_for_range() {
  local table="$1"
  local start_day="$2"
  local end_day="$3"

  case "$table" in
    monitoring_collect_tasks|question_monitor_runs|question_monitor_parse_results|monitoring_citation_facts|monitoring_brand_daily|monitoring_brand_platform_daily|monitoring_platform_access_snapshots)
      echo "business_date >= toDate32('$start_day') AND business_date < toDate32('$end_day')"
      ;;
    monitoring_user_marked_articles)
      echo "marked_at >= toDateTime64('$start_day 00:00:00', 6, '$TIMEZONE') AND marked_at < toDateTime64('$end_day 00:00:00', 6, '$TIMEZONE')"
      ;;
    monitoring_collect_dispatch_outbox|monitoring_collect_requests|monitoring_question_config_snapshots|monitoring_article_url_aliases)
      echo "created_at >= toDateTime64('$start_day 00:00:00', 6, '$TIMEZONE') AND created_at < toDateTime64('$end_day 00:00:00', 6, '$TIMEZONE')"
      ;;
    *)
      die "No date range filter configured for table: $table"
      ;;
  esac
}

pg_date_expr() {
  case "$1" in
    monitoring_collect_tasks|question_monitor_runs|question_monitor_parse_results|monitoring_citation_facts|monitoring_brand_daily|monitoring_brand_platform_daily|monitoring_platform_access_snapshots)
      echo "business_date"
      ;;
    monitoring_user_marked_articles)
      echo "(marked_at AT TIME ZONE '$TIMEZONE')::date"
      ;;
    monitoring_collect_dispatch_outbox|monitoring_collect_requests|monitoring_question_config_snapshots|monitoring_article_url_aliases)
      echo "(created_at AT TIME ZONE '$TIMEZONE')::date"
      ;;
    *)
      die "No PostgreSQL date expression configured for table: $1"
      ;;
  esac
}

pg_where_for_date() {
  local table="$1"
  local day="$2"
  local next_day
  next_day="$(TZ="$TIMEZONE" date -d "$day +1 day" +%F)"

  case "$table" in
    monitoring_collect_tasks|question_monitor_runs|question_monitor_parse_results|monitoring_citation_facts|monitoring_brand_daily|monitoring_brand_platform_daily|monitoring_platform_access_snapshots)
      echo "business_date = DATE '$day'"
      ;;
    monitoring_user_marked_articles)
      echo "marked_at >= (TIMESTAMP '$day 00:00:00' AT TIME ZONE '$TIMEZONE') AND marked_at < (TIMESTAMP '$next_day 00:00:00' AT TIME ZONE '$TIMEZONE')"
      ;;
    monitoring_collect_dispatch_outbox|monitoring_collect_requests|monitoring_question_config_snapshots|monitoring_article_url_aliases)
      echo "created_at >= (TIMESTAMP '$day 00:00:00' AT TIME ZONE '$TIMEZONE') AND created_at < (TIMESTAMP '$next_day 00:00:00' AT TIME ZONE '$TIMEZONE')"
      ;;
    *)
      die "No PostgreSQL date filter configured for table: $table"
      ;;
  esac
}

pg_where_for_range() {
  local table="$1"
  local start_day="$2"
  local end_day="$3"

  case "$table" in
    monitoring_collect_tasks|question_monitor_runs|question_monitor_parse_results|monitoring_citation_facts|monitoring_brand_daily|monitoring_brand_platform_daily|monitoring_platform_access_snapshots)
      echo "business_date >= DATE '$start_day' AND business_date < DATE '$end_day'"
      ;;
    monitoring_user_marked_articles)
      echo "marked_at >= (TIMESTAMP '$start_day 00:00:00' AT TIME ZONE '$TIMEZONE') AND marked_at < (TIMESTAMP '$end_day 00:00:00' AT TIME ZONE '$TIMEZONE')"
      ;;
    monitoring_collect_dispatch_outbox|monitoring_collect_requests|monitoring_question_config_snapshots|monitoring_article_url_aliases)
      echo "created_at >= (TIMESTAMP '$start_day 00:00:00' AT TIME ZONE '$TIMEZONE') AND created_at < (TIMESTAMP '$end_day 00:00:00' AT TIME ZONE '$TIMEZONE')"
      ;;
    *)
      die "No PostgreSQL date range filter configured for table: $table"
      ;;
  esac
}

validate_day() {
  local day="$1"
  [[ "$day" =~ ^[0-9]{4}-[0-9]{2}-[0-9]{2}$ ]] || die "Invalid date: $day"
  TZ="$TIMEZONE" date -d "$day" +%F >/dev/null || die "Invalid date: $day"
}

create_table_from_source() {
  local table="$1"
  local partition_expr order_expr source
  partition_expr="$(date_expr "$table")"
  order_expr="$(order_by_expr "$table")"
  source="$(pg_source "$table")"

  ch_multi <<SQL
CREATE DATABASE IF NOT EXISTS ${ARCHIVE_DB};
CREATE TABLE IF NOT EXISTS ${ARCHIVE_DB}.${table}
ENGINE = MergeTree
PARTITION BY ${partition_expr}
ORDER BY ${order_expr}
AS
SELECT *
FROM ${source}
WHERE 0;
SQL
}

create_stage_table() {
  local table="$1"
  local stage="$2"
  local select_where="$3"

  ch_multi <<SQL
DROP TABLE IF EXISTS ${ARCHIVE_DB}.${stage};
CREATE TABLE ${ARCHIVE_DB}.${stage} AS ${ARCHIVE_DB}.${table};
SQL

  psql_copy_to_stage "$table" "$stage" "$select_where"
}

setup_tables() {
  log "Creating archive database and tables in ${ARCHIVE_DB}"
  for table in "${TABLES[@]}"; do
    log "setup ${table}"
    create_table_from_source "$table"
  done
  log "setup complete"
}

cleanup_work_tables() {
  local table
  while IFS= read -r table; do
    [[ -n "$table" ]] || continue
    log "cleanup work table ${table}"
    ch_query "DROP TABLE IF EXISTS ${ARCHIVE_DB}.${table}"
  done < <(ch_query "SHOW TABLES FROM ${ARCHIVE_DB} LIKE '__stage_%' FORMAT TabSeparatedRaw")

  while IFS= read -r table; do
    [[ -n "$table" ]] || continue
    log "cleanup old table ${table}"
    ch_query "DROP TABLE IF EXISTS ${ARCHIVE_DB}.${table}"
  done < <(ch_query "SHOW TABLES FROM ${ARCHIVE_DB} LIKE '__old_%' FORMAT TabSeparatedRaw")
}

stage_count() {
  local table="$1"
  ch_query "SELECT count() FROM ${ARCHIVE_DB}.${table} FORMAT TabSeparatedRaw"
}

stage_count_where() {
  local table="$1"
  local where_clause="$2"
  ch_query "SELECT count() FROM ${ARCHIVE_DB}.${table} WHERE ${where_clause} FORMAT TabSeparatedRaw"
}

source_count() {
  local table="$1"
  local where_clause="${2:-1}"
  ch_query "SELECT count() FROM $(pg_source "$table") WHERE ${where_clause} FORMAT TabSeparatedRaw"
}

import_full_table() {
  local table="$1"
  local bounds min_day max_day chunk_start chunk_end after_max where_clause stage total_days

  bounds="$(source_date_bounds "$table")"
  min_day="${bounds%%$'\t'*}"
  max_day="${bounds##*$'\t'}"

  if [[ -z "$min_day" || -z "$max_day" ]]; then
    log "full import ${table}: source has no rows"
    return
  fi

  log "full import ${table}: day range ${min_day}..${max_day}"
  after_max="$(TZ="$TIMEZONE" date -d "$max_day +1 day" +%F)"
  chunk_start="$min_day"
  total_days=0

  while [[ "$chunk_start" < "$after_max" ]]; do
    chunk_end="$(TZ="$TIMEZONE" date -d "$chunk_start +${FULL_IMPORT_CHUNK_DAYS} day" +%F)"
    if [[ "$chunk_end" > "$after_max" ]]; then
      chunk_end="$after_max"
    fi

    where_clause="$(pg_where_for_range "$table" "$chunk_start" "$chunk_end")"
    stage="__stage_full_${chunk_start//-/}_${table}_$(TZ="$TIMEZONE" date +%Y%m%d%H%M%S)_$$"

    log "full import staging ${table}: ${chunk_start}..<${chunk_end}"
    create_stage_table "$table" "$stage" "$where_clause"
    replace_partitions_from_stage "$table" "$stage" "$chunk_start" "$chunk_end"

    total_days=$((total_days + ($(TZ="$TIMEZONE" date -d "$chunk_end" +%s) - $(TZ="$TIMEZONE" date -d "$chunk_start" +%s)) / 86400))
    chunk_start="$chunk_end"
  done
  log "full import ${table}: imported ${total_days} day partitions"
}

source_date_bounds() {
  local table="$1"
  local expr
  expr="$(pg_date_expr "$table")"

  psql_query "SELECT COALESCE(min(${expr})::text, ''), COALESCE(max(${expr})::text, '') FROM public.${table};"
}

replace_partitions_from_stage() {
  local table="$1"
  local stage="$2"
  local start_day="$3"
  local end_day="$4"
  local day rows where_clause

  day="$start_day"
  while [[ "$day" < "$end_day" ]]; do
    where_clause="$(where_for_date "$table" "$day")"
    rows="$(stage_count_where "$stage" "$where_clause")"

    if [[ "$rows" == "0" ]]; then
      log "full import ${table} day=${day} rows=0; dropping target partition after successful staging"
      ch_query "ALTER TABLE ${ARCHIVE_DB}.${table} DROP PARTITION '${day}'"
    else
      log "full import replacing partition ${day} in ${table}, rows=${rows}"
      ch_query "ALTER TABLE ${ARCHIVE_DB}.${table} REPLACE PARTITION '${day}' FROM ${ARCHIVE_DB}.${stage}"
    fi

    day="$(TZ="$TIMEZONE" date -d "$day +1 day" +%F)"
  done

  ch_query "DROP TABLE IF EXISTS ${ARCHIVE_DB}.${stage}"
}

initial_full() {
  local force="${1:-}"
  local marker="$STATE_DIR/initial_full_import.completed"

  if [[ -f "$marker" && "$force" != "--force" ]]; then
    die "Initial full import already completed at $(cat "$marker"). Refusing to replace permanent archive. Use --force only if you are intentionally rebuilding from PostgreSQL's current retention window."
  fi

  setup_tables
  cleanup_work_tables
  log "Starting initial full import"
  for table in "${TABLES[@]}"; do
    import_full_table "$table"
  done

  TZ="$TIMEZONE" date '+%F %T %Z' > "$marker"
  log "Initial full import complete. Marker: $marker"
}

import_date_table() {
  local table="$1"
  local day="$2"
  local where_clause stage rows partition

  where_clause="$(pg_where_for_date "$table" "$day")"
  stage="__stage_${day//-/}_${table}_$(TZ="$TIMEZONE" date +%Y%m%d%H%M%S)_$$"
  partition="$day"

  log "date import staging ${table} day=${day}"
  create_stage_table "$table" "$stage" "$where_clause"
  rows="$(stage_count "$stage")"

  if [[ "$rows" == "0" ]]; then
    log "date import ${table} day=${day} rows=0; dropping target partition only after successful empty staging"
    ch_multi <<SQL
ALTER TABLE ${ARCHIVE_DB}.${table} DROP PARTITION '${partition}';
DROP TABLE ${ARCHIVE_DB}.${stage};
SQL
    return
  fi

  log "date import replacing partition ${partition} in ${table}, rows=${rows}"
  ch_multi <<SQL
ALTER TABLE ${ARCHIVE_DB}.${table} REPLACE PARTITION '${partition}' FROM ${ARCHIVE_DB}.${stage};
DROP TABLE ${ARCHIVE_DB}.${stage};
SQL
}

import_date() {
  local day="$1"
  local marker="$STATE_DIR/sync_${day}.completed"

  validate_day "$day"
  setup_tables
  cleanup_work_tables

  log "Starting date import day=${day}"
  for table in "${TABLES[@]}"; do
    import_date_table "$table" "$day"
  done

  TZ="$TIMEZONE" date '+%F %T %Z' > "$marker"
  log "Date import complete day=${day}. Marker: $marker"
}

previous_day() {
  TZ="$TIMEZONE" date -d yesterday +%F
}

status() {
  setup_tables >/dev/null
  log "Archive table row counts"
  for table in "${TABLES[@]}"; do
    printf '%s\t%s\n' "$table" "$(stage_count "$table")"
  done
}

main() {
  local mode="${1:-}"
  [[ -n "$mode" ]] || {
    usage
    exit 2
  }

  require_root_or_docker_access "$@"
  load_config

  case "$mode" in
    setup)
      setup_tables
      ;;
    initial-full)
      initial_full "${2:-}"
      ;;
    previous-day)
      import_date "$(previous_day)"
      ;;
    date)
      [[ -n "${2:-}" ]] || die "date mode requires YYYY-MM-DD"
      import_date "$2"
      ;;
    status)
      status
      ;;
    -h|--help|help)
      usage
      ;;
    *)
      usage
      exit 2
      ;;
  esac
}

main "$@"