amphetamine 0.1.2

// Amphetamine is built on Mach, libproc, AppKit and IOKit, so it is macOS-only
// by construction. Failing here turns a wall of objc2 and linker errors on
// another platform into one legible line. A bare `#![cfg]` would instead report
// a missing `main`, which explains nothing.
#[cfg(not(target_os = "macos"))]
compile_error!(
    "Amphetamine is macOS-only: it is built on Mach, libproc, AppKit and IOKit. \
     There is no cross-platform equivalent of what it does."
);

mod apps;
mod caches;
mod cli;
mod config;
mod focus;
mod guard;
mod manage;
mod pick;
mod privilege;
mod proc;
mod sysinfo;
mod ui;

use anyhow::{Context, Result};
use clap::Parser;
use cli::{Cli, Cmd, ConfigCmd};
use owo_colors::OwoColorize;
use std::process::ExitCode;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant};

fn main() -> ExitCode {
    match run() {
        Ok(()) => ExitCode::SUCCESS,
        Err(e) => {
            eprintln!("{} {e:#}", "error:".red().bold());
            ExitCode::FAILURE
        }
    }
}

fn run() -> Result<()> {
    match Cli::parse().command.unwrap_or(Cmd::Status {
        json: false,
        top: 10,
    }) {
        Cmd::Status { json, top } => status(json, top),
        Cmd::Boost {
            dry_run,
            no_caches,
            only_caches,
            force,
        } => boost(dry_run, no_caches, only_caches, force),
        Cmd::Focus { minutes, dry_run } => focus_session(minutes, dry_run),
        Cmd::Restore => restore(),
        Cmd::Setup { print, remove } => setup(print, remove),
        Cmd::Add { names, list } => edit_list(list.into(), &names, false),
        Cmd::Rm { names, list } => edit_list(list.into(), &names, true),
        Cmd::Pick { list } => pick_list(list.into()),
        Cmd::Config { action } => config_cmd(action.unwrap_or(ConfigCmd::Show)),
    }
}

/// Shared reporting for `add` and `rm`, so both explain themselves identically.
fn edit_list(list: manage::List, names: &[String], removing: bool) -> Result<()> {
    let table = proc::Table::load()?;
    let running = apps::list(&table);
    let changes = match removing {
        true => manage::remove(list, names, &running)?,
        false => manage::add(list, names, &running)?,
    };

    ui::title(list.label());
    for c in &changes {
        match c {
            manage::Change::Added(n) => ui::good(&format!("added {n}")),
            manage::Change::Removed(n) => ui::good(&format!("removed {n}")),
            manage::Change::AlreadyPresent(n) => ui::skipped(&format!("{n} is already there")),
            manage::Change::NotPresent(n) => ui::skipped(&format!("{n} was not in the list")),
            manage::Change::Pointless { name, why } => {
                ui::warn(&format!("{name} not added — {why}"));
            }
        }
    }
    hint_after_edit(list, &changes);
    Ok(())
}

fn pick_list(list: manage::List) -> Result<()> {
    let cfg = config::load()?.unwrap_or_default();
    let table = proc::Table::load()?;
    let running = apps::list(&table);
    let current = match list {
        manage::List::Close => &cfg.apps.close,
        manage::List::Protect => &cfg.apps.protect,
        manage::List::Demote => &cfg.focus.demote,
    };

    let changes = pick::run(list, &running, &table, current)?;
    if changes.is_empty() {
        return Ok(());
    }
    ui::title(list.label());
    for c in &changes {
        match c {
            manage::Change::Added(n) => ui::good(&format!("added {n}")),
            manage::Change::Removed(n) => ui::good(&format!("removed {n}")),
            _ => {}
        }
    }
    hint_after_edit(list, &changes);
    Ok(())
}

