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

service_name=${1:-}
if [[ ! $service_name =~ ^[a-z0-9][a-z0-9-]*$ ]]; then
  echo "usage: $0 <service-name>" >&2
  exit 2
fi

if [[ $EUID -ne 0 && ${HOST_MEMORY_ALLOW_NON_ROOT:-0} != 1 ]]; then
  echo "host memory configuration must run as root" >&2
  exit 1
fi

meminfo=${HOST_MEMORY_MEMINFO:-/proc/meminfo}
if [[ ! -r $meminfo ]]; then
  echo "cannot read physical memory information: $meminfo" >&2
  exit 1
fi
memory_total_bytes=$(awk '$1 == "MemTotal:" { printf "%.0f", $2 * 1024; exit }' "$meminfo")
if [[ ! $memory_total_bytes =~ ^[1-9][0-9]*$ ]]; then
  echo "failed to read MemTotal from $meminfo" >&2
  exit 1
fi
memory_total_mib=$((memory_total_bytes / 1048576))

# The memory class now only selects swappiness and the recommended swap size.
# The cgroup limits themselves are derived proportionally from host RAM below,
# because a fixed table plateaus: it handed a 16 GiB host the same sub-gigabyte
# API budget as an 8 GiB host while 14 GiB sat unused.
two_gib=2147483648
four_gib=4294967296
eight_gib=8589934592
sixteen_gib=17179869184
if ((memory_total_bytes <= two_gib)); then
  memory_class=small
  automatic_swappiness=80
elif ((memory_total_bytes <= four_gib)); then
  memory_class=medium
  automatic_swappiness=60
elif ((memory_total_bytes <= eight_gib)); then
  memory_class=large
  automatic_swappiness=40
elif ((memory_total_bytes <= sixteen_gib)); then
  memory_class=xlarge
  automatic_swappiness=40
else
  memory_class=huge
  automatic_swappiness=30
fi

# Swap is no longer treated as optional on large hosts. A cgroup with no swap
# cannot reclaim anonymous memory at all, so once it crosses MemoryHigh the
# kernel throttles every allocation without ever being able to bring usage back
# under the soft limit. Swap is what lets that reclaim make progress; it is a
# stall-avoidance device here, not a way to oversubscribe RAM.
recommended_swap_mib=$((memory_total_mib / 8))
((recommended_swap_mib < 2048)) && recommended_swap_mib=2048
((recommended_swap_mib > 4096)) && recommended_swap_mib=4096
recommended_swap_bytes=$((recommended_swap_mib * 1048576))

cpu_count=${HOST_MEMORY_CPU_COUNT:-}
if [[ -z $cpu_count ]]; then
  cpu_count=$(getconf _NPROCESSORS_ONLN 2>/dev/null || true)
fi
if [[ ! $cpu_count =~ ^[1-9][0-9]*$ ]]; then
  echo "failed to determine online CPU count" >&2
  exit 1
fi

read_percent_env() {
  local env_name=$1
  local default_value=$2
  local maximum=$3
  local value=${!env_name:-$default_value}
  if [[ ! $value =~ ^[0-9]+$ ]] || ((value < 1 || value > maximum)); then
    echo "$env_name must be an integer from 1 to $maximum" >&2
    exit 2
  fi
  printf '%s' "$value"
}

# Headroom left to the OS, page cache, the openresty container and anything else
# resident on the box. Everything below is carved out of what remains.
reserve_percent=$(read_percent_env HOST_MEMORY_RESERVE_PERCENT 30 90)
reserve_mib=$((memory_total_mib * reserve_percent / 100))
((reserve_mib < 768)) && reserve_mib=768
((reserve_mib > 4096)) && reserve_mib=4096
allocatable_mib=$((memory_total_mib - reserve_mib))
((allocatable_mib < 256)) && allocatable_mib=256

# Limits are rounded down to whole 64 MiB steps so the generated drop-ins stay
# stable across trivial MemTotal jitter (a reboot can shift it by a few MiB).
round_down_step() {
  local value=$1
  local rounded=$((value / 64 * 64))
  ((rounded < 64)) && rounded=64
  printf '%s' "$rounded"
}

share_mib() {
  local percent=$1
  local floor_mib=$2
  local ceiling_mib=$3
  local value=$((allocatable_mib * percent / 100))
  ((value < floor_mib)) && value=$floor_mib
  ((value > ceiling_mib)) && value=$ceiling_mib
  round_down_step "$value"
}

