#!/bin/bash
set -euo pipefail

CREDS_FILE="/usr/local/directadmin/conf/postgresql.conf"
PLUGIN_ROOT="$(cd "$(dirname "$0")/../.." && pwd)"
PLUGIN_SETTINGS_FILE="$PLUGIN_ROOT/plugin-settings.conf"

PGHOST_VAL=""
PGPORT_VAL="5432"
PGUSER_VAL="postgres"
PGPASSWORD_VAL=""
PG_VERSION_NUM=""

ALLOW_REMOTE_HOSTS="1"
RESTART_POSTGRESQL="0"

HBA_MANAGED_BLOCK_BEGIN="# DA_POSTGRESQL_PLUGIN_HBA_BEGIN"
HBA_MANAGED_BLOCK_END="# DA_POSTGRESQL_PLUGIN_HBA_END"
POSTGRESQL_CONF=""
HBA_FILE=""

CSF_ALLOW_FILE="/etc/csf/csf.allow"
CSF_INCLUDE_FILE="/etc/csf.allow.postgres"

PG_CONF_MARKER_BEGIN="# DA_POSTGRESQL_PLUGIN_BEGIN"
PG_CONF_MARKER_END="# DA_POSTGRESQL_PLUGIN_END"

declare -A REMOTE_IP_NOTES=()

parse_args() {
  while [ "$#" -gt 0 ]; do
    case "$1" in
      --restart-postgresql)
        RESTART_POSTGRESQL="1"
        ;;
      *)
        echo "Nieznany argument: $1" >&2
        exit 2
        ;;
    esac
    shift
  done
}

find_psql() {
  if command -v psql >/dev/null 2>&1; then
    command -v psql
    return
  fi

  local candidate
  candidate="$(ls -1d /usr/pgsql-*/bin/psql 2>/dev/null | sort -V | tail -n1 || true)"
  if [ -n "$candidate" ] && [ -x "$candidate" ]; then
    echo "$candidate"
    return
  fi

  echo ""
}

trim() {
  local value="$1"
  value="${value#"${value%%[![:space:]]*}"}"
  value="${value%"${value##*[![:space:]]}"}"
  printf '%s' "$value"
}

strip_inline_comment() {
  local value="$1"
  value="${value%%#*}"
  trim "$value"
}

unquote() {
  local value
  value="$(trim "$1")"

  if [ "${#value}" -ge 2 ] && [ "${value:0:1}" = "'" ] && [ "${value: -1}" = "'" ]; then
    value="${value:1:${#value}-2}"
  elif [ "${#value}" -ge 2 ] && [ "${value:0:1}" = '"' ] && [ "${value: -1}" = '"' ]; then
    value="${value:1:${#value}-2}"
  fi

  printf '%s' "$value"
}

bool_to_01() {
  local raw
  raw="$(echo "$1" | tr '[:upper:]' '[:lower:]')"
  raw="$(trim "$raw")"
  case "$raw" in
    1|true|yes|y|on|checked)
      echo "1"
      ;;
    *)
      echo "0"
      ;;
  esac
}

is_ipv4_strict() {
  local ip="$1"
  if [[ ! "$ip" =~ ^([0-9]{1,3}\.){3}[0-9]{1,3}$ ]]; then
    return 1
  fi

  local IFS='.'
  local o1 o2 o3 o4
  read -r o1 o2 o3 o4 <<< "$ip"
  for octet in "$o1" "$o2" "$o3" "$o4"; do
    if [ "$octet" -gt 255 ]; then
      return 1
    fi
    if [[ "$octet" =~ ^0[0-9]+$ ]]; then
      return 1
    fi
  done

  return 0
}

is_ipv4_cidr_strict() {
  local cidr="$1"
  if [[ ! "$cidr" =~ ^([^/]+)/([0-9]{1,2})$ ]]; then
    return 1
  fi

  local ip="${BASH_REMATCH[1]}"
  local prefix="${BASH_REMATCH[2]}"
  is_ipv4_strict "$ip" || return 1
  [ "$prefix" -ge 0 ] && [ "$prefix" -le 32 ]
}

skip_csf_remote_pattern() {
  [ "$1" = "0.0.0.0/0" ]
}

sanitize_note() {
  local note="$1"
  note="${note//$'\r'/ }"
  note="${note//$'\n'/ }"
  note="${note//$'\t'/ }"
  note="${note//#/ }"
  note="$(echo "$note" | sed -E 's/[[:space:]]+/ /g; s/^ //; s/ $//')"
  note="${note:0:180}"
  printf '%s' "$note"
}

