#!/usr/bin/env bash
set -euo pipefail

SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
# shellcheck source=/dev/null
source "$SCRIPT_DIR/libenv.sh"

ENV_FILE=$(limristem_mail_resolve_main_env_file)

declare -A LIMIT_ENV_MAP=(
  [postfix-client-connection-rate-limit]=LIMRISTEM_MAIL_POSTFIX_CLIENT_CONNECTION_RATE_LIMIT
  [postfix-client-message-rate-limit]=LIMRISTEM_MAIL_POSTFIX_CLIENT_MESSAGE_RATE_LIMIT
  [postfix-rate-time-unit]=LIMRISTEM_MAIL_POSTFIX_RATE_TIME_UNIT
  [api-auth-fail-limit]=LIMRISTEM_MAIL_API_AUTH_FAIL_LIMIT
  [api-auth-window-seconds]=LIMRISTEM_MAIL_API_AUTH_WINDOW_SECONDS
  [api-auth-block-seconds]=LIMRISTEM_MAIL_API_AUTH_BLOCK_SECONDS
  [rspamd-action-greylist]=LIMRISTEM_MAIL_RSPAMD_ACTION_GREYLIST
  [rspamd-action-add-header]=LIMRISTEM_MAIL_RSPAMD_ACTION_ADD_HEADER
  [rspamd-action-reject]=LIMRISTEM_MAIL_RSPAMD_ACTION_REJECT
  [rspamd-greylist-delay]=LIMRISTEM_MAIL_RSPAMD_GREYLIST_DELAY
  [rspamd-greylist-expire]=LIMRISTEM_MAIL_RSPAMD_GREYLIST_EXPIRE
)

declare -A LIMIT_DEFAULT_MAP=(
  [postfix-client-connection-rate-limit]=30
  [postfix-client-message-rate-limit]=100
  [postfix-rate-time-unit]=60s
  [api-auth-fail-limit]=5
  [api-auth-window-seconds]=300
  [api-auth-block-seconds]=900
  [rspamd-action-greylist]=4
  [rspamd-action-add-header]=6
  [rspamd-action-reject]=15
  [rspamd-greylist-delay]=5m
  [rspamd-greylist-expire]=35d
)

usage() {
  cat <<'EOF'
Usage:
  manage-limits.sh show [--json]
  manage-limits.sh set <key> <value>
  manage-limits.sh set-many <key> <value> [<key> <value> ...]

Supported keys:
  postfix-client-connection-rate-limit
  postfix-client-message-rate-limit
  postfix-rate-time-unit
  api-auth-fail-limit
  api-auth-window-seconds
  api-auth-block-seconds
  rspamd-action-greylist
  rspamd-action-add-header
  rspamd-action-reject
  rspamd-greylist-delay
  rspamd-greylist-expire
EOF
}

require_root() {
  if [[ $EUID -ne 0 ]]; then
    echo "Run as root." >&2
    exit 1
  fi
}

load_env() {
  limristem_mail_load_env_file "$ENV_FILE"
}

set_env_value() {
  local key=$1
  local value=$2
  limristem_mail_upsert_env_value "$ENV_FILE" "$key" "$value"
}

show_limits() {
  local as_json=${1:-no}
  local key env_key value
  if [[ "$as_json" == "yes" ]]; then
    printf '{'
    local first=yes
    for key in "${!LIMIT_ENV_MAP[@]}"; do
      env_key=${LIMIT_ENV_MAP[$key]}
      value=${!env_key:-${LIMIT_DEFAULT_MAP[$key]}}
      if [[ "$first" == "yes" ]]; then
        first=no
      else
        printf ','
      fi
      python3 - "$key" "$value" <<'PY'
import json
import sys
print(json.dumps(sys.argv[1]) + ":" + json.dumps(sys.argv[2]), end="")
PY
    done
    printf '}\n'
    return 0
  fi
  for key in "${!LIMIT_ENV_MAP[@]}"; do
    env_key=${LIMIT_ENV_MAP[$key]}
    value=${!env_key:-${LIMIT_DEFAULT_MAP[$key]}}
    printf '%s=%s\n' "$key" "$value"
  done | sort
}

