#!/bin/zsh
# Keep this Mac awake — even with the lid closed — for a fixed number of minutes.
#
# Wraps `pmset -a disablesleep`, the only setting that survives a closed lid
# (caffeinate does not). That flag is persistent across reboots, so the script
# always clears it on the way out: normal finish, Ctrl-C, SIGTERM, SIGHUP, or
# logout/restart/shutdown. For the cases no trap can catch — kernel panic, held
# power button, kill -9 — a LaunchDaemon clears the flag at every boot. The
# first run on a new Mac installs that guard automatically (Dropbox syncs the
# script, not the daemon), so there's nothing to remember on a fresh machine.
#
# Usage:
#   keep-awake [minutes]         # default 60
#   keep-awake --install-guard   # (re)install the boot guard explicitly
#   keep-awake --uninstall-guard # remove it and stop auto-installing
set -e

GREEN='\033[1;32m'
YELLOW='\033[1;33m'
RED='\033[1;31m'
DIM='\033[2m'
RESET='\033[0m'

SELF="$(basename "$0")"
GUARD_LABEL="com.hahnfamily.keep-awake-guard"
GUARD_PLIST="/Library/LaunchDaemons/${GUARD_LABEL}.plist"
# Machine-local, deliberately NOT in Dropbox: opting out on one Mac must not
# disable the guard on every other Mac the script syncs to.
GUARD_OPTOUT="${XDG_STATE_HOME:-$HOME/.local/state}/keep-awake/no-guard"

usage() {
  echo "Usage: $SELF [minutes]            # default 60"
  echo "       $SELF --install-guard      # (re)install the boot guard"
  echo "       $SELF --uninstall-guard    # remove it, stop auto-installing"
}

# Seconds as "5h 20m 30s" — leading unit unpadded, the ones after it padded, so
# a ticking countdown keeps a steady width instead of jittering as each unit
# rolls over. Units above the largest non-zero one are dropped ("20m 30s",
# "30s"). Pass a second argument to stop at minutes ("5h 20m", "20m"), for the
# once-a-minute line printed when stdout isn't a terminal.
fmt_duration() {
  local total=$1 coarse="${2:-}"
  local h=$(( total / 3600 )) m=$(( total % 3600 / 60 )) s=$(( total % 60 ))
  if (( h )); then
    printf '%dh %02dm' $h $m
  elif (( m )) || [[ -n "$coarse" ]]; then
    printf '%dm' $m
  fi
  if [[ -z "$coarse" ]]; then
    if (( h || m )); then printf ' %02ds' $s; else printf '%ds' $s; fi
  fi
}

# True only when the plist exists AND launchd has it loaded, so a half-installed
# guard (plist copied in by hand, never bootstrapped) gets repaired, not skipped.
guard_installed() {
  [[ -f "$GUARD_PLIST" ]] && sudo launchctl print "system/${GUARD_LABEL}" >/dev/null 2>&1
}

# Every step checks its own status. `set -e` is suppressed for the whole
# function when it runs as an `if !` condition, so without these an early
# failure would fall through and print the success message anyway.
install_guard() {
  echo "${DIM}Installing boot guard…${RESET}"
  sudo tee "$GUARD_PLIST" >/dev/null <<PLIST || return 1
<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
<dict>
  <key>Label</key><string>${GUARD_LABEL}</string>
  <key>ProgramArguments</key>
  <array>
    <string>/usr/bin/pmset</string>
    <string>-a</string>
    <string>disablesleep</string>
    <string>0</string>
  </array>
  <key>RunAtLoad</key><true/>
</dict>
</plist>
PLIST
  sudo chown root:wheel "$GUARD_PLIST" || return 1
  sudo chmod 644 "$GUARD_PLIST" || return 1
  sudo launchctl bootout system "$GUARD_PLIST" 2>/dev/null || true
  sudo launchctl bootstrap system "$GUARD_PLIST" || return 1
  rm -f "$GUARD_OPTOUT"
  echo "${GREEN}Boot guard installed.${RESET} Sleep is re-enabled at every startup, so a"
  echo "panic or hard power-off can't leave the setting stuck on."
  echo "${DIM}Remove it with: $SELF --uninstall-guard${RESET}"
}

uninstall_guard() {
  sudo launchctl bootout system "$GUARD_PLIST" 2>/dev/null || true
  sudo rm -f "$GUARD_PLIST"
  mkdir -p "$(dirname "$GUARD_OPTOUT")"
  date > "$GUARD_OPTOUT"
  echo "${GREEN}Boot guard removed${RESET} — and it won't be auto-installed again"
  echo "on this Mac. Reinstate it with: $SELF --install-guard"
}