fn hint_after_edit(list: manage::List, changes: &[manage::Change]) {
    let touched = changes
        .iter()
        .any(|c| matches!(c, manage::Change::Added(_) | manage::Change::Removed(_)));
    if !touched {
        return;
    }
    println!();
    match list {
        manage::List::Demote => ui::note("preview with `amph focus --dry-run`"),
        _ => ui::note("preview with `amph boost --dry-run`"),
    }
}

fn setup(print: bool, remove: bool) -> Result<()> {
    if print {
        print!("{}", privilege::sudoers_rule(&privilege::current_user()));
        return Ok(());
    }
    if remove {
        privilege::uninstall()?;
        ui::title("Setup");
        ui::good(&format!("removed {}", privilege::SUDOERS_PATH));
        return Ok(());
    }
    if privilege::is_installed() {
        ui::title("Setup");
        ui::good(&format!("{} is already installed", privilege::SUDOERS_PATH));
        match privilege::can_restore() {
            true => ui::good("the grant works; focus mode can undo its own changes"),
            false => ui::warn("the file exists but the grant does not work"),
        }
        return Ok(());
    }

    ui::title("Setup");
    ui::note("Focus mode deprioritises apps so your editor gets more CPU. macOS lets");
    ui::note("any user lower a process's priority, but only root can raise it back,");
    ui::note("so without this Amphetamine would refuse to demote anything at all.");
    println!();
    ui::note(&format!("This installs {}:", privilege::SUDOERS_PATH));
    println!();
    for line in privilege::sudoers_rule(&privilege::current_user()).lines() {
        println!("    {}", line.dimmed());
    }
    println!();
    ui::note("The 0 is literal: the rule can only ever return a process to normal");
    ui::note("priority. It cannot deprioritise anything, run any other command, or");
    ui::note("open a shell. It is validated with visudo before being installed.");
    println!();
    ui::note("sudo will ask for your password once.");
    println!();

    privilege::install()?;
    ui::good(&format!("installed {}", privilege::SUDOERS_PATH));
    ui::good("verified: focus mode can now undo its own changes");
    ui::note("undo any time with `amph setup --remove`");
    Ok(())
}

/// Loads config, creating the starter file on first run.
///
/// Returns `None` after onboarding so the caller stops: a brand-new config has
/// empty lists, and silently doing nothing would look like a bug.
fn load_or_onboard() -> Result<Option<config::Config>> {
    if let Some(cfg) = config::load()? {
        return Ok(Some(cfg));
    }
    let path = config::path();
    config::init(&path)?;
    ui::title("First run");
    ui::good(&format!("wrote {}", path.display()));
    println!();
    ui::note("Nothing has been closed. Amphetamine will not quit an app until you");
    ui::note("name it in that file — the empty list is the safety mechanism.");
    println!();
    ui::note("Run `amph status` to see what is actually eating memory, add those");
    ui::note("apps to `close`, then run `amph boost` again.");
    Ok(None)
}

