#!/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
  [message-size-limit]=LIMRISTEM_MAIL_MESSAGE_SIZE_LIMIT
  [smtp-recipient-limit]=LIMRISTEM_MAIL_SMTP_RECIPIENT_LIMIT
  [smtp-client-connection-count-limit]=LIMRISTEM_MAIL_SMTP_CLIENT_CONNECTION_COUNT_LIMIT
  [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
  [send-rate-burst]=LIMRISTEM_MAIL_SEND_RATE_BURST
  [send-rate-per-minute]=LIMRISTEM_MAIL_SEND_RATE_PER_MINUTE
  [send-rate-per-hour]=LIMRISTEM_MAIL_SEND_RATE_PER_HOUR
  [send-rate-per-day]=LIMRISTEM_MAIL_SEND_RATE_PER_DAY
  [send-rate-per-week]=LIMRISTEM_MAIL_SEND_RATE_PER_WEEK
  [send-rate-per-month]=LIMRISTEM_MAIL_SEND_RATE_PER_MONTH
  [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
  [dkim-rotation-interval-days]=LIMRISTEM_MAIL_DKIM_ROTATION_INTERVAL_DAYS
  [dkim-overlap-days]=LIMRISTEM_MAIL_DKIM_OVERLAP_DAYS
  [dkim-auto-rotate-default]=LIMRISTEM_MAIL_DKIM_AUTO_ROTATE
)

declare -A LIMIT_DEFAULT_MAP=(
  [postfix-client-connection-rate-limit]=30
  [postfix-client-message-rate-limit]=100
  [postfix-rate-time-unit]=60s
  [message-size-limit]=52428800
  [smtp-recipient-limit]=100
  [smtp-client-connection-count-limit]=50
  [api-auth-fail-limit]=5
  [api-auth-window-seconds]=300
  [api-auth-block-seconds]=900
  [send-rate-burst]=20
  [send-rate-per-minute]=10
  [send-rate-per-hour]=100
  [send-rate-per-day]=500
  [send-rate-per-week]=2000
  [send-rate-per-month]=5000
  [rspamd-action-greylist]=4
  [rspamd-action-add-header]=6
  [rspamd-action-reject]=15
  [rspamd-greylist-delay]=5m
  [rspamd-greylist-expire]=35d
  [dkim-rotation-interval-days]=90
  [dkim-overlap-days]=14
  [dkim-auto-rotate-default]=no
)

# Preferred display order for panel/API consumers.
LIMIT_KEY_ORDER=(
  api-auth-fail-limit
  api-auth-window-seconds
  api-auth-block-seconds
  message-size-limit
  smtp-recipient-limit
  smtp-client-connection-count-limit
  postfix-client-connection-rate-limit
  postfix-client-message-rate-limit
  postfix-rate-time-unit
  send-rate-burst
  send-rate-per-minute
  send-rate-per-hour
  send-rate-per-day
  send-rate-per-week
  send-rate-per-month
  rspamd-action-greylist
  rspamd-action-add-header
  rspamd-action-reject
  rspamd-greylist-delay
  rspamd-greylist-expire
  dkim-rotation-interval-days
  dkim-overlap-days
  dkim-auto-rotate-default
)

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:
  api-auth-fail-limit / api-auth-window-seconds / api-auth-block-seconds
  message-size-limit / smtp-recipient-limit / smtp-client-connection-count-limit
  postfix-client-connection-rate-limit / postfix-client-message-rate-limit / postfix-rate-time-unit
  send-rate-burst / send-rate-per-minute / send-rate-per-hour / send-rate-per-day / send-rate-per-week / send-rate-per-month
  rspamd-action-greylist / rspamd-action-add-header / rspamd-action-reject
  rspamd-greylist-delay / rspamd-greylist-expire
  dkim-rotation-interval-days / dkim-overlap-days / dkim-auto-rotate-default
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"
}

