vaken 0.1.0

// Vaken wraps the macOS-only `caffeinate` command and sets macOS-only
// Info.plist keys (LSUIElement). The whole product is a macOS menu bar
// app — fail loud at compile time on any other target instead of
// pretending to build and then panicking when the user clicks Toggle.
#[cfg(not(target_os = "macos"))]
compile_error!("vaken only builds on macOS — it wraps the macOS-only `caffeinate` command");

use std::process::{Child, Command, Stdio};
use std::sync::{Arc, Mutex};

use muda::{Menu, MenuId, MenuItem};
use tray_icon::{Icon, TrayIcon, TrayIconBuilder, menu::MenuEvent};
use winit::application::ApplicationHandler;
use winit::event::{StartCause, WindowEvent};
use winit::event_loop::{ActiveEventLoop, ControlFlow, EventLoop};
use winit::window::WindowId;

// caffeinate flags: -d display sleep, -i idle sleep, -m disk sleep,
// -s system sleep, -u assert user-active. The combination is the standard
// "keep the machine fully awake until I say otherwise" stack. We also
// pass `-w <vaken_pid>` per-spawn (see App::vaken_pid + spawn site) so
// the caffeinate child self-exits whenever Vaken dies — kernel-level
// orphan prevention that catches SIGKILL / panic / hard crashes that no
// user-space cleanup hook could intercept.
const CAFFEINATE_FLAGS: &[&str] = &["-d", "-i", "-m", "-s", "-u"];

// Icons are baked into the binary so the app works regardless of CWD.
// A .app bundle launched from Finder/Dock has CWD = "/"; embedding avoids
// any runtime filesystem dependency.
const INACTIVE_ICON_BYTES: &[u8] = include_bytes!("../icons/sovbar.png");
const ACTIVE_ICON_BYTES: &[u8] = include_bytes!("../icons/vaken.png");

/// Wraps muda's menu events so they ride through the winit event loop.
/// Without this, the main thread would block on `MenuEvent::receiver()`
/// instead of pumping the AppKit run loop, and macOS would mark the
/// `.app` as "not responding" within ~10s of launch.
#[derive(Debug)]
enum UserEvent {
    Menu(MenuEvent),
}

fn decode_icon(bytes: &'static [u8]) -> Icon {
    let image = image::load_from_memory(bytes)
        .expect("decoding embedded icon")
        .into_rgba8();
    let (width, height) = image.dimensions();
    let rgba = image.into_raw();
    Icon::from_rgba(rgba, width, height).expect("constructing icon from embedded bytes")
}

/// Kill + reap the caffeinate child if one is running. Idempotent.
fn stop_caffeinate(slot: &Mutex<Option<Child>>) {
    if let Some(mut child) = slot.lock().unwrap().take() {
        let _ = child.kill();
        // Wait avoids leaving a zombie process; on macOS the system reaps
        // eventually, but explicit wait keeps `ps` clean while Vaken runs.
        let _ = child.wait();
    }
}

/// Startup cleanup: kill any `caffeinate` process matching vaken's
/// exact flag signature that has no `-w` watcher attached.
///
/// Pre-`-w-fix` builds of vaken (or any future code path that fails
/// to attach `-w <pid>`) could leave such a process running forever
/// after vaken crashed or was force-quit, permanently holding the
/// Mac awake even across reboots of vaken itself. The kernel-level
/// `-w` watcher prevents new orphans, but it cannot retroactively
/// clean up a pre-existing one. This sweep does that.
///
/// A caffeinate with `-w <pid>` in its args is left alone: it is
/// either managed by another live vaken (parallel install) or will
/// be reaped by the kernel when its watched pid dies. Best-effort —
/// failures (no `pgrep`, no permission to `kill`, etc.) are silent.
fn sweep_orphan_caffeinates() {
    let Ok(output) = Command::new("pgrep").args(["-x", "caffeinate"]).output() else {
        return;
    };
    let pids = String::from_utf8_lossy(&output.stdout);
    for pid in pids.lines().filter(|l| !l.trim().is_empty()) {
        let Ok(ps_out) = Command::new("ps").args(["-o", "args=", "-p", pid]).output() else {
            continue;
        };
        let args = String::from_utf8_lossy(&ps_out.stdout);
        let argv: Vec<&str> = args.split_whitespace().collect();
        let has_all_flags = CAFFEINATE_FLAGS.iter().all(|f| argv.contains(f));
        let has_w = argv.contains(&"-w");
        if has_all_flags && !has_w {
            let _ = Command::new("kill").arg(pid).output();
            eprintln!("vaken: killed orphan caffeinate pid={pid}");
        }
    }
}

struct App {
    caffeinate: Arc<Mutex<Option<Child>>>,
    /// Held so we can `set_text()` it when toggling awake mode.
    toggle_item: MenuItem,
    toggle_id: MenuId,
    quit_id: MenuId,
    /// Vaken's own PID, captured once at startup, passed to caffeinate
    /// as `-w <pid>` so the child self-terminates if Vaken dies for any
    /// reason — including SIGKILL or hard crash, which user-space
    /// cleanup hooks cannot catch.
    vaken_pid: String,
    inactive_icon: Icon,
    active_icon: Icon,
    /// Menu is built up-front but only handed to the `TrayIcon` once the
    /// run loop is ready (see `new_events` / `StartCause::Init`).
    /// `Option` so we can `.take()` it into the builder at that point.
    menu: Option<Menu>,
    /// Tray construction is deferred to `StartCause::Init` because on
    /// macOS the underlying `NSStatusItem` requires the `NSApplication`
    /// to be initialized — which only holds true after winit has booted
    /// the run loop. See the tray-icon crate's winit example.
    tray: Option<TrayIcon>,
}