fn status(json: bool, top: usize) -> Result<()> {
    let cfg = config::load()?.unwrap_or_default();
    let snap = sysinfo::snapshot()?;
    let table = proc::Table::load()?;
    let app_list = apps::list(&table);
    let buckets = caches::scan(&cfg, &app_list);

    let reclaimable: u64 = buckets.iter().map(|b| b.eligible).sum();
    let planned: u64 = apps::plan(&cfg, &app_list)
        .iter()
        .filter_map(|d| match d {
            apps::Decision::Close(a) => Some(a.rss),
            apps::Decision::Refuse(..) => None,
        })
        .sum();

    if json {
        let out = serde_json::json!({
            "memory": {
                "total": snap.memory.total,
                "footprint": snap.memory.footprint(),
                "free": snap.memory.free,
                "compressed": snap.memory.compressed,
                "pressure_pct": snap.memory.pressure_pct(),
            },
            "swap": {
                "total": snap.swap.total,
                "used": snap.swap.used,
                "used_pct": snap.swap.used_pct(),
            },
            "power": {
                "mode": snap.power.mode.to_string(),
                "throttled": snap.power.throttled,
            },
            "reclaimable": { "apps": planned, "caches": reclaimable },
            "top_apps": app_list.iter().filter(|a| !a.nested).take(top).map(|a| serde_json::json!({
                "name": a.name,
                "bundle_id": a.bundle_id,
                "pid": a.pid,
                "rss": a.rss,
                "protected": guard::protected_match(&a.identities()).is_some(),
            })).collect::<Vec<_>>(),
        });
        println!("{}", serde_json::to_string_pretty(&out)?);
        return Ok(());
    }

    let m = &snap.memory;
    ui::title("Memory");
    ui::row(
        "In use",
        &format!(
            "{:>9} of {:<9} {}  {:.0}%",
            ui::bytes(m.footprint()),
            ui::bytes(m.total),
            ui::bar(m.pressure_pct()),
            m.pressure_pct()
        ),
    );
    ui::row("Free", &ui::bytes(m.free));
    ui::row("Compressed", &ui::bytes(m.compressed));
    ui::row("Reclaimable", &ui::bytes(m.reclaimable()));

    let s = &snap.swap;
    ui::title("Swap");
    ui::row(
        "In use",
        &format!(
            "{:>9} of {:<9} {}  {:.0}%",
            ui::bytes(s.used),
            ui::bytes(s.total),
            ui::bar(s.used_pct()),
            s.used_pct()
        ),
    );
    if s.used_pct() >= 80.0 {
        ui::warn("Swap is nearly full, which is a real drag on responsiveness.");
        ui::note("Every touch of a swapped-out page is a disk read. Swap only drains");
        ui::note("as the processes owning those pages exit, so closing apps is the");
        ui::note("only way to get it back — `purge` and free RAM will not do it.");
        ui::note(&format!(
            "{} pages have been faulted back in since boot, which is the",
            m.decompressions
        ));
        ui::note("cumulative cost of that pressure.");
    }

    ui::title("Power");
    ui::row("Mode", &snap.power.mode.to_string());
    ui::row("Thermal", &snap.power.thermal_note);
    if snap.power.mode != sysinfo::PowerMode::High {
        ui::note("System Settings › Battery › Energy Mode can be set to High Power.");
    }
    if !snap.power.throttled {
        ui::note("Nothing is limiting your clock, so there is no clock speed to win");
        ui::note("back. What is winnable is core contention — see `amph focus`.");
    }

    // Helpers are already counted inside their parent's tree; listing them
    // again would double-count and bury the apps you can actually act on.
    let top_level: Vec<&apps::App> = app_list.iter().filter(|a| !a.nested).collect();
    ui::title(&format!("Heaviest apps  ({} running)", top_level.len()));
    for a in top_level.iter().take(top) {
        let tag = match guard::protected_match(&a.identities()) {
            Some(_) => "protected".dimmed().to_string(),
            None if guard::any_matches(&a.identities(), &cfg.apps.close) => {
                "in close list".green().to_string()
            }
            None => String::new(),
        };
        let name = match a.foreground {
            true => a.name.clone(),
            // Background agents have no Dock icon, so quitting one is far less
            // obvious than quitting a window you can see.
            false => format!("{} (background)", a.name),
        };
        println!(
            "  {:>9}  {:<34} {:>3} proc  {}",
            ui::bytes(a.rss),
            name.chars().take(34).collect::<String>(),
            table.tree(a.pid).len(),
            tag
        );
    }

    ui::title("Reclaimable now");
    match planned {
        0 if cfg.apps.close.is_empty() => {
            ui::skipped("Apps: close list is empty");
            ui::note("pick from the list above with `amph pick`, or `amph add Slack Spotify`");
        }
        0 => ui::skipped("Apps: nothing in your close list is running"),
        n => ui::good(&format!("Apps: {} across your close list", ui::bytes(n))),
    }
    let clearable = buckets
        .iter()
        .filter(|b| matches!(b.verdict, caches::Verdict::Clear))
        .count();
    match reclaimable {
        0 => ui::skipped("Caches: nothing stale enough to clear"),
        n => ui::good(&format!(
            "Caches: {} across {clearable} buckets",
            ui::bytes(n)
        )),
    }
    println!();
    ui::note("`amph boost --dry-run` shows exactly what would be touched.");
    Ok(())
}