validate_limit_value() {
  local key=$1
  local value=$2
  case "$key" in
    api-auth-fail-limit|api-auth-window-seconds|api-auth-block-seconds|message-size-limit|smtp-recipient-limit|smtp-client-connection-count-limit|postfix-client-connection-rate-limit|postfix-client-message-rate-limit|send-rate-burst|send-rate-per-minute|send-rate-per-hour|send-rate-per-day|send-rate-per-week|send-rate-per-month|dkim-rotation-interval-days|dkim-overlap-days|rspamd-action-greylist|rspamd-action-add-header|rspamd-action-reject)
      if [[ ! "$value" =~ ^[0-9]+$ ]] || (( value < 0 )); then
        echo "Invalid numeric value for $key: $value" >&2
        exit 1
      fi
      ;;
    postfix-rate-time-unit|rspamd-greylist-delay|rspamd-greylist-expire)
      if [[ ! "$value" =~ ^[0-9]+[smhd]?$ ]]; then
        echo "Invalid duration value for $key: $value" >&2
        exit 1
      fi
      ;;
    dkim-auto-rotate-default)
      case "${value,,}" in
        yes|no|true|false|1|0|on|off) ;;
        *)
          echo "Invalid boolean value for $key: $value (use yes/no)" >&2
          exit 1
          ;;
      esac
      ;;
  esac
}

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_KEY_ORDER[@]}"; 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_KEY_ORDER[@]}"; do
    env_key=${LIMIT_ENV_MAP[$key]}
    value=${!env_key:-${LIMIT_DEFAULT_MAP[$key]}}
    printf '%s=%s\n' "$key" "$value"
  done
}

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
  local burst=${LIMRISTEM_MAIL_SEND_RATE_BURST:-20}
  local per_min=${LIMRISTEM_MAIL_SEND_RATE_PER_MINUTE:-10}
  local per_hour=${LIMRISTEM_MAIL_SEND_RATE_PER_HOUR:-100}
  local per_day=${LIMRISTEM_MAIL_SEND_RATE_PER_DAY:-500}
  local per_week=${LIMRISTEM_MAIL_SEND_RATE_PER_WEEK:-2000}
  local per_month=${LIMRISTEM_MAIL_SEND_RATE_PER_MONTH:-5000}
  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
  # Per-authenticated-user outbound rate limits (anti-spam / abuse).
  cat > /etc/rspamd/local.d/ratelimit.conf <<EOF
# Generated by Limristem eMail manage-limits.sh
# Limits apply to authenticated SMTP users (selector = user).
rates {
  send_mail = {
    selector = "user";
    bucket = [
      {
        burst = ${burst};
        rate = "${per_min} / 1m";
      },
      {
        burst = ${burst};
        rate = "${per_hour} / 1h";
      },
      {
        burst = ${burst};
        rate = "${per_day} / 1d";
      },
      {
        burst = ${burst};
        rate = "${per_week} / 7d";
      },
      {
        burst = ${burst};
        rate = "${per_month} / 30d";
      }
    ];
  }
}
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 \
      /etc/rspamd/local.d/ratelimit.conf
    chmod 640 \
      /etc/rspamd/local.d/actions.conf \
      /etc/rspamd/local.d/greylist.conf \
      /etc/rspamd/local.d/ratelimit.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}"
    postconf -e "message_size_limit = ${LIMRISTEM_MAIL_MESSAGE_SIZE_LIMIT:-52428800}"
    postconf -e "mailbox_size_limit = 0"
    postconf -e "smtpd_recipient_limit = ${LIMRISTEM_MAIL_SMTP_RECIPIENT_LIMIT:-100}"
    postconf -e "smtpd_client_connection_count_limit = ${LIMRISTEM_MAIL_SMTP_CLIENT_CONNECTION_COUNT_LIMIT:-50}"
    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

  # Rebuild per-mailbox Redis profiles so overrides re-merge against new globals.
  if [[ -x "${LIMRISTEM_MAIL_BASE_DIR:-/opt/limristem-mail}/limristem-mail" ]]; then
    sudo -u limristem-mail -n "${LIMRISTEM_MAIL_BASE_DIR:-/opt/limristem-mail}/limristem-mail" accounts sync-send-limits --json >/dev/null 2>&1 \
      || "${LIMRISTEM_MAIL_BASE_DIR:-/opt/limristem-mail}/limristem-mail" accounts sync-send-limits --json >/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
  validate_limit_value "$key" "$value"
  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
    validate_limit_value "$key" "$value"
    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