detect_rspamd_user() {
  if id _rspamd >/dev/null 2>&1; then
    printf '_rspamd\n'
  elif id rspamd >/dev/null 2>&1; then
    printf 'rspamd\n'
  fi
}

write_rspamd_limits() {
  local rspamd_service_user
  mkdir -p /etc/rspamd/local.d
  cat > /etc/rspamd/local.d/actions.conf <<EOF
reject = ${LIMRISTEM_MAIL_RSPAMD_ACTION_REJECT:-15};
add_header = ${LIMRISTEM_MAIL_RSPAMD_ACTION_ADD_HEADER:-6};
greylist = ${LIMRISTEM_MAIL_RSPAMD_ACTION_GREYLIST:-4};
EOF
  cat > /etc/rspamd/local.d/greylist.conf <<EOF
enabled = ${LIMRISTEM_MAIL_RSPAMD_GREYLIST_ENABLED:-true};
timeout = ${LIMRISTEM_MAIL_RSPAMD_GREYLIST_DELAY:-5m};
expire = ${LIMRISTEM_MAIL_RSPAMD_GREYLIST_EXPIRE:-35d};
key_prefix = "grey";
EOF
  rspamd_service_user=$(detect_rspamd_user || true)
  if [[ -n "$rspamd_service_user" ]]; then
    chown root:"$rspamd_service_user" /etc/rspamd/local.d/actions.conf /etc/rspamd/local.d/greylist.conf
    chmod 640 /etc/rspamd/local.d/actions.conf /etc/rspamd/local.d/greylist.conf
  fi
}

apply_limits() {
  if command -v postconf >/dev/null 2>&1; then
    postconf -e "anvil_rate_time_unit = ${LIMRISTEM_MAIL_POSTFIX_RATE_TIME_UNIT:-60s}"
    postconf -e "smtpd_client_connection_rate_limit = ${LIMRISTEM_MAIL_POSTFIX_CLIENT_CONNECTION_RATE_LIMIT:-30}"
    postconf -e "smtpd_client_message_rate_limit = ${LIMRISTEM_MAIL_POSTFIX_CLIENT_MESSAGE_RATE_LIMIT:-100}"
    systemctl reload postfix >/dev/null 2>&1 || systemctl restart postfix >/dev/null 2>&1 || true
  fi

  if [[ "${LIMRISTEM_MAIL_ENABLE_RSPAMD:-yes}" == "yes" ]]; then
    write_rspamd_limits
    systemctl restart rspamd >/dev/null 2>&1 || true
  fi

  systemctl restart --no-block limristem-mail >/dev/null 2>&1 || true
}

set_limit() {
  local key=$1
  local value=$2
  local env_key=${LIMIT_ENV_MAP[$key]:-}
  if [[ -z "$env_key" ]]; then
    echo "Unknown limit key: $key" >&2
    exit 1
  fi
  set_env_value "$env_key" "$value"
  load_env
  apply_limits
}

set_limits_batch() {
  if (( $# == 0 || $# % 2 != 0 )); then
    echo "set-many requires <key> <value> pairs" >&2
    exit 1
  fi
  local key value env_key
  while (( $# > 0 )); do
    key=$1
    value=$2
    env_key=${LIMIT_ENV_MAP[$key]:-}
    if [[ -z "$env_key" ]]; then
      echo "Unknown limit key: $key" >&2
      exit 1
    fi
    set_env_value "$env_key" "$value"
    shift 2
  done
  load_env
  apply_limits
}

require_root
load_env

command=${1:-}
case "$command" in
  show)
    if [[ ${2:-} == "--json" ]]; then
      show_limits yes
    else
      show_limits no
    fi
    ;;
  set)
    set_limit "${2:?key required}" "${3:?value required}"
    ;;
  set-many)
    shift
    set_limits_batch "$@"
    ;;
  *)
    usage >&2
    exit 1
    ;;
esac