fn boost(dry_run: bool, no_caches: bool, only_caches: bool, force: bool) -> Result<()> {
    let Some(cfg) = load_or_onboard()? else {
        return Ok(());
    };

    let before = sysinfo::snapshot()?;
    let table = proc::Table::load()?;
    let app_list = apps::list(&table);

    if dry_run {
        ui::title("Dry run — nothing will be changed");
    }

    if !only_caches {
        ui::title("Apps");
        let plan = apps::plan(&cfg, &app_list);
        if plan.is_empty() {
            ui::skipped(&match cfg.apps.close.is_empty() {
                true => "close list is empty; add apps to your config".to_owned(),
                false => format!(
                    "none of your {} listed apps are running",
                    cfg.apps.close.len()
                ),
            });
        }
        let results = apps::execute(&plan, &cfg, dry_run, force);
        for r in &results {
            let label = format!("{} ({})", r.app.name, ui::bytes(r.app.rss));
            match &r.outcome {
                apps::Outcome::Closed if dry_run => ui::good(&format!("would close {label}")),
                apps::Outcome::Closed => ui::good(&format!("closed {label}")),
                apps::Outcome::StillRunning => {
                    ui::warn(&format!("{} is still running", r.app.name));
                    ui::note("it likely has an unsaved-work dialog open; left alone");
                }
                apps::Outcome::Refused(why) => ui::skipped(&format!("{} — {why}", r.app.name)),
                apps::Outcome::Failed(why) => ui::warn(&format!("{} — {why}", r.app.name)),
            }
        }
        let reclaimed: u64 = results.iter().map(|r| r.freed).sum();
        if reclaimed > 0 {
            ui::good(&format!(
                "{} attributed to closed apps",
                ui::bytes(reclaimed)
            ));
        }
    }

    if !no_caches && cfg.caches.enabled {
        ui::title("Caches");
        let buckets = caches::scan(&cfg, &app_list);
        let result = caches::sweep(&buckets, dry_run);

        for b in buckets.iter().take(6).filter(|b| b.eligible > 0) {
            let verb = if dry_run { "would clear" } else { "cleared" };
            ui::good(&format!("{verb} {} from {}", ui::bytes(b.eligible), b.name));
        }
        let shielded: Vec<&caches::Bucket> = buckets
            .iter()
            .filter(|b| matches!(b.verdict, caches::Verdict::Skipped(_)))
            .collect();
        // Largest first: the buckets worth explaining are the ones holding the
        // most, since those are what someone would wonder about.
        let mut biggest: Vec<&&caches::Bucket> = shielded.iter().collect();
        biggest.sort_by_key(|b| std::cmp::Reverse(b.total));
        for b in biggest.iter().take(3).filter(|b| b.total > 64 << 20) {
            if let caches::Verdict::Skipped(why) = &b.verdict {
                ui::skipped(&format!(
                    "kept {} in {} — {why}",
                    ui::bytes(b.total),
                    b.name
                ));
            }
        }
        if shielded.len() > 3 {
            ui::skipped(&format!("{} more buckets left alone", shielded.len() - 3));
        }
        match result.files {
            0 => ui::skipped(&format!(
                "nothing older than {} days to clear",
                cfg.caches.min_age_days
            )),
            n => ui::good(&format!(
                "{} across {n} files{}",
                ui::bytes(result.freed),
                if dry_run { " (dry run)" } else { "" }
            )),
        }
        for e in result.errors.iter().take(5) {
            ui::skipped(&format!("skipped {e}"));
        }
    }

    if !dry_run {
        // Let the kernel finish reclaiming before measuring, otherwise the
        // delta undercounts what was actually returned.
        std::thread::sleep(Duration::from_millis(1500));
        let after = sysinfo::snapshot()?;
        ui::title("Result");
        let freed = after.memory.free.saturating_sub(before.memory.free);
        ui::row("Memory freed", &ui::bytes(freed));
        ui::row(
            "Swap now",
            &format!(
                "{} of {} ({:.0}%)",
                ui::bytes(after.swap.used),
                ui::bytes(after.swap.total),
                after.swap.used_pct()
            ),
        );
        if after.swap.used < before.swap.used {
            ui::good(&format!(
                "swap drained by {}",
                ui::bytes(before.swap.used - after.swap.used)
            ));
        }
    }
    Ok(())
}