# Shares are sized so the worst concurrent case (a rollout running two API
# processes next to the worker and the daily job) still fits inside the
# allocatable budget. They are also capped: a single-threaded JS process with a
# multi-gigabyte live heap is pathological regardless of host size, and past
# that point a larger limit only lets a leak run longer before it is stopped.
# The worker gets the higher ceiling because a diff legitimately peaks at
# roughly six times the bundle size (DIFF_PEAK_BUNDLE_MULTIPLIER in
# src/worker/taskRunner.ts — keep the two in sync).
# Each percentage is resolved into its own variable first: `exit` inside a
# command substitution only leaves the subshell, so a validation failure nested
# in an argument would be swallowed instead of aborting the run.
api_share_percent=$(read_percent_env HOST_MEMORY_API_SHARE_PERCENT 25 90)
worker_share_percent=$(read_percent_env HOST_MEMORY_WORKER_SHARE_PERCENT 25 90)
dailyjob_share_percent=$(read_percent_env HOST_MEMORY_DAILYJOB_SHARE_PERCENT 12 90)
api_max_mib=$(share_mib "$api_share_percent" 512 2560)
worker_max_mib=$(share_mib "$worker_share_percent" 512 3072)
dailyjob_max_mib=$(share_mib "$dailyjob_share_percent" 384 1536)

proc_swaps=${HOST_MEMORY_PROC_SWAPS:-/proc/swaps}
swap_bytes=0
if [[ -r $proc_swaps ]]; then
  swap_bytes=$(awk 'NR > 1 { total += $3 * 1024 } END { printf "%.0f", total + 0 }' "$proc_swaps")
else
  echo "warning: cannot read active swap table: $proc_swaps" >&2
fi
if ((swap_bytes > 0)); then
  swap_state=present
else
  swap_state=absent
fi

# The gap between MemoryHigh and MemoryMax is the throttle band: above the soft
# limit the kernel makes the allocating task sleep, in proportion to how far it
# has overshot, on every charge. For a single-threaded runtime that band stalls
# request handling, timers and even signal delivery.
#
# With swap the band is useful, because reclaim can push cold anonymous pages
# out and pull usage back under the soft limit. Without swap it is a trap: the
# only reclaimable memory is page cache, so usage stays above the soft limit and
# the process is throttled indefinitely while systemd still reports it healthy.
# So on a swapless host keep the band narrow and let MemoryMax do the work —
# OOMPolicy=stop plus Restart=always turns that into a ~2 s restart, which beats
# an open-ended hang that needs a human.
if [[ $swap_state == present ]]; then
  default_high_percent=85
else
  default_high_percent=95
fi
high_percent=$(read_percent_env HOST_MEMORY_HIGH_PERCENT "$default_high_percent" 99)

derive_high_mib() {
  local max_mib=$1
  local value
  value=$(round_down_step $((max_mib * high_percent / 100)))
  # round_down_step can land back on max_mib for very small limits; always keep
  # at least one step of band so MemoryHigh stays a soft limit.
  ((value >= max_mib)) && value=$((max_mib - 64))
  ((value < 64)) && value=64
  printf '%s' "$value"
}

# Half the hard limit, capped: enough for reclaim to unload a cold heap without
# letting a runaway quietly relocate itself into swap and thrash there.
derive_swap_max_mib() {
  local max_mib=$1
  local value=$((max_mib / 2))
  ((value > 1024)) && value=1024
  round_down_step "$value"
}

resolve_limit_mib() {
  local env_name=$1
  local default_value=$2
  local allow_zero=${3:-false}
  local value=${!env_name:-$default_value}
  if [[ ! $value =~ ^[0-9]+$ ]] ||
    { [[ $allow_zero != true ]] && ((value == 0)); }; then
    echo "$env_name must be a positive integer MiB value" >&2
    exit 2
  fi
  printf '%s' "$value"
}

# Hard limits resolve first so an explicit MemoryMax override also moves the
# soft limit with it, keeping the band proportional instead of accidentally
# widening it back into the stall zone.
api_max_mib=$(resolve_limit_mib HOST_MEMORY_API_MAX_MIB "$api_max_mib")
worker_max_mib=$(resolve_limit_mib HOST_MEMORY_WORKER_MAX_MIB "$worker_max_mib")
dailyjob_max_mib=$(resolve_limit_mib HOST_MEMORY_DAILYJOB_MAX_MIB "$dailyjob_max_mib")