impl App {
    fn new() -> Result<Self, Box<dyn std::error::Error>> {
        // Defensive cleanup: any caffeinate left over from a pre-`-w-fix`
        // vaken (or a crash that somehow bypassed the watcher) gets
        // killed before we spawn a new one. Without this, a fresh vaken
        // restart cannot fix a "Mac is stuck awake" condition.
        sweep_orphan_caffeinates();

        let toggle_item = MenuItem::new("Enable Awake Mode", true, None);
        // Custom Quit (rather than `PredefinedMenuItem::quit`) so we can
        // kill caffeinate before exiting — otherwise the Mac stays awake
        // until reboot when the user quits while awake mode is on.
        let quit_item = MenuItem::new("Quit Vaken", true, None);
        let toggle_id = toggle_item.id().clone();
        let quit_id = quit_item.id().clone();

        let menu = Menu::new();
        menu.append(&toggle_item)?;
        menu.append(&quit_item)?;

        Ok(Self {
            caffeinate: Arc::new(Mutex::new(None)),
            toggle_item,
            toggle_id,
            quit_id,
            vaken_pid: std::process::id().to_string(),
            inactive_icon: decode_icon(INACTIVE_ICON_BYTES),
            active_icon: decode_icon(ACTIVE_ICON_BYTES),
            menu: Some(menu),
            tray: None,
        })
    }

    fn ensure_tray(&mut self) {
        if self.tray.is_some() {
            return;
        }
        let menu = self
            .menu
            .take()
            .expect("menu must still be present when ensure_tray runs");
        let tray = TrayIconBuilder::new()
            .with_menu(Box::new(menu))
            .with_tooltip("Vaken")
            .with_icon(self.inactive_icon.clone())
            .build()
            .expect("building tray icon");
        self.tray = Some(tray);
    }

    fn handle_menu(&mut self, event: MenuEvent, event_loop: &ActiveEventLoop) {
        if event.id == self.quit_id {
            stop_caffeinate(&self.caffeinate);
            event_loop.exit();
            return;
        }
        if event.id != self.toggle_id {
            return;
        }

        let mut proc = self.caffeinate.lock().unwrap();
        match proc.take() {
            // OFF → ENABLE
            None => match Command::new("caffeinate")
                .args(CAFFEINATE_FLAGS)
                // -w <pid>: caffeinate exits when this pid exits. Vaken
                // dies → kernel notifies caffeinate → caffeinate exits.
                // Crash-safe regardless of *how* Vaken dies.
                .arg("-w")
                .arg(&self.vaken_pid)
                .stdin(Stdio::null())
                .stdout(Stdio::null())
                .stderr(Stdio::null())
                .spawn()
            {
                Ok(child) => {
                    *proc = Some(child);
                    self.toggle_item.set_text("Disable Awake Mode");
                    if let Some(tray) = &self.tray {
                        let _ = tray.set_icon(Some(self.active_icon.clone()));
                        let _ = tray.set_tooltip(Some("Vaken — Awake Mode"));
                    }
                }
                Err(err) => {
                    eprintln!("Failed to launch caffeinate: {err}");
                    // Surface failure in the menu so the user knows why
                    // nothing changed; stderr is invisible for a menu-bar app.
                    self.toggle_item.set_text("Enable Awake Mode (failed)");
                }
            },
            // ON → DISABLE
            Some(mut child) => {
                let _ = child.kill();
                let _ = child.wait();
                self.toggle_item.set_text("Enable Awake Mode");
                if let Some(tray) = &self.tray {
                    let _ = tray.set_icon(Some(self.inactive_icon.clone()));
                    let _ = tray.set_tooltip(Some("Vaken — Sleepy Mode"));
                }
            }
        }
    }
}

impl ApplicationHandler<UserEvent> for App {
    fn new_events(&mut self, _event_loop: &ActiveEventLoop, cause: StartCause) {
        // StartCause::Init fires exactly once, after the NSApplication
        // is up but before any user events — the canonical hook to
        // construct the tray icon on macOS.
        if matches!(cause, StartCause::Init) {
            self.ensure_tray();
        }
    }

    fn resumed(&mut self, _event_loop: &ActiveEventLoop) {
        // No-op: menu-bar app has no windows to recreate on resume.
    }

    fn window_event(&mut self, _: &ActiveEventLoop, _: WindowId, _: WindowEvent) {
        // No-op: no windows.
    }

    fn user_event(&mut self, event_loop: &ActiveEventLoop, event: UserEvent) {
        match event {
            UserEvent::Menu(e) => self.handle_menu(e, event_loop),
        }
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let event_loop = EventLoop::<UserEvent>::with_user_event().build()?;
    // Wait parks the loop until the OS wakes it — perfect for a menu
    // bar app that only acts on user input. (Poll would burn CPU 24/7.)
    event_loop.set_control_flow(ControlFlow::Wait);

    // Route muda menu events through the winit event loop so the main
    // thread wakes on each click and our handler runs. `set_event_handler`
    // replaces muda's default `MenuEvent::receiver()` channel, which we
    // don't use.
    let proxy = event_loop.create_proxy();
    MenuEvent::set_event_handler(Some(move |event| {
        let _ = proxy.send_event(UserEvent::Menu(event));
    }));

    let mut app = App::new()?;
    event_loop.run_app(&mut app)?;
    Ok(())
}