fn focus_session(minutes: Option<u64>, dry_run: bool) -> Result<()> {
    let Some(cfg) = load_or_onboard()? else {
        return Ok(());
    };
    let table = proc::Table::load()?;
    let app_list = apps::list(&table);
    let session = focus::Session::start(&cfg, &app_list, &table, dry_run)?;

    ui::title("Focus");
    if session.report.is_empty() {
        ui::skipped("demote list is empty; add noisy apps to [focus] in your config");
        return Ok(());
    }
    for d in &session.report {
        match &d.note {
            Some(why) => ui::skipped(&format!("{} — {why}", d.name)),
            None => ui::good(&format!(
                "{} {} ({} processes, {}) to nice {}",
                if dry_run {
                    "would deprioritise"
                } else {
                    "deprioritised"
                },
                d.name,
                d.pids,
                ui::bytes(d.rss),
                cfg.focus.nice_level.clamp(1, 20)
            )),
        }
    }
    if !session.restore_ready && !dry_run {
        println!();
        ui::warn("Nothing was changed: there is no way to undo a demotion yet.");
        ui::note("Run `amph setup` once to grant permission to restore priorities.");
        return Ok(());
    }
    if session.holding_sleep {
        ui::good("holding off idle sleep");
    }
    if dry_run {
        return Ok(());
    }

    let stop = Arc::new(AtomicBool::new(false));
    ctrlc::set_handler({
        let stop = stop.clone();
        move || stop.store(true, Ordering::SeqCst)
    })?;

    let deadline = minutes.map(|m| Instant::now() + Duration::from_secs(m * 60));
    println!();
    match minutes {
        Some(m) => ui::note(&format!("holding for {m} minutes — Ctrl-C to end early")),
        None => ui::note("holding — Ctrl-C to end"),
    }

    while !stop.load(Ordering::SeqCst) {
        if deadline.is_some_and(|d| Instant::now() >= d) {
            break;
        }
        std::thread::sleep(Duration::from_millis(200));
    }

    // Explicit rather than incidental: this is what puts the machine back, and
    // any failure here needs to be seen rather than swallowed by Drop.
    let mut session = session;
    let n = session.demoted_count();
    let errors = session.restore_now();
    ui::title("Session ended");
    ui::good(&format!(
        "restored {} of {n} processes to normal priority",
        n - errors.len()
    ));
    for e in errors.iter().take(5) {
        ui::warn(e);
    }
    if !errors.is_empty() {
        ui::note("run `amph restore` to retry");
    }
    Ok(())
}

fn restore() -> Result<()> {
    let cfg = config::load()?.unwrap_or_default();
    ui::title("Restore");
    if !privilege::can_restore() {
        ui::warn("cannot restore priorities without the privilege grant");
        ui::note("run `amph setup` first");
        return Ok(());
    }
    let table = proc::Table::load()?;
    let app_list = apps::list(&table);
    let restored = focus::restore_all(&cfg, &app_list, &table);
    if restored.is_empty() {
        ui::good("nothing was deprioritised");
    }
    for (name, n) in restored {
        ui::good(&format!("{name}: {n} processes back to normal priority"));
    }
    Ok(())
}

