caffeinate

#!/bin/bash

# Linux equivalent of macOS caffeinate
# Supports: -d, -i, -s, -u, -t, -w, -v
# Excludes: -m (disk sleep prevention)
# Battery-aware, systemd-native, resource-light

set -euo pipefail

readonly PROGNAME=$(basename "$0")
readonly PLACEHOLDER_SLEEP_CMD="sleep 86400"

# Globals
INHIBIT_WHAT=()
DBUS_ACTIONS=()
TIMEOUT_SEC=0
WAIT_PID=0
CMD_ARGS=()
SIMULATE_USER=false
AC_REQUIRED=false
ON_AC=true
VERBOSE=false
SCREENSAVER_COOKIE=0

# --- Helpers ---
log_info() {
    if [[ "$VERBOSE" == true ]]; then
        echo >&2 "[INFO] $*"
    fi
}

log_warn() {
    echo >&2 "Warning: $*"
}

check_tool() {
    if ! command -v "$1" >/dev/null 2>&1; then
        echo >&2 "$PROGNAME: required tool '$1' not found"
        exit 1
    fi
}

check_ac_power() {
    ON_AC=false

    # Try upower first (most reliable)
    if command -v upower >/dev/null 2>&1; then
        local ac_path
        # Look for any AC/line power device
        ac_path=$(upower -e 2>/dev/null | grep -E '/(line_power|AC)' | head -n1)
        if [[ -n "$ac_path" ]] && upower -i "$ac_path" 2>/dev/null | grep -q "online.*yes"; then
            ON_AC=true
            return
        fi
    fi

    # Fallback to sysfs - check multiple possible AC adapter names
    for ac_pattern in "AC" "ADP" "ACAD" "ACPI" "ADP0" "ADP1"; do
        for ac in /sys/class/power_supply/${ac_pattern}*/online; do
            if [[ -f "$ac" ]] && [[ "$(cat "$ac" 2>/dev/null)" == "1" ]]; then
                ON_AC=true
                return
            fi
        done
    done

    # Additional fallback: check if any power supply is online
    if [[ -d /sys/class/power_supply ]]; then
        for psu in /sys/class/power_supply/*/online; do
            if [[ -f "$psu" ]] && [[ "$(cat "$psu" 2>/dev/null)" == "1" ]]; then
                # Check if this is likely an AC adapter (not battery)
                local type_file="${psu%/online}/type"
                if [[ -f "$type_file" ]] && [[ "$(cat "$type_file" 2>/dev/null)" != "Battery" ]]; then
                    ON_AC=true
                    return
                fi
            fi
        done
    fi

    # Last resort: assume AC if we can't determine (safety first for -s flag)
    # This ensures -s flag works on VMs and containers where power detection fails
    if [[ "$AC_REQUIRED" == true ]]; then
        log_warn "Could not determine AC power status - assuming AC for safety"
        ON_AC=true
    fi
}

has_gui_session() {
    # Check for GUI session using multiple indicators
    [[ -n "${XDG_SESSION_TYPE:-}" && "${XDG_SESSION_TYPE}" != "tty" ]] || \
    [[ -n "${WAYLAND_DISPLAY:-}" ]] || \
    [[ -n "${DISPLAY:-}" ]] || \
    [[ -n "${KDE_FULL_SESSION:-}" ]] || \
    [[ -n "${GNOME_DESKTOP_SESSION_ID:-}" ]] || \
    [[ -n "${XFCE_DESKTOP_SESSION_ID:-}" ]] || \
    [[ -n "${MATE_DESKTOP_SESSION_ID:-}" ]] || \
    [[ -n "${CINNAMON_DESKTOP_SESSION_ID:-}" ]] || \
    [[ -n "${LXDE_DESKTOP_SESSION_ID:-}" ]] || \
    [[ -n "${LXQT_DESKTOP_SESSION_ID:-}" ]] || \
    [[ -n "${COSMIC_DESKTOP_SESSION_ID:-}" ]]
}

simulate_user_activity() {
    if ! has_gui_session; then
        return
    fi

    # Try GNOME ScreenSaver
    if command -v gdbus >/dev/null 2>&1; then
        gdbus call --session \
            --dest org.gnome.ScreenSaver \
            --object-path /org/gnome/ScreenSaver \
            --method org.gnome.ScreenSaver.SimulateUserActivity \
            >/dev/null 2>&1 && return
    fi

    # Try freedesktop ScreenSaver (KDE and others)
    if command -v dbus-send >/dev/null 2>&1; then
        dbus-send --session --type=method_call \
            --dest=org.freedesktop.ScreenSaver \
            /ScreenSaver org.freedesktop.ScreenSaver.SimulateUserActivity \
            >/dev/null 2>&1 && return
    fi

    # Try KDE-specific screen saver control
    if command -v qdbus >/dev/null 2>&1; then
        qdbus org.kde.screensaver /ScreenSaver SimulateUserActivity >/dev/null 2>&1 && return
    fi

    # Try X11 screen saver extension (fallback for older systems)
    if command -v xset >/dev/null 2>&1; then
        xset s reset >/dev/null 2>&1 && return
    fi

    log_warn "Could not simulate user activity (no compatible screen saver interface detected)."
}

inhibit_display_sleep() {
    if ! has_gui_session; then
        return
    fi

    # Try to inhibit using org.freedesktop.ScreenSaver.Inhibit method
    if command -v dbus-send >/dev/null 2>&1; then
        local cookie
        cookie=$(dbus-send --session --type=method_call --print-reply --dest=org.freedesktop.ScreenSaver \
            /ScreenSaver org.freedesktop.ScreenSaver.Inhibit \
            string:"caffeinate" string:"Prevent display sleep" 2>/dev/null | \
            awk '/uint32/ {print $2}')

        if [[ -n "$cookie" ]] && [[ "$cookie" =~ ^[0-9]+$ ]]; then
            SCREENSAVER_COOKIE="$cookie"
            log_info "Inhibited display sleep via freedesktop ScreenSaver (cookie: $SCREENSAVER_COOKIE)"
            return
        fi
    fi

    # Try GNOME-specific inhibition
    if command -v gdbus >/dev/null 2>&1; then
        local cookie
        cookie=$(gdbus call --session \
            --dest org.gnome.ScreenSaver \
            --object-path /org/gnome/ScreenSaver \
            --method org.gnome.ScreenSaver.Inhibit \
            "caffeinate" "Prevent display sleep" 2>/dev/null | \
            awk -F'[()]' '/uint32/ {print $2}')

        if [[ -n "$cookie" ]] && [[ "$cookie" =~ ^[0-9]+$ ]]; then
            SCREENSAVER_COOKIE="$cookie"
            log_info "Inhibited display sleep via GNOME ScreenSaver (cookie: $SCREENSAVER_COOKIE)"
            return
        fi
    fi

    # Try KDE-specific inhibition
    if command -v qdbus >/dev/null 2>&1; then
        local cookie
        cookie=$(qdbus org.kde.screensaver /ScreenSaver Inhibit \
            "caffeinate" "Prevent display sleep" 2>/dev/null)

        if [[ -n "$cookie" ]] && [[ "$cookie" =~ ^[0-9]+$ ]]; then
            SCREENSAVER_COOKIE="$cookie"
            log_info "Inhibited display sleep via KDE ScreenSaver (cookie: $cookie)"
            return
        fi
    fi

    # Try X11 DPMS inhibition (for older systems)
    if command -v xset >/dev/null 2>&1; then
        xset s off -display "$DISPLAY" >/dev/null 2>&1
        xset -dpms -display "$DISPLAY" >/dev/null 2>&1
        log_info "Inhibited display sleep via X11 DPMS"
        return
    fi

    # Fallback: simulate user activity periodically
    log_info "Falling back to periodic user activity simulation"
    simulate_user_activity
}

parse_args() {
    local OPTIND opt
    while getopts "disuht:w:v" opt; do
        case "$opt" in
            d) DBUS_ACTIONS+=("display") ;;
            i) INHIBIT_WHAT+=("idle") ;;
            s) INHIBIT_WHAT+=("sleep"); AC_REQUIRED=true ;;
            u) SIMULATE_USER=true; DBUS_ACTIONS+=("user") ;;
            t)
                if [[ "$OPTARG" =~ ^[0-9]+$ ]] && (( OPTARG > 0 )); then
                    TIMEOUT_SEC="$OPTARG"
                else
                    echo >&2 "$PROGNAME: invalid timeout: $OPTARG"
                    exit 1
                fi
                ;;
            w)
                if [[ "$OPTARG" =~ ^[0-9]+$ ]] && kill -0 "$OPTARG" 2>/dev/null; then
                    WAIT_PID="$OPTARG"
                else
                    echo >&2 "$PROGNAME: invalid or non-existent PID: $OPTARG"
                    exit 1
                fi
                ;;
            v) VERBOSE=true ;;
            h)
                cat <<EOF
Usage: $PROGNAME [-disuv] [-t timeout] [-w pid] [utility [args...]]

Prevent the system from sleeping on behalf of a utility.

Flags:
  -d      Prevent the display from sleeping.
  -i      Prevent the system from idle sleeping.
  -s      Prevent the system from sleeping (AC power only).
  -u      Simulate user activity (wakes display, resets idle timer).
  -t N    Timeout in seconds (ignored if command or -w is used).
  -w PID  Wait for process PID to exit (ignores -t).
  -v      Verbose mode.
  -h      Show this help.

Note: -t has no effect when running a command or using -w.
      -m (disk sleep) is not supported on Linux.

Examples:
  $PROGNAME                      # Prevent idle sleep until exit
  $PROGNAME -t 1800              # Keep awake for 30 minutes
  $PROGNAME -u -t 60             # Simulate activity for 60s
  $PROGNAME -i make              # Run 'make' with idle sleep prevented
  $PROGNAME -s -w 1234           # Prevent sleep (if on AC) until PID 1234 exits
EOF
                exit 0
                ;;
            *)
                exit 1
                ;;
        esac
    done
    shift $((OPTIND - 1))
    CMD_ARGS=("$@")
}

wait_for_pid() {
    local pid="$1"
    while kill -0 "$pid" 2>/dev/null; do
        sleep 2
    done
}

main() {
    parse_args "$@"

    if [[ ${#INHIBIT_WHAT[@]} -eq 0 ]] && [[ ${#CMD_ARGS[@]} -eq 0 ]] && [[ $WAIT_PID -eq 0 ]]; then
        INHIBIT_WHAT=("idle")
    fi

    if [[ "$AC_REQUIRED" == true ]]; then
        check_ac_power
        if [[ "$ON_AC" == false ]]; then
            log_warn "On battery: ignoring -s (full sleep prevention)."
            INHIBIT_WHAT=("${INHIBIT_WHAT[@]/sleep}")
            if [[ ${#INHIBIT_WHAT[@]} -eq 0 ]] && [[ ${#CMD_ARGS[@]} -eq 0 ]] && [[ $WAIT_PID -eq 0 ]]; then
                echo "$PROGNAME: nothing to do."
                exit 0
            fi
        fi
    fi

    if [[ "$SIMULATE_USER" == true ]]; then
        log_info "Simulating user activity (display wake + idle reset)"
        simulate_user_activity
    fi

    if [[ " ${DBUS_ACTIONS[*]} " == *" display "* ]]; then
        log_info "Inhibiting display sleep (best-effort)"
        inhibit_display_sleep
    fi

    inhibit_cmd=()
    if [[ ${#INHIBIT_WHAT[@]} -gt 0 ]]; then
        IFS=":"; what_str="${INHIBIT_WHAT[*]}"; unset IFS
        inhibit_cmd=( systemd-inhibit --mode=block --what="$what_str" )
        log_info "Inhibiting: $what_str"
    fi

    local mode="idle"
    local timeout_used=false
    if [[ ${#CMD_ARGS[@]} -gt 0 ]]; then
        mode="command"
        if (( TIMEOUT_SEC > 0 )); then
            log_info "-t ignored because a command was provided"
        fi
        TIMEOUT_SEC=0
    elif [[ $WAIT_PID -gt 0 ]]; then
        mode="waitpid"
        if (( TIMEOUT_SEC > 0 )); then
            log_info "-t ignored because -w was used"
        fi
        TIMEOUT_SEC=0
    else
        timeout_used=true
    fi

    if [[ "$mode" == "command" ]]; then
        full_cmd=( "${inhibit_cmd[@]}" "${CMD_ARGS[@]}" )
    else
        placeholder_cmd=( $PLACEHOLDER_SLEEP_CMD )
        if [[ ${#inhibit_cmd[@]} -gt 0 ]]; then
            full_cmd=( "${inhibit_cmd[@]}" "${placeholder_cmd[@]}" )
        else
            full_cmd=( "${placeholder_cmd[@]}" )
        fi
    fi

    cleanup() {
        pkill -P $$ -f "$PLACEHOLDER_SLEEP_CMD" 2>/dev/null || true
        # Release screensaver inhibition if we have a cookie
        if [[ "$SCREENSAVER_COOKIE" -gt 0 ]]; then
            # Try freedesktop ScreenSaver first
            dbus-send --session --type=method_call --dest=org.freedesktop.ScreenSaver \
                /ScreenSaver org.freedesktop.ScreenSaver.UnInhibit \
                uint32:"$SCREENSAVER_COOKIE" 2>/dev/null || true

            # Try GNOME ScreenSaver
            gdbus call --session \
                --dest org.gnome.ScreenSaver \
                --object-path /org/gnome/ScreenSaver \
                --method org.gnome.ScreenSaver.UnInhibit \
                "$SCREENSAVER_COOKIE" 2>/dev/null || true

            # Try KDE ScreenSaver
            qdbus org.kde.screensaver /ScreenSaver UnInhibit \
                "$SCREENSAVER_COOKIE" 2>/dev/null || true

            log_info "Released display sleep inhibition (cookie: $SCREENSAVER_COOKIE)"
        fi

        # Restore X11 DPMS settings if we disabled them
        if command -v xset >/dev/null 2>&1 && [[ -n "${DISPLAY:-}" ]]; then
            xset s on -display "$DISPLAY" 2>/dev/null || true
            xset +dpms -display "$DISPLAY" 2>/dev/null || true
        fi
    }
    trap cleanup EXIT INT TERM

    if [[ "$mode" == "command" ]]; then
        if ! "${full_cmd[@]}"; then
            log_warn "Command or systemd-inhibit failed (exit code: $?)"
        fi
    elif [[ "$mode" == "waitpid" ]]; then
        "${full_cmd[@]}" &
        wait_for_pid "$WAIT_PID"
    else
        if (( TIMEOUT_SEC > 0 )); then
            if command -v timeout >/dev/null 2>&1; then
                timeout "$TIMEOUT_SEC" "${full_cmd[@]}" 2>/dev/null || true
            else
                log_info "Using manual sleep fallback (timeout not available)"
                "${full_cmd[@]}" &
                sleep "$TIMEOUT_SEC"
                cleanup
            fi
        else
            "${full_cmd[@]}"
        fi
    fi
}

check_tool systemd-inhibit
check_tool grep
check_tool sleep

main "$@"