#!/bin/bash -p

# omarchy:summary=Manage sleep and idle inhibition during an update
# omarchy:args=<start|stop>
# omarchy:hidden=true

if [[ $- != *p* ]]; then
  echo "Refusing an unsafe Bash startup." >&2
  exit 126
fi

security_entrypoint=$(/usr/bin/readlink -e -- "${BASH_SOURCE[0]}") || exit 126
source "${security_entrypoint%/*}/omarchy-security-functions" || exit 126
omarchy_security_require_privileged_bash_startup || exit 126
set -e
omarchy_security_sanitize_bash_environment "$0" "$@"
omarchy_security_require_source_root "$0"
# omarchy update authorizes once and revokes when it is done. Revoking here
# would throw that away and prompt again for the rest of the update.
update_sudo_session=0
if [[ ${OMARCHY_UPDATE_SUDO_SESSION:-} == "1" ]]; then
  update_sudo_session=1
  omarchy_security_revoke_sudo_timestamp() { :; }
fi
# Traps first, so a signal or failure during the entry revocation still
# exits through the cleanup path.
omarchy_security_install_sudo_cleanup_traps
omarchy_security_revoke_sudo_timestamp || exit 1
omarchy_security_enable_no_update_sudo

stay_awake_state="$HOME/.local/state/omarchy/indicators/stay-awake"
caller_uid=""
state_base=""
state_dir=""
idle_owner_file=""
inhibit_pid_file=""
launch_control_file=""
launch_pending=0
launch_token=""

fail_state_boundary() {
  echo "Refusing to use an unsafe Omarchy update inhibitor state path." >&2
  return 1
}

