#!/usr/bin/env bash
# =============================================================================
# onx-cgroup-apply — Apply cgroups v2 per-user resource limits
#
# Purpose:
#   Creates/updates a systemd user-slice drop-in that enforces CPU, memory,
#   IO and process limits for a hosting account via cgroups v2.
#
#   Writes:
#     /etc/systemd/system/user-${UID}.slice.d/onox-limits.conf
#   Then reloads systemd and applies the properties at runtime.
#
# Input (stdin JSON):
#   {
#     "username":              "onx_xxxx",
#     "cpu_limit_percent":     80,      -- % of one CPU core (100 = 1 full core)
#     "memory_limit_mb":       1024,    -- hard memory cap in MiB
#     "io_read_mbps":          50,      -- IO read bandwidth cap in MiB/s
#     "io_write_mbps":         30,      -- IO write bandwidth cap in MiB/s
#     "entry_processes_limit": 20,      -- max concurrent entry processes
#     "nproc_limit":           100      -- total task/thread limit
#   }
#
# Output (stdout JSON):
#   {
#     "applied": true,
#     "username": "onx_xxxx",
#     "uid": 1001,
#     "cpu_limit_percent": 80,
#     "memory_limit_mb": 1024,
#     "io_read_mbps": 50,
#     "io_write_mbps": 30,
#     "nproc_limit": 100,
#     "effective_immediately": true,
#     "drop_in": "/etc/systemd/system/user-1001.slice.d/onox-limits.conf"
#   }
#
# Exit codes: 0=ok 1=invalid-input 2=preflight-fail 3=exec-fail
#
# Sudoers entry needed:
#   apache ALL=(root) NOPASSWD: /usr/local/onoxsoft/bin/onx-cgroup-apply
#
# Deployed to: /usr/local/onoxsoft/bin/onx-cgroup-apply
# =============================================================================

set -euo pipefail

SCRIPT_DIR="$(dirname "$(readlink -f "$0")")"
# shellcheck source=_lib/common.sh
source "${SCRIPT_DIR}/_lib/common.sh"

# ── Dependencies ──────────────────────────────────────────────────────────────
command -v jq         >/dev/null 2>&1 || { printf '{"error":"jq required"}\n' >&2; exit 2; }
command -v systemctl  >/dev/null 2>&1 || { printf '{"error":"systemctl required (systemd system)"}\n' >&2; exit 2; }
command -v id         >/dev/null 2>&1 || { printf '{"error":"id required"}\n' >&2; exit 2; }
require_root

# ── Read & parse stdin ────────────────────────────────────────────────────────
INPUT=$(cat)
onx_require_json "${INPUT}"

USERNAME=$(onx_json_get "${INPUT}" "username")
CPU_PCT=$(onx_json_get  "${INPUT}" "cpu_limit_percent"     "100")
MEM_MB=$(onx_json_get   "${INPUT}" "memory_limit_mb"       "1024")
IO_READ=$(onx_json_get  "${INPUT}" "io_read_mbps"          "50")
IO_WRITE=$(onx_json_get "${INPUT}" "io_write_mbps"         "30")
NPROC=$(onx_json_get    "${INPUT}" "nproc_limit"           "100")

# ── Input validation ──────────────────────────────────────────────────────────
onx_validate_username "${USERNAME}"
[[ "${CPU_PCT}"  =~ ^[0-9]+$ ]] || onx_die 1 "cpu_limit_percent must be a positive integer"
[[ "${MEM_MB}"   =~ ^[0-9]+$ ]] || onx_die 1 "memory_limit_mb must be a positive integer"
[[ "${IO_READ}"  =~ ^[0-9]+$ ]] || onx_die 1 "io_read_mbps must be a positive integer"
[[ "${IO_WRITE}" =~ ^[0-9]+$ ]] || onx_die 1 "io_write_mbps must be a positive integer"
[[ "${NPROC}"    =~ ^[0-9]+$ ]] || onx_die 1 "nproc_limit must be a positive integer"
(( CPU_PCT >= 10 && CPU_PCT <= 10000 )) || onx_die 1 "cpu_limit_percent out of range [10-10000]"
(( MEM_MB >= 128 ))                      || onx_die 1 "memory_limit_mb must be >= 128"
(( NPROC >= 20 ))                        || onx_die 1 "nproc_limit must be >= 20"

# ── Preflight ─────────────────────────────────────────────────────────────────
id "${USERNAME}" &>/dev/null || onx_die 2 "Linux user does not exist: ${USERNAME}"