load_plugin_settings() {
  local plain_allow=""
  local legacy_allow=""

  if [ ! -r "$PLUGIN_SETTINGS_FILE" ]; then
    ALLOW_REMOTE_HOSTS="1"
    return
  fi

  while IFS= read -r raw_line || [ -n "$raw_line" ]; do
    raw_line="$(trim "$raw_line")"
    [ -n "$raw_line" ] || continue
    [[ "$raw_line" =~ ^# ]] && continue
    [[ "$raw_line" == *"="* ]] || continue

    local key="${raw_line%%=*}"
    local value="${raw_line#*=}"
    key="$(trim "$key")"
    value="$(unquote "$(strip_inline_comment "$value")")"

    case "$key" in
      allow_remote_hosts)
        plain_allow="$value"
        ;;
      PG_PLUGIN_ALLOW_REMOTE_HOSTS)
        legacy_allow="$value"
        ;;
    esac
  done < "$PLUGIN_SETTINGS_FILE"

  if [ -n "$plain_allow" ]; then
    ALLOW_REMOTE_HOSTS="$(bool_to_01 "$plain_allow")"
  elif [ -n "$legacy_allow" ]; then
    ALLOW_REMOTE_HOSTS="$(bool_to_01 "$legacy_allow")"
  else
    ALLOW_REMOTE_HOSTS="1"
  fi
}

load_credentials() {
  [ -r "$CREDS_FILE" ] || {
    echo "Brak pliku poświadczeń: $CREDS_FILE" >&2
    exit 1
  }

  local line
  line="$(grep -Ev '^[[:space:]]*($|#)' "$CREDS_FILE" | head -n1 || true)"
  [ -n "$line" ] || {
    echo "Pusty plik poświadczeń: $CREDS_FILE" >&2
    exit 1
  }

  if [[ "$line" == *=* ]]; then
    while IFS='=' read -r key value; do
      key="$(echo "$key" | tr -d '[:space:]' | tr '[:lower:]' '[:upper:]')"
      value="$(trim "$value")"
      value="$(unquote "$value")"
      case "$key" in
        PG_HOST) PGHOST_VAL="$value" ;;
        PG_PORT) PGPORT_VAL="$value" ;;
        PG_USER) PGUSER_VAL="$value" ;;
        PG_PASSWORD) PGPASSWORD_VAL="$value" ;;
      esac
    done < <(grep -Ev '^[[:space:]]*($|#)' "$CREDS_FILE")
  else
    local rest dbname
    PGHOST_VAL="${line%%:*}"
    rest="${line#*:}"
    PGPORT_VAL="${rest%%:*}"
    rest="${rest#*:}"
    dbname="${rest%%:*}"
    rest="${rest#*:}"
    PGUSER_VAL="${rest%%:*}"
    PGPASSWORD_VAL="${rest#*:}"

    [ -n "$dbname" ] || dbname="postgres"
  fi

  [ -n "$PGHOST_VAL" ] || PGHOST_VAL="localhost"
  [ -n "$PGPORT_VAL" ] || PGPORT_VAL="5432"
  [ -n "$PGUSER_VAL" ] || PGUSER_VAL="postgres"
  [ -n "$PGPASSWORD_VAL" ] || {
    echo "Brak hasła PostgreSQL w $CREDS_FILE" >&2
    exit 1
  }
}

psql_exec() {
  PGPASSWORD="$PGPASSWORD_VAL" "$PSQL_BIN" -h "$PGHOST_VAL" -p "$PGPORT_VAL" -U "$PGUSER_VAL" -d postgres -v ON_ERROR_STOP=1 -At -F $'\t' "$@"
}

load_server_version_num() {
  PG_VERSION_NUM="$(psql_exec -c 'SHOW server_version_num;' | tr -d '[:space:]')"
  [[ "$PG_VERSION_NUM" =~ ^[0-9]+$ ]] || {
    echo "Nieprawidłowy server_version_num: ${PG_VERSION_NUM}" >&2
    exit 1
  }
}

