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

# -----------------------------------------------------------------------------
# disable-lid-sleep-proxmox.sh
#
# Purpose:
#   Configure a Proxmox/Debian host (especially laptops) to ignore the lid switch
#   and prevent any suspend/hibernate when the lid is closed.
#
# What it does:
#   1) Configures systemd-logind to ignore lid close events
#   2) Masks sleep/suspend/hibernate targets (hard stop)
#   3) Disables console blanking via GRUB (consoleblank=0)
#   4) Installs a systemd oneshot service to disable setterm blanking/powersave
#
# Safe features:
#   - Makes backups before editing files
#   - Idempotent: re-running won't duplicate config
#   - Verifies systemd + grub tools exist
#
# Usage:
#   sudo ./disable-lid-sleep-proxmox.sh
#
# Optional:
#   sudo ./disable-lid-sleep-proxmox.sh --undo
# -----------------------------------------------------------------------------

SCRIPT_NAME="$(basename "$0")"
STAMP="$(date +%Y%m%d-%H%M%S)"

require_root() {
  if [[ "${EUID}" -ne 0 ]]; then
    echo "ERROR: Must be run as root. Use: sudo ./${SCRIPT_NAME}"
    exit 1
  fi
}

has_cmd() { command -v "$1" >/dev/null 2>&1; }

backup_file() {
  local f="$1"
  if [[ -f "$f" ]]; then
    cp -a "$f" "${f}.bak.${STAMP}"
    echo "Backup created: ${f}.bak.${STAMP}"
  fi
}

ensure_line_in_file() {
  local file="$1"
  local line="$2"

  mkdir -p "$(dirname "$file")"
  touch "$file"

  if ! grep -Fxq "$line" "$file"; then
    echo "$line" >> "$file"
    echo "Added to $file: $line"
  else
    echo "Already present in $file: $line"
  fi
}

set_or_replace_keyvalue() {
  # Sets key=value in a config file (logind.conf style), uncommenting if needed.
  local file="$1"
  local key="$2"
  local value="$3"

  mkdir -p "$(dirname "$file")"
  touch "$file"

  if grep -Eq "^[#[:space:]]*${key}=" "$file"; then
    sed -i -E "s|^[#[:space:]]*(${key}=).*|\1${value}|" "$file"
    echo "Set $key=$value in $file"
  else
    echo "${key}=${value}" >> "$file"
    echo "Appended $key=$value to $file"
  fi
}

mask_sleep_targets() {
  local targets=(
    sleep.target
    suspend.target
    hibernate.target
    hybrid-sleep.target
  )

  for t in "${targets[@]}"; do
    systemctl mask "$t" >/dev/null 2>&1 || true
    echo "Masked: $t"
  done
}

unmask_sleep_targets() {
  local targets=(
    sleep.target
    suspend.target
    hibernate.target
    hybrid-sleep.target
  )

  for t in "${targets[@]}"; do
    systemctl unmask "$t" >/dev/null 2>&1 || true
    echo "Unmasked: $t"
  done
}

install_dpms_service() {
  local svc="/etc/systemd/system/disable-dpms.service"

  backup_file "$svc"

  cat > "$svc" <<'EOF'
[Unit]
Description=Disable monitor sleep / console blanking on boot

[Service]
Type=oneshot
ExecStart=/usr/bin/setterm -blank 0 -powersave off -powerdown 0

[Install]
WantedBy=multi-user.target
EOF

  systemctl daemon-reload
  systemctl enable --now disable-dpms.service
  echo "Installed & enabled: disable-dpms.service"
}

remove_dpms_service() {
  local svc="/etc/systemd/system/disable-dpms.service"
  if [[ -f "$svc" ]]; then
    systemctl disable --now disable-dpms.service >/dev/null 2>&1 || true
    rm -f "$svc"
    systemctl daemon-reload
    echo "Removed: disable-dpms.service"
  else
    echo "disable-dpms.service not found; skipping."
  fi
}

configure_logind() {
  local f="/etc/systemd/logind.conf"
  backup_file "$f"

  set_or_replace_keyvalue "$f" "HandleLidSwitch" "ignore"
  set_or_replace_keyvalue "$f" "HandleLidSwitchExternalPower" "ignore"
  set_or_replace_keyvalue "$f" "HandleLidSwitchDocked" "ignore"
  set_or_replace_keyvalue "$f" "LidSwitchIgnoreInhibited" "no"

  systemctl restart systemd-logind
  echo "Restarted: systemd-logind"
}

