Automatically switch power plan

power-plan-auto-switcher.rs

use std::char::decode_utf16;
use std::collections::BTreeSet;
use std::ffi::c_void;
use std::mem::{self, ManuallyDrop};
use std::sync::mpsc::{channel, Receiver, RecvTimeoutError, Sender};
use std::sync::OnceLock;
use std::thread;
use std::time::Duration;

use windows::core::{Error, Interface, Result, GUID};
use windows::Win32::Foundation::{BOOL, ERROR_INVALID_HANDLE, HWND, LPARAM};
use windows::Win32::System::Com::{
    CoCreateInstance, CoInitializeEx, CLSCTX_INPROC_SERVER, COINIT_MULTITHREADED,
};
use windows::Win32::System::Power::PowerSetActiveScheme;
use windows::Win32::UI::Accessibility::{SetWinEventHook, HWINEVENTHOOK};
use windows::Win32::UI::Shell::{
    IVirtualDesktopManager, SHQueryUserNotificationState, VirtualDesktopManager, QUNS_BUSY,
    QUNS_RUNNING_D3D_FULL_SCREEN,
};
use windows::Win32::UI::WindowsAndMessaging::{
    DispatchMessageW, EnumWindows, GetMessageW, GetWindow, GetWindowLongW, GetWindowTextLengthW,
    GetWindowTextW, IsWindowVisible, RealGetWindowClassA, CHILDID_SELF, EVENT_OBJECT_CREATE,
    EVENT_OBJECT_DESTROY, EVENT_OBJECT_FOCUS, GWL_EXSTYLE, GW_OWNER, MSG, OBJID_WINDOW,
    WINEVENT_OUTOFCONTEXT, WS_EX_APPWINDOW, WS_EX_NOACTIVATE, WS_EX_TOOLWINDOW,
};

static VDM: OnceLock<usize> = OnceLock::new();
static RECHECK_SENDER: OnceLock<Sender<(HWND, u32)>> = OnceLock::new();
static SWITCHER_SENDER: OnceLock<Sender<bool>> = OnceLock::new();
const IGNORED_CLASSES: &[&[u8]] = &[b"Windows.", b"ATL", b"GDI+", b"Console"];

unsafe extern "system" fn eventproc(
    _hwineventhook: HWINEVENTHOOK,
    event: u32,
    hwnd: HWND,
    idobject: i32,
    idchild: i32,
    _ideventthread: u32,
    _dwmseventtime: u32,
) {
    match event {
        EVENT_OBJECT_CREATE | EVENT_OBJECT_DESTROY
            if idobject == OBJID_WINDOW.0 && idchild as u32 == CHILDID_SELF =>
        {
            RECHECK_SENDER.get().unwrap().send((hwnd, event)).unwrap()
        }
        EVENT_OBJECT_FOCUS => SWITCHER_SENDER.get().unwrap().send(false).unwrap(),
        _ => (),
    }
}

unsafe fn recheck_thread(receiver: Receiver<(HWND, u32)>) {
    let mut rechecks = BTreeSet::new();
    loop {
        match if rechecks.is_empty() {
            receiver.recv().map_err(|_| RecvTimeoutError::Disconnected)
        } else {
            receiver.recv_timeout(Duration::from_millis(16))
        } {
            Ok((hwnd, event)) => match event {
                EVENT_OBJECT_DESTROY => {
                    rechecks.remove(&hwnd.0);
                }
                EVENT_OBJECT_CREATE => match is_app_window(hwnd) {
                    Some(true) => handle_app_start(hwnd),
                    Some(false) => (),
                    None => {
                        rechecks.insert(hwnd.0);
                    }
                },
                _ => (),
            },
            Err(RecvTimeoutError::Timeout) => {
                for hwnd in &rechecks {
                    let hwnd = HWND(*hwnd);
                    if is_app_window(hwnd).unwrap_or(true) {
                        handle_app_start(hwnd);
                    }
                }
                rechecks.clear();
            }
            Err(RecvTimeoutError::Disconnected) => break,
        }
    }
}

unsafe fn handle_app_start(hwnd: HWND) {
    let mut text = [0; 32];
    let Ok(len) = GetWindowTextW(hwnd, &mut text).try_into() else {
        return;
    };
    let text = &text[..len];
    let text: String = decode_utf16(text.iter().copied())
        .map(|r| r.unwrap_or(char::REPLACEMENT_CHARACTER))
        .collect();
    let mut class = [0; 32];
    let Ok(len) = RealGetWindowClassA(hwnd, &mut class).try_into() else {
        return;
    };
    let class = &class[..len];
    let class = String::from_utf8_lossy(class);
    eprintln!("{text} ({class}) started");

    SWITCHER_SENDER.get().unwrap().send(true).unwrap();
}

