#!/usr/bin/env bash
# /usr/sbin/secubox-tuning-apply — apply swap + cgroup MemoryHigh tuning.
#
# Idempotent. Reads config from /etc/secubox/tuning/.
#
# Usage:
#   secubox-tuning-apply               # apply both swap and cgroup caps
#   secubox-tuning-apply swap          # swap only
#   secubox-tuning-apply cgroup        # cgroup caps only
#   secubox-tuning-apply --dry-run     # show what would be done

set -euo pipefail

CONFIG_FILE="/etc/secubox/tuning/config.toml"
LXC_CAPS_FILE="/etc/secubox/tuning/lxc-memory-high.toml"

CYAN='\033[0;36m'
GREEN='\033[0;32m'
RED='\033[0;31m'
YELLOW='\033[1;33m'
NC='\033[0m'

log()  { echo -e "${CYAN}[tuning]${NC} $*"; }
ok()   { echo -e "${GREEN}[  ok  ]${NC} $*"; }
warn() { echo -e "${YELLOW}[ warn ]${NC} $*"; }
err()  { echo -e "${RED}[fail  ]${NC} $*" >&2; }

DRY_RUN=0
DO_SWAP=1
DO_CGROUP=1

for arg in "$@"; do
  case "$arg" in
    --dry-run) DRY_RUN=1 ;;
    swap)      DO_CGROUP=0 ;;
    cgroup)    DO_SWAP=0 ;;
    -h|--help)
      sed -n '1,/^set -/p' "$0" | grep '^#' | sed 's/^# \?//'
      exit 0
      ;;
    *) err "unknown arg: $arg"; exit 2 ;;
  esac
done

run() {
  if [[ $DRY_RUN -eq 1 ]]; then
    log "DRY-RUN: $*"
  else
    eval "$@"
  fi
}

# Minimal TOML reader using python3 (avoids depending on toml-cli).
toml_get() {
  local file="$1" key="$2" default="${3:-}"
  python3 -c "
import sys
try:
    import tomllib
except ImportError:
    import tomli as tomllib
try:
    d = tomllib.loads(open('$file').read())
    for k in '$key'.split('.'):
        d = d[k]
    print(d)
except (KeyError, FileNotFoundError):
    print('$default')
" 2>/dev/null || echo "$default"
}

toml_table_keys() {
  local file="$1" table="$2"
  python3 -c "
import sys
try:
    import tomllib
except ImportError:
    import tomli as tomllib
d = tomllib.loads(open('$file').read())
for k in d.get('$table', {}):
    print(k)
" 2>/dev/null
}

apply_swap() {
  log "swap phase"
  local swap_path size min_free_gb
  swap_path=$(toml_get "$CONFIG_FILE" swap.path /data/swapfile)
  size=$(toml_get "$CONFIG_FILE" swap.size 4G)
  min_free_gb=$(toml_get "$CONFIG_FILE" swap.minimum_data_free_gb 8)

  if swapon --show=NAME --noheadings | grep -q '.'; then
    ok "swap already active: $(swapon --show=NAME --noheadings | tr '\n' ' ')"
    return 0
  fi

  if [[ ! -d "$(dirname "$swap_path")" ]]; then
    warn "swap parent dir $(dirname "$swap_path") missing, skipping"
    return 0
  fi

  local free_gb
  free_gb=$(df --output=avail -BG "$(dirname "$swap_path")" | tail -1 | tr -d 'G ')
  if [[ "$free_gb" -lt "$min_free_gb" ]]; then
    warn "/data free space ${free_gb}G < required ${min_free_gb}G; skipping swap creation"
    return 0
  fi

  if [[ -f "$swap_path" ]]; then
    log "$swap_path exists, activating"
    run "swapon $swap_path"
  else
    log "creating $swap_path ($size)"
    run "fallocate -l $size $swap_path"
    run "chmod 600 $swap_path"
    run "mkswap $swap_path"
    run "swapon $swap_path"
  fi

  if ! grep -q "^$swap_path " /etc/fstab; then
    log "persisting in /etc/fstab"
    run "echo '$swap_path none swap sw 0 0' >> /etc/fstab"
  fi
  ok "swap: $(swapon --show)"
}

apply_cgroup() {
  log "cgroup MemoryHigh phase"
  if [[ ! -f "$LXC_CAPS_FILE" ]]; then
    warn "$LXC_CAPS_FILE not found, skipping"
    return 0
  fi

  local containers
  containers=$(toml_table_keys "$LXC_CAPS_FILE" caps)
  if [[ -z "$containers" ]]; then
    warn "no caps defined in $LXC_CAPS_FILE"
    return 0
  fi

  # On SecuBox boards LXC containers are started by lxc-autostart at boot,
  # NOT by systemctl start lxc@<name>.service. So drop-ins under
  # /etc/systemd/system/lxc@<name>.service.d/ are ineffective. Apply the
  # cap two ways instead:
  #   1. Direct write to the container's existing cgroup
  #      (/sys/fs/cgroup/lxc.payload.<name>/memory.high) for immediate
  #      effect, no restart.
  #   2. Add lxc.cgroup2.memory.high=<cap> to the container's LXC config
  #      (/var/lib/lxc/<name>/config OR /data/lxc/<name>/config — try
  #      both, take the first that exists) so the cap is reapplied on
  #      every lxc-start.

  local applied=0 persisted=0 missed=0
  for ctn in $containers; do
    local cap cgrp lxc_conf
    cap=$(toml_get "$LXC_CAPS_FILE" "caps.$ctn")

    # (1) immediate cgroup write
    cgrp="/sys/fs/cgroup/lxc.payload.${ctn}"
    if [[ -d "$cgrp" ]]; then
      run "echo $cap > $cgrp/memory.high"
      applied=$((applied+1))
      printf "  %-12s cgroup memory.high=%s (runtime)" "$ctn" "$cap"
    else
      missed=$((missed+1))
      printf "  %-12s container not running (no cgroup)" "$ctn"
    fi

    # (2) persistent LXC config edit. Look in both standard and /data paths.
    lxc_conf=""
    for path in "/data/lxc/${ctn}/config" "/var/lib/lxc/${ctn}/config"; do
      if [[ -f "$path" ]]; then
        lxc_conf="$path"
        break
      fi
    done
    if [[ -n "$lxc_conf" ]]; then
      if [[ $DRY_RUN -eq 1 ]]; then
        printf " + DRY-RUN: would set memory.high=%s in %s\n" "$cap" "$lxc_conf"
      else
        # Replace existing memory.high line if present, otherwise append.
        if grep -q "^lxc\.cgroup2\.memory\.high\s*=" "$lxc_conf"; then
          sed -i "s|^lxc\.cgroup2\.memory\.high\s*=.*|lxc.cgroup2.memory.high = ${cap}|" "$lxc_conf"
        else
          echo "lxc.cgroup2.memory.high = ${cap}" >> "$lxc_conf"
        fi
        printf " + persisted in %s\n" "$lxc_conf"
        persisted=$((persisted+1))
      fi
    else
      printf "\n"
    fi
  done

  ok "cgroup caps: $applied runtime + $persisted persisted + $missed missed"
}

log "secubox-tuning-apply (dry_run=$DRY_RUN, swap=$DO_SWAP, cgroup=$DO_CGROUP)"
[[ $DO_SWAP   -eq 1 ]] && apply_swap
[[ $DO_CGROUP -eq 1 ]] && apply_cgroup
ok "done"