api_high_mib=$(resolve_limit_mib HOST_MEMORY_API_HIGH_MIB "$(derive_high_mib "$api_max_mib")")
worker_high_mib=$(resolve_limit_mib HOST_MEMORY_WORKER_HIGH_MIB "$(derive_high_mib "$worker_max_mib")")
dailyjob_high_mib=$(resolve_limit_mib HOST_MEMORY_DAILYJOB_HIGH_MIB "$(derive_high_mib "$dailyjob_max_mib")")

api_swap_max_mib=$(resolve_limit_mib HOST_MEMORY_API_SWAP_MAX_MIB "$(derive_swap_max_mib "$api_max_mib")" true)
worker_swap_max_mib=$(resolve_limit_mib HOST_MEMORY_WORKER_SWAP_MAX_MIB "$(derive_swap_max_mib "$worker_max_mib")" true)
dailyjob_swap_max_mib=$(resolve_limit_mib HOST_MEMORY_DAILYJOB_SWAP_MAX_MIB "$(derive_swap_max_mib "$dailyjob_max_mib")" true)

if ((api_high_mib >= api_max_mib)); then
  echo "HOST_MEMORY_API_HIGH_MIB must be lower than HOST_MEMORY_API_MAX_MIB" >&2
  exit 2
fi
if ((worker_high_mib >= worker_max_mib)); then
  echo "HOST_MEMORY_WORKER_HIGH_MIB must be lower than HOST_MEMORY_WORKER_MAX_MIB" >&2
  exit 2
fi
if ((dailyjob_high_mib >= dailyjob_max_mib)); then
  echo "HOST_MEMORY_DAILYJOB_HIGH_MIB must be lower than HOST_MEMORY_DAILYJOB_MAX_MIB" >&2
  exit 2
fi

# A rollout runs the primary and the candidate API side by side, so the worst
# concurrent case charges the API budget twice.
projected_peak_mib=$((2 * api_max_mib + worker_max_mib + dailyjob_max_mib))
if ((projected_peak_mib > allocatable_mib)); then
  echo "warning: worst-case unit total exceeds the allocatable budget: ${projected_peak_mib} > ${allocatable_mib} MiB" >&2
fi

desired_swappiness=${HOST_MEMORY_SWAPPINESS:-$automatic_swappiness}
swappiness_policy=automatic
if [[ -n ${HOST_MEMORY_SWAPPINESS+x} ]]; then
  swappiness_policy=override
fi
if [[ ! $desired_swappiness =~ ^[0-9]+$ ]] ||
  ((desired_swappiness < 0 || desired_swappiness > 200)); then
  echo "HOST_MEMORY_SWAPPINESS must be an integer from 0 to 200" >&2
  exit 2
fi

minimum_swap_bytes=${HOST_MEMORY_MIN_SWAP_BYTES:-$recommended_swap_bytes}
if [[ ! $minimum_swap_bytes =~ ^[0-9]+$ ]]; then
  echo "HOST_MEMORY_MIN_SWAP_BYTES must be a non-negative integer" >&2
  exit 2
fi
# swapon reports usable space, which is the file size minus its header page, so
# an exactly-sized 2 GiB swap file always measures a few KiB short. Without this
# slack every run would emit a warning nobody can act on.
swap_size_tolerance_bytes=$((4 * 1024 * 1024))

sysctl_bin=${HOST_MEMORY_SYSCTL_BIN:-$(command -v sysctl || true)}
[[ -n $sysctl_bin && -x $sysctl_bin ]] || {
  echo "sysctl executable not found" >&2
  exit 1
}

sysctl_config=${HOST_MEMORY_SYSCTL_CONFIG:-/etc/sysctl.d/99-zz-${service_name}-memory.conf}

if ((swap_bytes == 0 && minimum_swap_bytes > 0)); then
  echo "warning: no active swap; MemoryHigh has been narrowed to ${high_percent}% of MemoryMax so the services fail fast instead of stalling. Run ensure-swapfile.sh to provision ${recommended_swap_mib} MiB." >&2
elif ((swap_bytes + swap_size_tolerance_bytes < minimum_swap_bytes)); then
  echo "warning: active swap is below the recommended minimum: ${swap_bytes} < ${minimum_swap_bytes} bytes" >&2
fi

current_swappiness=$("$sysctl_bin" -n vm.swappiness | tr -d '[:space:]')
if [[ ! $current_swappiness =~ ^[0-9]+$ ]]; then
  echo "failed to read vm.swappiness" >&2
  exit 1
fi

