#!/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

source "${BASH_SOURCE[0]%/*}/omarchy-security-functions" || exit 126
omarchy_security_require_privileged_bash_startup || exit 126
set -e
omarchy_security_sanitize_bash_environment "$0" "$@"
omarchy_security_revoke_sudo_timestamp || exit 1
omarchy_security_install_sudo_cleanup_traps
omarchy_security_enable_no_update_sudo

state_dir="${XDG_RUNTIME_DIR:-/tmp/omarchy-$UID}/omarchy-update-stay-awake"
idle_owner_file="$state_dir/idle-owner"
inhibit_pid_file="$state_dir/inhibit-pid"
stay_awake_state="$HOME/.local/state/omarchy/indicators/stay-awake"

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
  printf '%s\n' "${stat_fields[19]}"
}

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%% *}"
}

stop() {
  local inhibit_pid=""
  local recorded_start_time=""
  local current_start_time=""
  local idle_owner=""
  local current_idle_owner=""

  if [[ -s $idle_owner_file ]]; then
    idle_owner=$(<"$idle_owner_file")
    if [[ -f $stay_awake_state ]]; then
      current_idle_owner=$(<"$stay_awake_state")
    fi
    if [[ -n $idle_owner && $current_idle_owner == "$idle_owner" ]]; then
      omarchy-toggle-idle allow-idle >/dev/null 2>&1 || true
    fi
    rm -f "$idle_owner_file"
  fi

  if [[ -s $inhibit_pid_file ]]; then
    read -r inhibit_pid recorded_start_time <"$inhibit_pid_file" || true
    if [[ $inhibit_pid =~ ^[0-9]+$ ]]; then
      current_start_time=$(process_start_time "$inhibit_pid" 2>/dev/null || true)
    fi
    if [[ -n $recorded_start_time && $current_start_time == "$recorded_start_time" ]]; then
      kill "$inhibit_pid" >/dev/null 2>&1 || true

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

      current_start_time=$(process_start_time "$inhibit_pid" 2>/dev/null || true)
      if [[ $current_start_time == "$recorded_start_time" ]] &&
        [[ $(process_state "$inhibit_pid" 2>/dev/null || true) != "Z" ]]; then
        echo "Failed to stop the Omarchy update sleep inhibitor." >&2
        return 1
      fi
    fi
    rm -f "$inhibit_pid_file"
  fi

  rmdir "$state_dir" 2>/dev/null || true
}

start() {
  local idle_owner="$$:$RANDOM:$RANDOM"
  local launcher_pid=""

  stop
  mkdir -p "$state_dir"

  # sudo authenticates in the foreground, then backgrounds the inhibitor.
  # The held command drops back to this user before publishing its PID, so stop
  # can release the inhibitor without another privileged operation or ticket.
  local hold_command=(
    /usr/bin/systemd-inhibit --what=sleep:idle --who=omarchy-update
    --why="Omarchy update in progress" --mode=block
    /usr/bin/setpriv --reuid "$UID" --regid "$(id -g)" --clear-groups
    /usr/bin/bash -p -c '
      read -r process_stat <"/proc/$$/stat"
      process_stat=${process_stat##*) }
      read -r -a fields <<< "$process_stat"
      printf "%s %s\n" "$$" "${fields[19]}" >"$1"
      exec /usr/bin/sleep infinity
    ' omarchy-update-inhibitor "$inhibit_pid_file"
  )

  if (( EUID == 0 )); then
    ( [[ -z ${OMARCHY_UPDATE_LOCK_FD:-} ]] || exec {OMARCHY_UPDATE_LOCK_FD}>&-
      exec "${hold_command[@]}" ) &
    launcher_pid=$!
  elif [[ -t 0 ]]; then
    /usr/bin/sudo -N -b -- "${hold_command[@]}"
  else
    ( [[ -z ${OMARCHY_UPDATE_LOCK_FD:-} ]] || exec {OMARCHY_UPDATE_LOCK_FD}>&-
      exec pkexec "${hold_command[@]}" ) &
    launcher_pid=$!
  fi

  # For graphical authentication the launcher may wait for a password. Wait for
  # either the user-owned held command to become ready or the launcher to fail.
  # sudo -b backgrounds internally, so Bash has no launcher PID in the TTY path;
  # $! there can still refer to a completed startup process substitution.
  local readiness_attempts=0
  while [[ ! -s $inhibit_pid_file ]]; do
    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=$((readiness_attempts + 1))
    if [[ -t 0 ]] && (( EUID != 0 && readiness_attempts >= 100 )); then
      echo "The update sleep inhibitor did not become ready." >&2
      return 1
    fi
    sleep 0.05
  done

  if [[ ! -f $stay_awake_state ]]; then
    printf '%s\n' "$idle_owner" >"$idle_owner_file"
    mkdir -p "$(dirname "$stay_awake_state")"
    if omarchy-toggle-idle stay-awake >/dev/null 2>&1; then
      printf '%s\n' "$idle_owner" >"$stay_awake_state"
    else
      rm -f "$idle_owner_file"
    fi
  fi
}

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