cleanup_hba_managed_entries() {
  local hba_file="$1"
  local tmp_file
  tmp_file="$(mktemp)"

  awk -v begin="$HBA_MANAGED_BLOCK_BEGIN" -v end="$HBA_MANAGED_BLOCK_END" '
    $0 == begin {skip=1; next}
    $0 == end {skip=0; next}
    skip == 1 {next}
    {
      lower = tolower($0)
      if (lower ~ /^[[:space:]]*include([[:space:]]|_if_exists)/ && lower ~ /pg_hba_da_plugin\.conf/) {
        next
      }
      print
    }
  ' "$hba_file" > "$tmp_file"

  install -m 600 "$tmp_file" "$hba_file"
  chown postgres:postgres "$hba_file" 2>/dev/null || true
  rm -f "$tmp_file"
}

write_hba_inline_block() {
  local hba_file="$1"
  local include_file="$2"
  local tmp_file
  tmp_file="$(mktemp)"

  cp -p "$hba_file" "${hba_file}.bak.$(date +%s)"
  cat "$hba_file" > "$tmp_file"

  {
    echo ""
    echo "$HBA_MANAGED_BLOCK_BEGIN"
    echo "# Managed by DirectAdmin PostgreSQL plugin (inline mode)."
    cat "$include_file"
    echo "$HBA_MANAGED_BLOCK_END"
  } >> "$tmp_file"

  install -m 600 "$tmp_file" "$hba_file"
  chown postgres:postgres "$hba_file" 2>/dev/null || true
  rm -f "$tmp_file"
}

append_remote_ip_note() {
  local ip="$1"
  local note="$2"

  if [ -z "${REMOTE_IP_NOTES[$ip]+x}" ]; then
    REMOTE_IP_NOTES["$ip"]="$note"
    return
  fi

  if [ -z "$note" ]; then
    return
  fi

  local existing="${REMOTE_IP_NOTES[$ip]}"
  if [ -z "$existing" ]; then
    REMOTE_IP_NOTES["$ip"]="$note"
    return
  fi

  if [[ "$existing" != *"$note"* ]]; then
    REMOTE_IP_NOTES["$ip"]="$existing | $note"
  fi
}

expand_hba_addresses() {
  local host="$1"
  case "$host" in
    localhost)
      echo "127.0.0.1/32"
      echo "::1/128"
      return
      ;;
    127.0.0.1)
      echo "127.0.0.1/32"
      return
      ;;
    ::1)
      echo "::1/128"
      return
      ;;
  esac

  if is_ipv4_strict "$host"; then
    echo "$host/32"
    return
  fi

  if is_ipv4_cidr_strict "$host"; then
    echo "$host"
  fi
}

build_hba_include() {
  local include_file="$1"
  local tmp_file
  tmp_file="$(mktemp)"

  {
    echo "# Auto-generated by DirectAdmin PostgreSQL plugin"
    echo "# Generated at: $(date -u +'%Y-%m-%d %H:%M:%S UTC')"
    echo "# Method: scram-sha-256"
  } > "$tmp_file"

  local sql_query
  sql_query="SELECT d.datname, h.role_name, h.host_pattern, COALESCE(h.note, '')
FROM da_plugin.access_hosts h
JOIN pg_roles r ON r.rolname = h.role_name
JOIN pg_database d
  ON d.datistemplate = false
 AND (
      (COALESCE(h.db_name, '') <> '' AND d.datname = h.db_name)
      OR
      (COALESCE(h.db_name, '') = '' AND left(d.datname, length(h.da_user) + 1) = h.da_user || '_')
 )
WHERE COALESCE(h.db_name, '') <> ''
   OR has_database_privilege(h.role_name, d.datname, 'CONNECT')
ORDER BY d.datname, h.role_name, h.host_pattern;"

  local query_output
  query_output="$(psql_exec -c "$sql_query")" || {
    local status=$?
    rm -f "$tmp_file"
    echo "Nie udało się odczytać rekordów hostów z da_plugin.access_hosts." >&2
    return "$status"
  }

  local access_host_rows="0"
  local generated_rules="0"

  while IFS=$'\t' read -r db role host note; do
    [ -n "$db" ] || continue
    [ -n "$role" ] || continue
    [ -n "$host" ] || continue
    access_host_rows=$((access_host_rows + 1))

    local host_note
    host_note="$(sanitize_note "$note")"
    local include_host="0"
    local is_remote="0"

    case "$host" in
      localhost|127.0.0.1|::1)
        include_host="1"
        ;;
      *)
        if [ "$ALLOW_REMOTE_HOSTS" = "1" ] && { is_ipv4_strict "$host" || is_ipv4_cidr_strict "$host"; }; then
          include_host="1"
          is_remote="1"
        fi
        ;;
    esac

    [ "$include_host" = "1" ] || continue

    if [ "$is_remote" = "1" ]; then
      append_remote_ip_note "$host" "$host_note"
    fi

    local note_label="$host_note"
    [ -n "$note_label" ] || note_label="(brak komentarza)"
    printf '# role=%s db=%s host=%s comment=%s\n' "$role" "$db" "$host" "$note_label" >> "$tmp_file"

    while IFS= read -r addr; do
      [ -n "$addr" ] || continue
      printf 'host\t%s\t%s\t%s\tscram-sha-256\n' "$db" "$role" "$addr" >> "$tmp_file"
      generated_rules=$((generated_rules + 1))
    done < <(expand_hba_addresses "$host")
  done <<< "$query_output"

  install -m 600 "$tmp_file" "$include_file"
  chown postgres:postgres "$include_file" 2>/dev/null || true
  rm -f "$tmp_file"

  echo "Read access host rows: $access_host_rows"
  echo "Generated pg_hba host rules: $generated_rules"
}