ARG="${1:-60}"
case "$ARG" in
  --install-guard)   install_guard;   exit 0 ;;
  --uninstall-guard) uninstall_guard; exit 0 ;;
  -h|--help)         usage;           exit 0 ;;
esac

if [[ ! "$ARG" =~ ^[0-9]+$ ]] || (( ARG < 1 )); then
  echo "${RED}Minutes must be a positive whole number.${RESET}"
  usage
  exit 1
fi
MINUTES=$ARG

if [[ "$(pmset -g | awk '/SleepDisabled/ { print $2 }')" == "1" ]]; then
  echo "${YELLOW}Note: sleep is already disabled — an earlier run may not have cleaned up."
  echo "This run will clear it when it finishes.${RESET}"
fi

# Prime the sudo credential cache, then keep it warm in the background. A long
# countdown outlives the default 5-minute timeout, and the restore at the end
# must not stop to ask for a password (it may be running under SIGTERM).
echo "${DIM}Requesting admin rights…${RESET}"
sudo -v

# First run on this Mac: install the boot guard now, while the sudo credentials
# are already in hand, so a new machine is protected without anyone remembering
# a flag. A failure here is not worth aborting the actual job over.
if [[ ! -f "$GUARD_OPTOUT" ]] && ! guard_installed; then
  echo "${YELLOW}No boot guard on this Mac yet — installing it once, now.${RESET}"
  if ! install_guard; then
    echo "${RED}Boot guard install failed.${RESET} Continuing anyway: this run still"
    echo "restores sleep on exit, just not after a panic or hard power-off."
  fi
fi

# Started BEFORE the traps below so this subshell can't inherit them and try to
# run the restore itself when we kill it.
# Waits in 1-second slices rather than one `sleep 50`: killing this subshell
# does not kill the `sleep` it is blocked in, and that orphan keeps stdout open
# — enough to stall anything reading this script's output through a pipe. Short
# slices also let it notice a vanished parent (kill -9) within a second instead
# of refreshing sudo for another minute.
(
  trap - EXIT INT TERM HUP
  while :; do
    for _ in {1..50}; do
      sleep 1
      kill -0 $$ 2>/dev/null || exit 0
    done
    sudo -n -v 2>/dev/null || exit 0
  done
) &
KEEPALIVE_PID=$!

sudo pmset -a disablesleep 1

RESTORED=0
restore() {
  if (( RESTORED )); then return 0; fi
  RESTORED=1
  if [[ -n "${KEEPALIVE_PID:-}" ]]; then
    kill "$KEEPALIVE_PID" 2>/dev/null || true
  fi
  sudo -n pmset -a disablesleep 0 2>/dev/null || sudo pmset -a disablesleep 0 || true
  if [[ -t 1 ]]; then printf '\r\033[K'; fi
  echo "${GREEN}Sleep re-enabled.${RESET} Closing the lid suspends as usual."
}

trap restore EXIT
trap 'exit 130' INT
trap 'exit 143' TERM
trap 'exit 129' HUP

END=$(( $(date +%s) + MINUTES * 60 ))
UNTIL="$(date -r "$END" '+%-I:%M %p')"

echo "${GREEN}Sleep disabled for $(fmt_duration $(( MINUTES * 60 )) coarse)${RESET} — the lid can stay closed until ${UNTIL}."
echo "${DIM}Ctrl-C to stop early and re-enable sleep immediately.${RESET}"

# Always tick once a second, even when only printing once a minute: zsh defers
# a signal trap until the current foreground command returns, and a `sleep 60`
# here would delay the restore by up to a minute — long enough for a shutdown
# to SIGKILL us after its SIGTERM and leave the flag stuck on.
LAST_MIN=-1
while :; do
  REMAINING=$(( END - $(date +%s) ))
  if (( REMAINING <= 0 )); then break; fi
  MIN=$(( (REMAINING + 59) / 60 ))
  if [[ -t 1 ]]; then
    printf '\r\033[K  %s remaining' "$(fmt_duration $REMAINING)"
  elif (( MIN != LAST_MIN )); then
    LAST_MIN=$MIN
    printf '  %s remaining\n' "$(fmt_duration $(( MIN * 60 )) coarse)"
  fi
  sleep 1
done

if [[ -t 1 ]]; then printf '\r\033[K'; fi
echo "${GREEN}Time's up.${RESET}"