# Verify cgroups v2 is mounted
mountpoint -q /sys/fs/cgroup 2>/dev/null || onx_die 2 "/sys/fs/cgroup not mounted (cgroups required)"
if [ -f /sys/fs/cgroup/cgroup.controllers ]; then
    grep -q "cpu memory" /sys/fs/cgroup/cgroup.controllers 2>/dev/null \
        || onx_log "warning: cpu/memory controllers may not be enabled in cgroups v2"
fi

# ── Resolve UID ───────────────────────────────────────────────────────────────
UID_VAL=$(id -u "${USERNAME}")

# ── Calculate derived values ──────────────────────────────────────────────────
# CPUQuota must be integer % string e.g. "80%"
CPU_QUOTA="${CPU_PCT}%"

# MemoryHigh = 90% of MemoryMax (soft warning before OOM)
MEM_HIGH_MB=$(( MEM_MB * 90 / 100 ))
(( MEM_HIGH_MB < 128 )) && MEM_HIGH_MB=128

IO_READ_BYTES=$(( IO_READ * 1024 * 1024 ))
IO_WRITE_BYTES=$(( IO_WRITE * 1024 * 1024 ))

# Determine block device for IO throttling (use /home's device, or first non-loop device)
HOME_DEV=""
if mountpoint -q /home 2>/dev/null; then
    HOME_DEV=$(stat -c '%d' /home 2>/dev/null || true)
    if [[ -n "${HOME_DEV}" ]]; then
        MAJOR=$(( HOME_DEV >> 8 ))
        MINOR=$(( HOME_DEV & 0xff ))
        HOME_DEV_STR="${MAJOR}:${MINOR}"
    fi
fi
# Fallback: skip IO throttling if device not determinable
IO_READ_CONF=""
IO_WRITE_CONF=""
if [[ -n "${HOME_DEV_STR:-}" ]]; then
    IO_READ_CONF="IOReadBandwidthMax=${HOME_DEV_STR} ${IO_READ_BYTES}"
    IO_WRITE_CONF="IOWriteBandwidthMax=${HOME_DEV_STR} ${IO_WRITE_BYTES}"
fi

# ── Create drop-in directory and config ───────────────────────────────────────
DROPIN_DIR="/etc/systemd/system/user-${UID_VAL}.slice.d"
DROPIN_FILE="${DROPIN_DIR}/onox-limits.conf"

mkdir -p "${DROPIN_DIR}"

cat > "${DROPIN_FILE}" <<DROPIN
# Generated by onox-cgroup-apply — do not edit manually
# Account: ${USERNAME}  UID: ${UID_VAL}
[Slice]
CPUQuota=${CPU_QUOTA}
MemoryMax=${MEM_MB}M
MemoryHigh=${MEM_HIGH_MB}M
TasksMax=${NPROC}
${IO_READ_CONF:+${IO_READ_CONF}}
${IO_WRITE_CONF:+${IO_WRITE_CONF}}
DROPIN

# ── Reload systemd + apply properties at runtime ──────────────────────────────
systemctl daemon-reload || onx_die 3 "systemctl daemon-reload failed"

# Apply runtime (best-effort; slice may not be active yet for idle accounts)
EFFECTIVE_NOW="false"
if systemctl is-active --quiet "user-${UID_VAL}.slice" 2>/dev/null; then
    systemctl set-property "user-${UID_VAL}.slice" \
        "CPUQuota=${CPU_QUOTA}" \
        "MemoryMax=${MEM_MB}M" \
        "MemoryHigh=${MEM_HIGH_MB}M" \
        "TasksMax=${NPROC}" 2>/dev/null || true
    EFFECTIVE_NOW="true"
fi

onx_log "cgroup-apply: ${USERNAME} (uid=${UID_VAL}) cpu=${CPU_QUOTA} mem=${MEM_MB}M nproc=${NPROC}"

# ── Output ────────────────────────────────────────────────────────────────────
printf '{"applied":true,"username":"%s","uid":%s,"cpu_limit_percent":%s,"memory_limit_mb":%s,"io_read_mbps":%s,"io_write_mbps":%s,"nproc_limit":%s,"effective_immediately":%s,"drop_in":"%s"}\n' \
    "${USERNAME}" "${UID_VAL}" "${CPU_PCT}" "${MEM_MB}" "${IO_READ}" "${IO_WRITE}" \
    "${NPROC}" "${EFFECTIVE_NOW}" "${DROPIN_FILE}"