sorted_remote_ips() {
  if [ "${#REMOTE_IP_NOTES[@]}" -eq 0 ]; then
    return 0
  fi

  printf '%s\n' "${!REMOTE_IP_NOTES[@]}" | sort -V
}

write_postgresql_conf_block() {
  local conf_file="$1"
  local listen_value="localhost"
  local has_remote_ips="0"

  if [ "$ALLOW_REMOTE_HOSTS" = "1" ]; then
    listen_value="*"
  fi

  if [ "$ALLOW_REMOTE_HOSTS" = "1" ] && [ "${#REMOTE_IP_NOTES[@]}" -gt 0 ]; then
    has_remote_ips="1"
  fi

  local tmp_file
  tmp_file="$(mktemp)"

  awk -v begin="$PG_CONF_MARKER_BEGIN" -v end="$PG_CONF_MARKER_END" '
    $0 == begin {skip=1; next}
    $0 == end {skip=0; next}
    skip != 1 {print}
  ' "$conf_file" > "$tmp_file"

  {
    echo ""
    echo "$PG_CONF_MARKER_BEGIN"
    echo "# Managed by DirectAdmin PostgreSQL plugin."
    echo "listen_addresses = '$listen_value'"
    echo "# Zdefiniowane hosty IPv4 z pluginu:"
    if [ "$has_remote_ips" = "1" ]; then
      while IFS= read -r ip; do
        [ -n "$ip" ] || continue
        local note="${REMOTE_IP_NOTES[$ip]}"
        if [ -n "$note" ]; then
          printf '# %s ; %s\n' "$ip" "$note"
        else
          printf '# %s\n' "$ip"
        fi
      done < <(sorted_remote_ips)
    else
      echo "# brak zdalnych hostów"
    fi
    echo "$PG_CONF_MARKER_END"
  } >> "$tmp_file"

  install -m 600 "$tmp_file" "$conf_file"
  chown postgres:postgres "$conf_file" 2>/dev/null || true
  rm -f "$tmp_file"
}

ensure_csf_include_directive() {
  local allow_file="$1"
  local include_file="$2"

  [ -f "$allow_file" ] || return 1

  local include_line="Include $include_file"
  if ! grep -Fxq "$include_line" "$allow_file"; then
    echo "$include_line" >> "$allow_file"
  fi
  return 0
}

write_csf_include_file() {
  local include_file="$1"
  local tmp_file
  tmp_file="$(mktemp)"

  {
    echo "# Auto-generated by DirectAdmin PostgreSQL plugin"
    echo "# Generated at: $(date -u +'%Y-%m-%d %H:%M:%S UTC')"
  } > "$tmp_file"

  if [ "$ALLOW_REMOTE_HOSTS" = "1" ] && [ "${#REMOTE_IP_NOTES[@]}" -gt 0 ]; then
    while IFS= read -r ip; do
      [ -n "$ip" ] || continue
      if skip_csf_remote_pattern "$ip"; then
        continue
      fi
      local note="${REMOTE_IP_NOTES[$ip]}"
      if [ -n "$note" ]; then
        printf '%s # DA PostgreSQL: %s\n' "$ip" "$note" >> "$tmp_file"
      else
        printf '%s # DA PostgreSQL\n' "$ip" >> "$tmp_file"
      fi
    done < <(sorted_remote_ips)
  fi

  install -m 600 "$tmp_file" "$include_file"
  chown root:root "$include_file" 2>/dev/null || true
  rm -f "$tmp_file"
}