unsafe fn switcher_thread(receiver: Receiver<bool>) -> Result<()> {
    const HIGH_PERFORMANCE: GUID = GUID::from_values(
        0x8c5e7fda,
        0xe8bf,
        0x4a96,
        [0x9a, 0x85, 0xa6, 0xe2, 0x3a, 0x8c, 0x63, 0x5c],
    );
    const POWER_SAVER: GUID = GUID::from_values(
        0xa1841308,
        0x3541,
        0x4fab,
        [0xbc, 0x81, 0xf7, 0x15, 0x56, 0xf2, 0x0b, 0x4a],
    );
    let mut is_hp = false;
    let mut is_ff = false;
    loop {
        match if is_hp {
            receiver.recv_timeout(Duration::from_secs(8))
        } else {
            receiver.recv().map_err(|_| RecvTimeoutError::Disconnected)
        } {
            Ok(true) => {
                if !is_hp {
                    activate_power_plan(&HIGH_PERFORMANCE)?;
                    is_hp = true;
                    eprintln!("Switched to High Performance");
                }
            }
            Ok(false) => {
                if !is_hp {
                    let has_ff = has_fullscreen();
                    if has_ff {
                        activate_power_plan(&HIGH_PERFORMANCE)?;
                        is_hp = true;
                        eprintln!("Fullscreen application detected; switched to High Performance");
                    }
                    is_ff = has_ff;
                }
            }
            Err(RecvTimeoutError::Timeout) => {
                let has_ff = has_fullscreen();
                if !has_ff {
                    activate_power_plan(&POWER_SAVER)?;
                    is_hp = false;
                    eprintln!("Switched to Power Saver");
                }
                if has_ff && !is_ff {
                    eprintln!("Fullscreen application detected; keep High Performance");
                }
                is_ff = has_ff;
            }
            Err(RecvTimeoutError::Disconnected) => break Ok(()),
        }
    }
}

unsafe fn activate_power_plan(guid: &GUID) -> Result<()> {
    PowerSetActiveScheme(None, Some(guid as *const GUID))
}

unsafe fn hook() -> Result<()> {
    VDM.set(create_vdm()?.into_raw() as usize).unwrap();

    debug_assert!(EVENT_OBJECT_CREATE < EVENT_OBJECT_DESTROY);
    debug_assert!(EVENT_OBJECT_DESTROY < EVENT_OBJECT_FOCUS);
    let hook = SetWinEventHook(
        EVENT_OBJECT_CREATE,
        EVENT_OBJECT_FOCUS,
        None,
        Some(eventproc),
        0,
        0,
        WINEVENT_OUTOFCONTEXT,
    );
    if hook.is_invalid() {
        return Err(Error::from(ERROR_INVALID_HANDLE));
    }

    let (sender, receiver) = channel();
    RECHECK_SENDER.set(sender).unwrap();
    mem::forget(thread::spawn(move || {
        recheck_thread(receiver);
    }));

    let (sender, receiver) = channel();
    SWITCHER_SENDER.set(sender).unwrap();
    mem::forget(thread::spawn(move || {
        switcher_thread(receiver).unwrap();
    }));

    Ok(())
}

unsafe fn event_loop() {
    let mut msg = MSG::default();
    while GetMessageW(&mut msg as *mut MSG, None, 0, 0).as_bool() {
        DispatchMessageW(&msg as *const MSG);
    }
}

unsafe extern "system" fn enum_proc(hwnd: HWND, lparam: LPARAM) -> BOOL {
    let target = lparam.0 as *mut isize;
    if hwnd.0 == *target {
        *target = 0;
        false.into()
    } else {
        true.into()
    }
}

unsafe fn is_app_window(hwnd: HWND) -> Option<bool> {
    if hwnd.0 <= 0 {
        return Some(false);
    }
    let mut target = hwnd.clone();
    let _ = EnumWindows(
        Some(enum_proc),
        LPARAM((&mut target.0) as *mut isize as isize),
    );
    let vdm = get_vdm();
    let ex_sty = GetWindowLongW(hwnd, GWL_EXSTYLE) as u32;
    let mut class = [0; 32];
    RealGetWindowClassA(hwnd, &mut class);
    let is_app_unsure = GetWindow(hwnd, GW_OWNER).0 == 0
        && (ex_sty & WS_EX_NOACTIVATE.0) == 0
        && (ex_sty & WS_EX_TOOLWINDOW.0) == 0
        && GetWindowTextLengthW(hwnd) != 0
        && !IGNORED_CLASSES
            .iter()
            .any(|ignore| class.starts_with(ignore))
        && (target.0 == 0 || class.starts_with(b"CicMarshalWndClass\0"));
    let is_app_sure =
        (ex_sty & WS_EX_APPWINDOW.0) != 0 || IsWindowVisible(hwnd).as_bool() && is_app_unsure;
    let is_on_current_vd = vdm
        .IsWindowOnCurrentVirtualDesktop(hwnd)
        .unwrap_or_default()
        .as_bool();
    if is_on_current_vd {
        if is_app_sure {
            Some(true)
        } else if is_app_unsure {
            None
        } else {
            Some(false)
        }
    } else {
        Some(false)
    }
}

unsafe fn has_fullscreen() -> bool {
    SHQueryUserNotificationState()
        .is_ok_and(|quns| matches!(quns, QUNS_BUSY | QUNS_RUNNING_D3D_FULL_SCREEN))
}

unsafe fn create_vdm() -> Result<IVirtualDesktopManager> {
    CoCreateInstance(
        &VirtualDesktopManager as *const GUID,
        None,
        CLSCTX_INPROC_SERVER,
    )
}

unsafe fn get_vdm() -> IVirtualDesktopManager {
    let ptr = *VDM.get().unwrap() as *mut c_void;
    (&*ManuallyDrop::new(IVirtualDesktopManager::from_raw(ptr))).clone()
}

fn main() {
    unsafe {
        CoInitializeEx(None, COINIT_MULTITHREADED).unwrap();
        hook().unwrap();
        event_loop()
    }
}