directory_is_private() {
  local directory="$1"
  local expected_owner="$2"
  local legacy="${3:-0}"
  local canonical=""
  local owner=""
  local mode=""

  [[ -d $directory && ! -L $directory ]] || return 1
  canonical=$(readlink -e -- "$directory") || return 1
  [[ $canonical == "$directory" ]] || return 1
  read -r owner mode < <(stat -Lc '%u %a' -- "$directory") || return 1
  [[ $owner == "$expected_owner" ]] || return 1
  [[ $mode == "700" ]] || (( legacy == 1 && (8#$mode & 022) == 0 ))
}

root_owned_parent_chain() {
  local directory="$1"
  local parent owner mode type canonical

  parent=$(/usr/bin/dirname -- "$directory") || return 1
  while :; do
    [[ -d $parent && ! -L $parent ]] || return 1
    canonical=$(/usr/bin/readlink -e -- "$parent") || return 1
    [[ $canonical == "$parent" ]] || return 1
    read -r owner mode type < <(/usr/bin/stat -Lc '%u %a %F' -- "$parent") || return 1
    [[ $owner == 0 && $type == "directory" ]] || return 1
    ! ((8#$mode & 022)) || return 1
    [[ $parent == / ]] && break
    parent=$(/usr/bin/dirname -- "$parent") || return 1
  done
}

runtime_directory_is_private() {
  local directory="$1" expected_owner="$2"
  directory_is_private "$directory" "$expected_owner" &&
    root_owned_parent_chain "$directory"
}

ensure_private_directory() {
  local directory="$1"

  if [[ ! -e $directory && ! -L $directory ]]; then
    mkdir -m 700 -- "$directory" 2>/dev/null || true
  fi
  # Older helpers used mkdir -p with the caller's umask. Tighten only a
  # canonical, caller-owned directory that other accounts cannot write.
  directory_is_private "$directory" "$caller_uid" 1 || return 1
  chmod 700 -- "$directory" || return 1
  directory_is_private "$directory" "$caller_uid"
}

initialize_state_boundary() {
  local canonical_tmp=""
  local tmp_owner=""
  local tmp_mode=""

  caller_uid="$EUID"
  [[ $caller_uid =~ ^[0-9]+$ ]] || return 1

  if [[ -n ${XDG_RUNTIME_DIR:-} ]]; then
    runtime_directory_is_private "$XDG_RUNTIME_DIR" "$caller_uid" || fail_state_boundary
    state_base="$XDG_RUNTIME_DIR"
  else
    [[ -d /tmp && ! -L /tmp ]] || fail_state_boundary
    canonical_tmp=$(readlink -e -- /tmp) || fail_state_boundary
    read -r tmp_owner tmp_mode < <(stat -Lc '%u %a' -- /tmp) || fail_state_boundary
    [[ $canonical_tmp == "/tmp" && $tmp_owner == "0" && $tmp_mode == "1777" ]] || fail_state_boundary

    state_base="/tmp/omarchy-$caller_uid"
    ensure_private_directory "$state_base" || fail_state_boundary
  fi

  state_dir="$state_base/omarchy-update-stay-awake"
  idle_owner_file="$state_dir/idle-owner"
  inhibit_pid_file="$state_dir/inhibit-pid"
  launch_control_file="$state_dir/launch-control"
}

state_file_is_private() {
  local state_file="$1"
  local legacy="${2:-0}"
  local owner=""
  local mode=""
  local links=""

  [[ -f $state_file && ! -L $state_file ]] || return 1
  read -r owner mode links < <(stat -Lc '%u %a %h' -- "$state_file") || return 1
  [[ $owner == "$caller_uid" && $links == "1" ]] || return 1
  [[ $mode == "600" ]] || (( legacy == 1 && (8#$mode & 022) == 0 ))
}

read_state_record() {
  local state_file="$1"
  local legacy="${2:-0}"
  local records=()
  local file_size=""
  local LC_ALL=C

  state_file_is_private "$state_file" "$legacy" || return 1
  mapfile -t records <"$state_file" || return 1
  (( ${#records[@]} == 1 )) || return 1
  file_size=$(stat -Lc '%s' -- "$state_file") || return 1
  (( file_size == ${#records[0]} + 1 )) || return 1
  state_file_is_private "$state_file" "$legacy" || return 1
  printf '%s\n' "${records[0]}"
}

atomic_write_state() {
  local state_file="$1"
  local record="$2"
  local temporary=""

  [[ $record != *$'\n'* ]] || return 1
  temporary=$(mktemp "$state_dir/.${state_file##*/}.XXXXXXXX") || return 1
  chmod 600 "$temporary" || {
    rm -f -- "$temporary"
    return 1
  }
  if ! printf '%s\n' "$record" >"$temporary" || ! state_file_is_private "$temporary"; then
    rm -f -- "$temporary"
    return 1
  fi
  if [[ -e $state_file || -L $state_file ]]; then
    state_file_is_private "$state_file" || {
      rm -f -- "$temporary"
      return 1
    }
  fi
  mv -fT -- "$temporary" "$state_file" || {
    rm -f -- "$temporary"
    return 1
  }
  state_file_is_private "$state_file"
}

rollback_pending_launch() {
  local control_fd=""
  local inhibit_record=""
  local inhibit_pid=""
  local recorded_start_time=""
  local recorded_owner=""
  local token=""

  (( launch_pending == 1 )) || return 0

  if [[ -e $launch_control_file || -L $launch_control_file ]]; then
    state_file_is_private "$launch_control_file" || return 1
    exec {control_fd}<>"$launch_control_file" || return 1
    /usr/bin/flock -x "$control_fd" || {
      exec {control_fd}>&-
      return 1
    }
    : >"/proc/self/fd/$control_fd"
    printf 'cancelled %s\n' "$launch_token" >&"$control_fd"
    rm -f -- "$launch_control_file"
    /usr/bin/flock -u "$control_fd"
    exec {control_fd}>&-
  fi

  inhibit_record=$(read_state_record "$inhibit_pid_file" 2>/dev/null || true)
  if [[ $inhibit_record =~ ^1\ ([1-9][0-9]{0,18})\ ([1-9][0-9]{0,18})\ ([0-9]{1,10})\ ($launch_token)$ ]]; then
    inhibit_pid="${BASH_REMATCH[1]}"
    recorded_start_time="${BASH_REMATCH[2]}"
    recorded_owner="${BASH_REMATCH[3]}"
    token="${BASH_REMATCH[4]}"
    if process_matches "$inhibit_pid" "$recorded_start_time" "$recorded_owner" "$token"; then
      discard_launched_inhibitor "$inhibit_pid" "$recorded_start_time" "$recorded_owner" "$token"
    fi
    [[ $(read_state_record "$inhibit_pid_file" 2>/dev/null || true) != "$inhibit_record" ]] ||
      rm -f -- "$inhibit_pid_file"
  fi

  # A signal may interrupt either parent-side atomic write before its rename.
  rm -f -- "$state_dir"/.launch-control.* "$state_dir"/.idle-owner.*
  stop_locked || return 1
  launch_pending=0
}

cleanup_pending_launch() {
  local status=$?

  trap - EXIT HUP INT TERM
  if ! rollback_pending_launch; then
    echo "Failed to roll back the pending Omarchy update sleep inhibitor." >&2
    (( status != 0 )) || status=1
  fi
  omarchy_security_exit_with_revoked_sudo "$status" \
    "Failed to invalidate sudo credentials after the update sleep inhibitor."
}

process_start_time() {
  local process_pid="$1"
  local process_stat=""
  local stat_fields=()

  [[ -r /proc/$process_pid/stat ]] || return 1
  process_stat=$(</proc/"$process_pid"/stat)
  process_stat="${process_stat##*) }"
  read -r -a stat_fields <<<"$process_stat"
  (( ${#stat_fields[@]} > 19 )) || return 1
  [[ ${stat_fields[19]} =~ ^[0-9]+$ ]] || return 1
  printf '%s\n' "${stat_fields[19]}"
}

process_owner() {
  local process_pid="$1"
  local token="$2"
  local owner=""

  if [[ -n $token ]]; then
    owner=$(stat -Lc '%u' -- "/proc/$process_pid") || return 1
  else
    # Legacy state can name sudo: its real UID remains the invoking account,
    # even when /proc ownership reflects its effective root credentials.
    owner=$(awk '/^Uid:/ { print $2; exit }' "/proc/$process_pid/status") || return 1
  fi
  [[ $owner =~ ^[0-9]+$ ]] || return 1
  printf '%s\n' "$owner"
}

process_has_token() {
  local process_pid="$1"
  local token="$2"
  local argument=""
  local expected="--why=Omarchy update in progress [$token]"

  [[ -r /proc/$process_pid/cmdline ]] || return 1
  while IFS= read -r -d '' argument; do
    [[ $argument == "$expected" ]] && return 0
  done <"/proc/$process_pid/cmdline"
  return 1
}

process_identity() {
  local process_pid="$1"
  local token="$2"
  local start_before=""
  local start_after=""
  local owner_before=""
  local owner_after=""

  start_before=$(process_start_time "$process_pid") || return 1
  owner_before=$(process_owner "$process_pid" "$token") || return 1
  [[ -z $token ]] || process_has_token "$process_pid" "$token" || return 1
  start_after=$(process_start_time "$process_pid") || return 1
  owner_after=$(process_owner "$process_pid" "$token") || return 1
  [[ $start_before == "$start_after" && $owner_before == "$owner_after" ]] || return 1
  printf '%s %s\n' "$start_before" "$owner_before"
}

process_matches() {
  local process_pid="$1"
  local expected_start="$2"
  local expected_owner="$3"
  local token="$4"
  local identity=""

  identity=$(process_identity "$process_pid" "$token" 2>/dev/null) || return 1
  [[ $identity == "$expected_start $expected_owner" ]]
}

process_state() {
  local process_pid="$1"
  local process_stat=""

  [[ -r /proc/$process_pid/stat ]] || return 1
  process_stat=$(</proc/"$process_pid"/stat)
  process_stat="${process_stat##*) }"
  printf '%s\n' "${process_stat%% *}"
}

discard_launched_inhibitor() {
  local inhibit_pid="$1"
  local recorded_start_time="$2"
  local recorded_owner="$3"
  local token="$4"

  signal_inhibitor "$inhibit_pid" "$recorded_start_time" "$recorded_owner" "$token" TERM || true
  for (( attempt = 0; attempt < 25; attempt++ )); do
    process_matches "$inhibit_pid" "$recorded_start_time" "$recorded_owner" "$token" || break
    [[ $(process_state "$inhibit_pid" 2>/dev/null || true) == "Z" ]] && break
    sleep 0.02
  done
  if process_matches "$inhibit_pid" "$recorded_start_time" "$recorded_owner" "$token" &&
    [[ $(process_state "$inhibit_pid" 2>/dev/null || true) != "Z" ]]; then
    signal_inhibitor "$inhibit_pid" "$recorded_start_time" "$recorded_owner" "$token" KILL || true
  fi

  if [[ ! -e /proc/$inhibit_pid ]] || [[ $(process_state "$inhibit_pid" 2>/dev/null || true) == "Z" ]]; then
    wait "$inhibit_pid" 2>/dev/null || true
  fi
}

signal_inhibitor() {
  local inhibit_pid="$1"
  local recorded_start_time="$2"
  local recorded_owner="$3"
  local token="$4"
  local signal="$5"

  [[ $recorded_owner == "$caller_uid" ]] || return 1
  [[ $signal == "TERM" || $signal == "KILL" ]] || return 1
  process_matches "$inhibit_pid" "$recorded_start_time" "$recorded_owner" "$token" || return 1
  builtin kill -s "$signal" -- "$inhibit_pid" 2>/dev/null
}

terminate_inhibitor() {
  local inhibit_pid="$1"
  local recorded_start_time="$2"
  local recorded_owner="$3"
  local token="$4"

  signal_inhibitor "$inhibit_pid" "$recorded_start_time" "$recorded_owner" "$token" TERM || true

  for (( attempt = 0; attempt < 50; attempt++ )); do
    process_matches "$inhibit_pid" "$recorded_start_time" "$recorded_owner" "$token" || return 0
    [[ $(process_state "$inhibit_pid" 2>/dev/null || true) != "Z" ]] || return 0
    sleep 0.02
  done

  process_matches "$inhibit_pid" "$recorded_start_time" "$recorded_owner" "$token" || return 0
  [[ $(process_state "$inhibit_pid" 2>/dev/null || true) == "Z" ]] && return 0
  return 1
}

stop_locked() {
  local inhibit_record=""
  local inhibit_pid=""
  local recorded_start_time=""
  local recorded_owner=""
  local token=""
  local idle_owner=""
  local current_idle_owner=""
  local current_start=""
  local failed=0
  local remove_inhibit_state=1

  if [[ ! -e $state_dir && ! -L $state_dir ]]; then
    return 0
  fi
  if ! ensure_private_directory "$state_dir"; then
    fail_state_boundary
    return 1
  fi

  if [[ -e $idle_owner_file || -L $idle_owner_file ]]; then
    idle_owner=$(read_state_record "$idle_owner_file" 1 2>/dev/null || true)
    if [[ $idle_owner =~ ^[0-9]+:[0-9]+:[0-9]+$ ]]; then
      if [[ -f $stay_awake_state ]]; then
        current_idle_owner=$(<"$stay_awake_state")
      fi
      if [[ $current_idle_owner == "$idle_owner" ]]; then
        omarchy-toggle-idle allow-idle >/dev/null 2>&1 || true
      fi
    else
      echo "Ignoring unsafe Omarchy update idle ownership state." >&2
      failed=1
    fi
    rm -f -- "$idle_owner_file"
  fi

  if [[ -e $inhibit_pid_file || -L $inhibit_pid_file ]]; then
    inhibit_record=$(read_state_record "$inhibit_pid_file" 1 2>/dev/null || true)
    if state_file_is_private "$inhibit_pid_file" &&
      [[ $inhibit_record =~ ^1\ ([1-9][0-9]{0,18})\ ([1-9][0-9]{0,18})\ ([0-9]{1,10})\ ([0-9a-f]{32})$ ]]; then
      inhibit_pid="${BASH_REMATCH[1]}"
      recorded_start_time="${BASH_REMATCH[2]}"
      recorded_owner="${BASH_REMATCH[3]}"
      token="${BASH_REMATCH[4]}"
    elif [[ $inhibit_record =~ ^([1-9][0-9]{0,18})\ ([1-9][0-9]{0,18})$ ]]; then
      # The first update starts the old helper and installs this one before
      # cleanup. Its protected same-account record predates launch tokens.
      inhibit_pid="${BASH_REMATCH[1]}"
      recorded_start_time="${BASH_REMATCH[2]}"
      recorded_owner="$caller_uid"
      if [[ -e /proc/$inhibit_pid && $(process_state "$inhibit_pid" 2>/dev/null || true) != "Z" ]] &&
        ! process_matches "$inhibit_pid" "$recorded_start_time" "$recorded_owner" ""; then
        current_start=$(process_start_time "$inhibit_pid" 2>/dev/null || true)
        if [[ -z $current_start || $current_start == "$recorded_start_time" ]]; then
          echo "Cannot verify permission to stop the legacy update inhibitor; retained its state for recovery." >&2
          failed=1
          remove_inhibit_state=0
        fi
      fi
    else
      echo "Ignoring unsafe Omarchy update sleep inhibitor state." >&2
      failed=1
    fi
    if [[ -n $inhibit_pid ]] && (( remove_inhibit_state == 1 )); then
      if process_matches "$inhibit_pid" "$recorded_start_time" "$recorded_owner" "$token"; then
        if ! terminate_inhibitor "$inhibit_pid" "$recorded_start_time" "$recorded_owner" "$token"; then
          echo "Failed to stop the Omarchy update sleep inhibitor." >&2
          failed=1
          remove_inhibit_state=0
        fi
      fi
    fi
    (( remove_inhibit_state == 0 )) || rm -f -- "$inhibit_pid_file"
  fi

  rmdir "$state_dir" 2>/dev/null || true
  (( failed == 0 ))
}

start_locked() {
  local idle_owner="$$:$RANDOM:$RANDOM"
  local token=""
  local launcher_pid=""
  local inhibit_record=""
  local inhibit_pid=""
  local inhibit_start_time=""
  local inhibit_owner=""
  local readiness_attempts=0

  stop_locked || return 1
  ensure_private_directory "$state_dir" || fail_state_boundary

  token=$(LC_ALL=C /usr/bin/od -An -N16 -tx1 /dev/urandom | /usr/bin/tr -d ' \n')
  [[ $token =~ ^[0-9a-f]{32}$ ]] || return 1
  launch_token="$token"
  launch_pending=1
  trap cleanup_pending_launch EXIT
  atomic_write_state "$launch_control_file" "active $token" || return 1

  # The child is identifiable by the token in its arguments from the moment it
  # starts. It stays a shell so stop's TERM ends it with a clean exit, since
  # systemd-inhibit reports a child killed by a signal as an error.
  local hold_command=(
    /usr/bin/systemd-inhibit --what=sleep:idle --who=omarchy-update
    --why="Omarchy update in progress [$token]" --mode=block
    /usr/bin/setpriv --reuid "$caller_uid" --regid "$(/usr/bin/id -g)" --clear-groups
    /usr/bin/bash -p -c '
      set -e
      umask 077
      state_dir=$1
      token=$2
      owner=$3
      control=$4
      temporary=""
      cleanup() { [[ -z $temporary ]] || /usr/bin/rm -f -- "$temporary"; }
      trap cleanup EXIT
      read -r process_stat <"/proc/$$/stat"
      process_stat=${process_stat##*) }
      read -r -a fields <<<"$process_stat"
      temporary=$(/usr/bin/mktemp "$state_dir/.inhibit-pid.XXXXXXXX")
      /usr/bin/chmod 600 "$temporary"
      printf "1 %s %s %s %s\n" "$$" "${fields[19]}" "$owner" "$token" >"$temporary"
      exec {control_fd}<"$control"
      /usr/bin/flock -x "$control_fd"
      IFS= read -r control_record <&"$control_fd"
      [[ $control_record == "active $token" ]]
      /usr/bin/mv -fT -- "$temporary" "$state_dir/inhibit-pid"
      temporary=""
      /usr/bin/flock -u "$control_fd"
      exec {control_fd}>&-
      trap - EXIT
      trap "kill \$!; exit 0" TERM
      /usr/bin/sleep infinity &
      wait
    ' "--why=Omarchy update in progress [$token]" "$state_dir" "$token" "$caller_uid" "$launch_control_file"
  )

  if (( EUID == 0 )); then
    (
      [[ -z ${OMARCHY_UPDATE_LOCK_FD:-} ]] || exec {OMARCHY_UPDATE_LOCK_FD}>&-
      exec {state_lock_fd}>&-
      exec "${hold_command[@]}"
    ) &
    launcher_pid=$!
  elif [[ -t 0 ]] || { (( update_sudo_session )) && /usr/bin/sudo -n -N /usr/bin/true 2>/dev/null; }; then
    # Inside an update without a terminal, reuse its authorization
    # non-interactively rather than asking again through polkit.
    sudo_options=(-N -b)
    [[ -t 0 ]] || sudo_options=(-n -N -b)
    if ! (
      [[ -z ${OMARCHY_UPDATE_LOCK_FD:-} ]] || exec {OMARCHY_UPDATE_LOCK_FD}>&-
      exec {state_lock_fd}>&-
      exec /usr/bin/sudo "${sudo_options[@]}" -- "${hold_command[@]}"
    ); then
      return 1
    fi
  else
    (
      [[ -z ${OMARCHY_UPDATE_LOCK_FD:-} ]] || exec {OMARCHY_UPDATE_LOCK_FD}>&-
      exec {state_lock_fd}>&-
      exec /usr/bin/pkexec "${hold_command[@]}"
    ) &
    launcher_pid=$!
  fi

  while :; do
    inhibit_record=$(read_state_record "$inhibit_pid_file" 2>/dev/null || true)
    if [[ $inhibit_record =~ ^1\ ([1-9][0-9]{0,18})\ ([1-9][0-9]{0,18})\ ([0-9]{1,10})\ ($token)$ ]]; then
      inhibit_pid="${BASH_REMATCH[1]}"
      inhibit_start_time="${BASH_REMATCH[2]}"
      inhibit_owner="${BASH_REMATCH[3]}"
      process_matches "$inhibit_pid" "$inhibit_start_time" "$inhibit_owner" "$token" && break
    fi
    if [[ -n $launcher_pid ]] && ! kill -0 "$launcher_pid" 2>/dev/null; then
      wait "$launcher_pid" || return 1
      echo "The update sleep inhibitor did not start." >&2
      return 1
    fi
    (( readiness_attempts += 1 ))
    if [[ -z $launcher_pid ]] && (( EUID != 0 && readiness_attempts >= 100 )); then
      echo "The update sleep inhibitor did not become ready." >&2
      return 1
    fi
    /usr/bin/sleep 0.05
  done

  if [[ -e $launch_control_file || -L $launch_control_file ]]; then
    state_file_is_private "$launch_control_file" || return 1
    rm -f -- "$launch_control_file"
  fi
  if [[ ! -f $stay_awake_state ]]; then
    atomic_write_state "$idle_owner_file" "$idle_owner" || return 1
    mkdir -p "$(dirname "$stay_awake_state")" || return 1
    # The idle toggle uses this marker. Publish its owner in the same write so
    # cancellation cannot leave an unowned Stay Awake setting behind.
    if ! printf '%s\n' "$idle_owner" >"$stay_awake_state"; then
      omarchy-toggle-idle allow-idle >/dev/null 2>&1 || true
      return 1
    fi
  fi
  launch_pending=0
  omarchy_security_install_sudo_cleanup_traps
}

case "${1:-}" in
  start | stop) ;;
  *)
    echo "Usage: omarchy-update-stay-awake <start|stop>" >&2
    exit 2
    ;;
esac

initialize_state_boundary
exec {state_lock_fd}<"$state_base" || fail_state_boundary
flock -x "$state_lock_fd" || fail_state_boundary
directory_is_private "$state_base" "$caller_uid" || fail_state_boundary

case "$1" in
  start)
    start_locked
    ;;
  stop)
    stop_locked
    ;;
esac