reload_postgres_conf() {
  psql_exec -c 'SELECT pg_reload_conf();' >/dev/null
}

restart_postgresql_service() {
  if ! command -v systemctl >/dev/null 2>&1; then
    echo "Nie znaleziono systemctl, nie można zrestartować PostgreSQL." >&2
    return 1
  fi

  local major=""
  if [[ "$PG_VERSION_NUM" =~ ^[0-9]+$ ]] && [ "$PG_VERSION_NUM" -ge 100000 ]; then
    major="$((PG_VERSION_NUM / 10000))"
  fi

  local candidates=()
  if [ -n "$major" ]; then
    candidates+=("postgresql-${major}" "postgresql@${major}-main")
  fi
  candidates+=("postgresql")

  local service
  for service in "${candidates[@]}"; do
    if systemctl list-unit-files --no-legend "${service}.service" 2>/dev/null | grep -q . || systemctl status "$service" >/dev/null 2>&1; then
      if systemctl restart "$service"; then
        echo "Restarted PostgreSQL service: $service"
        return 0
      fi
    fi
  done

  echo "Nie znaleziono usługi PostgreSQL do restartu. Próbowano: ${candidates[*]}" >&2
  return 1
}

reload_csf_if_available() {
  if command -v csf >/dev/null 2>&1; then
    csf -r >/dev/null 2>&1 || true
  fi
}

parse_args "$@"
load_plugin_settings
load_credentials

PSQL_BIN="$(find_psql)"
[ -n "$PSQL_BIN" ] || {
  echo "Nie znaleziono binarki psql" >&2
  exit 1
}

HBA_FILE="$(psql_exec -c 'SHOW hba_file;')"
[ -n "$HBA_FILE" ] || {
  echo "Nie udało się odczytać ścieżki do pg_hba.conf" >&2
  exit 1
}

POSTGRESQL_CONF="$(psql_exec -c 'SHOW config_file;')"
[ -n "$POSTGRESQL_CONF" ] || {
  echo "Nie udało się odczytać ścieżki do postgresql.conf" >&2
  exit 1
}

load_server_version_num

HBA_DIR="$(dirname "$HBA_FILE")"
HBA_INCLUDE_FILE="$HBA_DIR/pg_hba_da_plugin.conf"

build_hba_include "$HBA_INCLUDE_FILE"
cleanup_hba_managed_entries "$HBA_FILE"
write_hba_inline_block "$HBA_FILE" "$HBA_INCLUDE_FILE"

CSF_SYNCED="0"
CSF_SKIP_REASON=""

if [ "$ALLOW_REMOTE_HOSTS" = "1" ]; then
  write_postgresql_conf_block "$POSTGRESQL_CONF"

  if ensure_csf_include_directive "$CSF_ALLOW_FILE" "$CSF_INCLUDE_FILE"; then
    write_csf_include_file "$CSF_INCLUDE_FILE"
    reload_csf_if_available
    CSF_SYNCED="1"
  else
    CSF_SKIP_REASON="CSF allow file not found: $CSF_ALLOW_FILE"
    echo "Ostrzeżenie: nie znaleziono $CSF_ALLOW_FILE, pomijam aktualizację CSF." >&2
  fi
fi

reload_postgres_conf

if [ "$RESTART_POSTGRESQL" = "1" ]; then
  restart_postgresql_service
fi

echo "Synchronized pg_hba inline block in: $HBA_FILE"
echo "Generated pg_hba include file (diagnostic copy): $HBA_INCLUDE_FILE"
if [ "$ALLOW_REMOTE_HOSTS" = "1" ]; then
  echo "Synchronized postgresql.conf: $POSTGRESQL_CONF"
  echo "PostgreSQL restart required for listen_addresses changes"
  if [ "$CSF_SYNCED" = "1" ]; then
    echo "Synchronized CSF include: $CSF_INCLUDE_FILE"
  else
    echo "Skipped CSF synchronization: ${CSF_SKIP_REASON:-CSF allow file unavailable}"
  fi
else
  echo "Remote hosts disabled: skipped postgresql.conf and CSF synchronization"
fi