fn config_cmd(action: ConfigCmd) -> Result<()> {
    let path = config::path();
    match action {
        ConfigCmd::Path => println!("{}", path.display()),
        ConfigCmd::Init => {
            config::init(&path)?;
            ui::good(&format!("wrote {}", path.display()));
        }
        ConfigCmd::Raw => match path.exists() {
            true => print!("{}", std::fs::read_to_string(&path)?),
            false => {
                ui::skipped(&format!("no config at {}", path.display()));
                ui::note("run `amph config init` to create one");
            }
        },
        ConfigCmd::Edit => {
            if !path.exists() {
                config::init(&path)?;
            }
            let editor = std::env::var("VISUAL")
                .or_else(|_| std::env::var("EDITOR"))
                .unwrap_or_else(|_| "vi".into());
            // Split so EDITOR="code -w" works rather than being treated as a
            // single binary name with a space in it.
            let mut parts = editor.split_whitespace();
            let bin = parts.next().unwrap_or("vi");
            std::process::Command::new(bin)
                .args(parts)
                .arg(&path)
                .status()
                .with_context(|| format!("launching {editor}"))?;
        }
        ConfigCmd::Show => show_config(&path)?,
    }
    Ok(())
}

/// Renders the config against live state, so a listed app that is not running
/// or a grant that is not installed is visible rather than implied.
fn show_config(path: &std::path::Path) -> Result<()> {
    let Some(cfg) = config::load()? else {
        ui::title("Config");
        ui::skipped(&format!("no config at {}", path.display()));
        ui::note("run `amph config init`, or just `amph add <app>` to create one");
        return Ok(());
    };
    let table = proc::Table::load()?;
    let running = apps::list(&table);

    ui::title("Config");
    ui::row("File", &path.display().to_string());

    for (label, names, hint) in [
        ("Close on boost", &cfg.apps.close, "amph add <app>"),
        ("Extra protected", &cfg.apps.protect, "amph add -p <app>"),
        ("Deprioritise", &cfg.focus.demote, "amph add -d <app>"),
    ] {
        ui::title(label);
        if names.is_empty() {
            ui::skipped(&format!("empty — add with `{hint}`"));
            continue;
        }
        for name in names {
            let found = running
                .iter()
                .find(|a| guard::identity_matches(&a.identities(), name));
            match found {
                Some(a) => ui::good(&format!("{:<24} running · {}", name, ui::bytes(a.rss))),
                None => ui::skipped(&format!("{name:<24} not running")),
            }
        }
    }

    ui::title("Caches");
    ui::row(
        "Sweep",
        match cfg.caches.enabled {
            true => "enabled",
            false => "disabled",
        },
    );
    ui::row(
        "Age floor",
        &format!("older than {} days", cfg.caches.min_age_days),
    );
    ui::row(
        "Running apps",
        match cfg.caches.skip_running_apps {
            true => "skipped",
            false => "NOT skipped — risky",
        },
    );
    if !cfg.caches.allow.is_empty() {
        ui::row("Opted in", &cfg.caches.allow.join(", "));
    }
    if !cfg.caches.skip.is_empty() {
        ui::row("Also skipped", &cfg.caches.skip.join(", "));
    }

    ui::title("Focus");
    ui::row("Nice level", &cfg.focus.nice_level.clamp(1, 20).to_string());
    ui::row(
        "Idle sleep",
        match cfg.focus.prevent_sleep {
            true => "held off during a session",
            false => "left alone",
        },
    );
    match privilege::can_restore() {
        true => ui::good("privilege grant installed; focus mode can undo itself"),
        false => {
            ui::skipped("no privilege grant — focus mode will refuse to demote");
            ui::note("run `amph setup` once to enable it");
        }
    }
    Ok(())
}