active_changed=false
if [[ $current_swappiness != "$desired_swappiness" ]]; then
  "$sysctl_bin" -w "vm.swappiness=$desired_swappiness" >/dev/null
  active_changed=true
fi

effective_swappiness=$("$sysctl_bin" -n vm.swappiness | tr -d '[:space:]')
if [[ $effective_swappiness != "$desired_swappiness" ]]; then
  echo "failed to set vm.swappiness=$desired_swappiness (effective=$effective_swappiness)" >&2
  exit 1
fi

install -d -m 0755 "$(dirname "$sysctl_config")"
temporary_config=$(mktemp "${TMPDIR:-/tmp}/${service_name}-memory.XXXXXX")
trap 'rm -f "$temporary_config"' EXIT
printf '# Managed by %s host startup (RAM class: %s; policy: %s).\nvm.swappiness = %s\n' \
  "$service_name" "$memory_class" "$swappiness_policy" \
  "$desired_swappiness" >"$temporary_config"

persistent_changed=false
if [[ ! -f $sysctl_config ]] || ! cmp -s "$temporary_config" "$sysctl_config"; then
  install -m 0644 "$temporary_config" "$sysctl_config"
  persistent_changed=true
fi

systemd_dir=${HOST_MEMORY_SYSTEMD_DIR:-/etc/systemd/system}
limits_changed=false
write_memory_dropin() {
  local unit=$1
  local high_mib=$2
  local max_mib=$3
  local swap_max_mib=$4
  local dropin_dir="$systemd_dir/${service_name}-${unit}.service.d"
  local dropin="$dropin_dir/50-memory.conf"
  local temporary_dropin
  install -d -m 0755 "$dropin_dir"
  temporary_dropin=$(mktemp "${TMPDIR:-/tmp}/${service_name}-${unit}-memory.XXXXXX")
  printf '# Managed by %s host memory profile (%s, %s MiB RAM, %s CPUs, swap %s).\n[Service]\nMemoryHigh=%sM\nMemoryMax=%sM\nMemorySwapMax=%sM\n' \
    "$service_name" "$memory_class" "$memory_total_mib" "$cpu_count" \
    "$swap_state" "$high_mib" "$max_mib" "$swap_max_mib" >"$temporary_dropin"
  if [[ ! -f $dropin ]] || ! cmp -s "$temporary_dropin" "$dropin"; then
    install -m 0644 "$temporary_dropin" "$dropin"
    limits_changed=true
  fi
  rm -f "$temporary_dropin"
}

write_memory_dropin api "$api_high_mib" "$api_max_mib" "$api_swap_max_mib"
write_memory_dropin api-candidate "$api_high_mib" "$api_max_mib" "$api_swap_max_mib"
# Compatibility limit for the one-time migration from the persistent peer.
write_memory_dropin api-peer "$api_high_mib" "$api_max_mib" "$api_swap_max_mib"
write_memory_dropin worker "$worker_high_mib" "$worker_max_mib" "$worker_swap_max_mib"
write_memory_dropin dailyjob "$dailyjob_high_mib" "$dailyjob_max_mib" "$dailyjob_swap_max_mib"

printf 'host-memory service=%s memory_total_bytes=%s memory_class=%s cpu_count=%s swap_bytes=%s swap_state=%s minimum_swap_bytes=%s reserve_mib=%s allocatable_mib=%s high_percent=%s projected_peak_mib=%s swappiness=%s swappiness_policy=%s active_changed=%s persistent_changed=%s limits_changed=%s api_high_mib=%s api_max_mib=%s api_swap_max_mib=%s worker_high_mib=%s worker_max_mib=%s worker_swap_max_mib=%s dailyjob_high_mib=%s dailyjob_max_mib=%s dailyjob_swap_max_mib=%s config=%s systemd_dir=%s\n' \
  "$service_name" "$memory_total_bytes" "$memory_class" "$cpu_count" "$swap_bytes" \
  "$swap_state" "$minimum_swap_bytes" "$reserve_mib" "$allocatable_mib" \
  "$high_percent" "$projected_peak_mib" \
  "$effective_swappiness" "$swappiness_policy" \
  "$active_changed" "$persistent_changed" "$limits_changed" \
  "$api_high_mib" "$api_max_mib" "$api_swap_max_mib" \
  "$worker_high_mib" "$worker_max_mib" "$worker_swap_max_mib" \
  "$dailyjob_high_mib" "$dailyjob_max_mib" "$dailyjob_swap_max_mib" \
  "$sysctl_config" "$systemd_dir"