undo_logind() {
  # We don't "restore defaults" perfectly (depends on distro), but we do remove
  # the explicit settings by commenting them out if present.
  local f="/etc/systemd/logind.conf"
  if [[ -f "$f" ]]; then
    backup_file "$f"
    sed -i -E \
      -e 's|^(HandleLidSwitch=)|#\1|' \
      -e 's|^(HandleLidSwitchExternalPower=)|#\1|' \
      -e 's|^(HandleLidSwitchDocked=)|#\1|' \
      -e 's|^(LidSwitchIgnoreInhibited=)|#\1|' \
      "$f"
    systemctl restart systemd-logind
    echo "Commented out lid settings in $f and restarted logind."
  else
    echo "$f not found; skipping."
  fi
}

configure_grub_consoleblank() {
  if ! has_cmd update-grub && ! has_cmd grub-mkconfig; then
    echo "WARNING: Neither update-grub nor grub-mkconfig found; skipping GRUB step."
    return 0
  fi

  local f="/etc/default/grub.d/nolidsleep.cfg"
  mkdir -p /etc/default/grub.d
  backup_file "$f"

  # Ensure the commandline includes consoleblank=0 without duplicating.
  if [[ -f "$f" ]] && grep -q 'consoleblank=0' "$f"; then
    echo "GRUB consoleblank already configured in $f"
  else
    cat > "$f" <<'EOF'
# Keep the local console from blanking (useful on laptop Proxmox installs)
GRUB_CMDLINE_LINUX_DEFAULT="$GRUB_CMDLINE_LINUX_DEFAULT consoleblank=0"
EOF
    echo "Wrote GRUB config snippet: $f"
  fi

  if has_cmd update-grub; then
    update-grub
    echo "Ran: update-grub"
  else
    # Fallback (uncommon on Debian/Proxmox, but harmless)
    grub-mkconfig -o /boot/grub/grub.cfg
    echo "Ran: grub-mkconfig -o /boot/grub/grub.cfg"
  fi
}

undo_grub_consoleblank() {
  local f="/etc/default/grub.d/nolidsleep.cfg"
  if [[ -f "$f" ]]; then
    backup_file "$f"
    rm -f "$f"
    echo "Removed: $f"

    if has_cmd update-grub; then
      update-grub
      echo "Ran: update-grub"
    elif has_cmd grub-mkconfig; then
      grub-mkconfig -o /boot/grub/grub.cfg
      echo "Ran: grub-mkconfig -o /boot/grub/grub.cfg"
    else
      echo "WARNING: No grub update tool found; GRUB not regenerated."
    fi
  else
    echo "$f not found; skipping."
  fi
}

show_status() {
  echo
  echo "---- STATUS CHECKS ----"
  echo "logind.conf lid settings:"
  grep -E '^(#\s*)?(HandleLidSwitch|HandleLidSwitchExternalPower|HandleLidSwitchDocked|LidSwitchIgnoreInhibited)=' \
    /etc/systemd/logind.conf 2>/dev/null || echo "(no matching lines found)"

  echo
  echo "Masked sleep targets:"
  for t in sleep.target suspend.target hibernate.target hybrid-sleep.target; do
    systemctl is-enabled "$t" 2>/dev/null || true
  done

  echo
  echo "DPMS service:"
  systemctl status disable-dpms.service --no-pager -l 2>/dev/null || echo "(service not installed)"

  echo
  echo "GRUB snippet:"
  if [[ -f /etc/default/grub.d/nolidsleep.cfg ]]; then
    cat /etc/default/grub.d/nolidsleep.cfg
  else
    echo "(not present)"
  fi
  echo "-----------------------"
  echo
}

main() {
  require_root

  local mode="apply"
  if [[ "${1:-}" == "--undo" ]]; then
    mode="undo"
  fi

  if [[ "$mode" == "apply" ]]; then
    echo "Applying lid-close ignore + anti-sleep configuration..."
    configure_logind
    mask_sleep_targets
    configure_grub_consoleblank
    install_dpms_service
    show_status
    echo "DONE. Reboot recommended to ensure GRUB consoleblank=0 takes effect:"
    echo "  reboot"
  else
    echo "Undoing configuration..."
    undo_logind
    unmask_sleep_targets
    undo_grub_consoleblank
    remove_dpms_service
    show_status
    echo "DONE. Reboot recommended:"
    echo "  reboot"
  fi
}

main "$@"