use crate::{
FocusTrackerConfig, FocusTrackerError, FocusTrackerResult, FocusedWindow, icon_cache::IconCache,
};
use focus_tracker_core::IconConfig;
use std::future::Future;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use super::utils;
use windows_sys::Win32::{
Foundation::{HWND, WPARAM},
Graphics::Gdi::{
BI_RGB, BITMAPINFO, BITMAPINFOHEADER, CreateCompatibleDC, DIB_RGB_COLORS, DeleteDC,
DeleteObject, GetDIBits, SelectObject,
},
UI::WindowsAndMessaging::{
DestroyIcon, GCLP_HICON, GCLP_HICONSM, GetClassLongPtrW, GetIconInfo, ICON_BIG, ICON_SMALL,
ICONINFO, SendMessageW, WM_GETICON,
},
};
#[derive(Default)]
struct FocusState {
hwnd: isize,
process_id: u32,
process_name: String,
window_title: Option<String>,
}
impl FocusState {
fn has_changed(&self, hwnd: HWND, process_id: u32, title: &Option<String>) -> bool {
let hwnd_value = hwnd as isize;
hwnd_value != self.hwnd
|| process_id != self.process_id
|| self.window_title.as_deref() != title.as_deref()
}
fn focus_changed(&self, hwnd: HWND) -> bool {
hwnd as isize != self.hwnd
}
fn update(&mut self, hwnd: HWND, process_id: u32, process_name: String, title: Option<String>) {
self.hwnd = hwnd as isize;
self.process_id = process_id;
self.process_name = process_name;
self.window_title = title;
}
fn clear(&mut self) {
self.hwnd = 0;
self.process_id = 0;
self.process_name.clear();
self.window_title = None;
}
}
#[inline]
fn should_stop(stop_signal: Option<&AtomicBool>) -> bool {
stop_signal.is_some_and(|stop| stop.load(Ordering::Relaxed))
}
fn poll_focus_change(
prev_state: &mut FocusState,
icon_cache: &mut IconCache,
config: &FocusTrackerConfig,
) -> Option<FocusedWindow> {
let Some(hwnd) = utils::get_foreground_window() else {
if prev_state.hwnd != 0 {
prev_state.clear();
}
return None;
};
let (title, process_name) = match utils::get_window_info(hwnd) {
Ok(info) => info,
Err(e) => {
tracing::debug!("Failed to get window info: {e}");
return None;
}
};
if config
.windows_ignore_rules
.matches(&process_name, title.as_deref())
{
return None;
}
let process_id = utils::get_window_process_id(hwnd).unwrap_or_default();
if !prev_state.has_changed(hwnd, process_id, &title) {
return None;
}
let icon = if prev_state.focus_changed(hwnd) {
resolve_icon(icon_cache, hwnd, process_id, &process_name, &config.icon)
} else {
icon_cache.get(&process_name).map(Arc::clone)
};
let focused = FocusedWindow {
process_id,
process_name: process_name.clone(),
window_title: title.clone(),
icon,
};
prev_state.update(hwnd, process_id, process_name, title);
Some(focused)
}
pub(crate) async fn track_focus<F, Fut>(
mut on_focus: F,
stop_signal: Option<&AtomicBool>,
config: &FocusTrackerConfig,
) -> FocusTrackerResult<()>
where
F: FnMut(FocusedWindow) -> Fut,
Fut: Future<Output = FocusTrackerResult<()>>,
{
if !utils::is_interactive_session()? {
return Err(FocusTrackerError::NotInteractiveSession);
}
let mut prev_state = FocusState::default();
let mut icon_cache = IconCache::new(config.icon_cache_capacity);
loop {
if should_stop(stop_signal) {
tracing::debug!("Stop signal received, exiting focus tracking loop");
break;
}
let pending = poll_focus_change(&mut prev_state, &mut icon_cache, config);
if let Some(focused) = pending {
on_focus(focused).await?;
}
tokio::time::sleep(config.poll_interval).await;
}
Ok(())
}
fn resolve_icon(
cache: &mut IconCache,
hwnd: HWND,
process_id: u32,
process_name: &str,
icon_config: &IconConfig,
) -> Option<Arc<image::RgbaImage>> {
if let Some(cached) = cache.get(process_name) {
return Some(Arc::clone(cached));
}
let image = match extract_window_icon(hwnd, icon_config) {
Ok(img) => img,
Err(e) => {
tracing::debug!("Failed to extract window icon: {e}, trying exe fallback");
match extract_exe_icon(process_id, icon_config) {
Ok(img) => img,
Err(e2) => {
tracing::debug!("Failed to extract exe icon: {e2}");
return None;
}
}
}
};
let icon = Arc::new(image);
cache.insert(process_name.to_owned(), Arc::clone(&icon));
Some(icon)
}
struct DcGuard(windows_sys::Win32::Graphics::Gdi::HDC);
impl Drop for DcGuard {
fn drop(&mut self) {
if !self.0.is_null() {
unsafe { DeleteDC(self.0) };
}
}
}
struct IconBitmapGuard {
hdc: windows_sys::Win32::Graphics::Gdi::HDC,
old_bitmap: windows_sys::Win32::Graphics::Gdi::HGDIOBJ,
hbm_color: windows_sys::Win32::Graphics::Gdi::HBITMAP,
hbm_mask: windows_sys::Win32::Graphics::Gdi::HBITMAP,
}
impl Drop for IconBitmapGuard {
fn drop(&mut self) {
unsafe {
if !self.hdc.is_null() {
SelectObject(self.hdc, self.old_bitmap);
}
if !self.hbm_color.is_null() {
DeleteObject(self.hbm_color);
}
if !self.hbm_mask.is_null() {
DeleteObject(self.hbm_mask);
}
}
}
}
fn acquire_icon_handle(hwnd: HWND) -> Result<(isize, bool), FocusTrackerError> {
let hicon = unsafe { SendMessageW(hwnd, WM_GETICON, ICON_BIG as WPARAM, 0) } as isize;
if hicon != 0 {
return Ok((hicon, false));
}
let hicon = unsafe { SendMessageW(hwnd, WM_GETICON, ICON_SMALL as WPARAM, 0) } as isize;
if hicon != 0 {
return Ok((hicon, false));
}
let hicon = unsafe { GetClassLongPtrW(hwnd, GCLP_HICON) } as isize;
if hicon != 0 {
return Ok((hicon, false));
}
let hicon = unsafe { GetClassLongPtrW(hwnd, GCLP_HICONSM) } as isize;
if hicon != 0 {
return Ok((hicon, false));
}
Err(FocusTrackerError::platform("no icon found for window"))
}
fn extract_window_icon(
hwnd: HWND,
icon_config: &IconConfig,
) -> FocusTrackerResult<image::RgbaImage> {
let (hicon, owned) = acquire_icon_handle(hwnd)?;
let result = icon_handle_to_image(hicon, icon_config);
if owned {
unsafe { DestroyIcon(hicon as _) };
}
result
}
fn extract_exe_icon(
process_id: u32,
icon_config: &IconConfig,
) -> FocusTrackerResult<image::RgbaImage> {
use windows_sys::Win32::UI::Shell::ExtractIconExW;
let exe_path = utils::get_process_exe_path(process_id)?;
let mut path_z = exe_path;
path_z.push(0);
let mut hicon_large = std::ptr::null_mut();
let count = unsafe {
ExtractIconExW(
path_z.as_ptr(),
0,
&mut hicon_large,
std::ptr::null_mut(),
1,
)
};
if count == 0 || hicon_large.is_null() {
return Err(FocusTrackerError::platform(
"no icon found in process executable",
));
}
let result = icon_handle_to_image(hicon_large as isize, icon_config);
unsafe { DestroyIcon(hicon_large as _) };
result
}
fn icon_handle_to_image(
hicon: isize,
icon_config: &IconConfig,
) -> FocusTrackerResult<image::RgbaImage> {
let mut icon_info: ICONINFO = unsafe { std::mem::zeroed() };
if unsafe { GetIconInfo(hicon as _, &mut icon_info) } == 0 {
return Err(FocusTrackerError::platform("failed to get icon info"));
}
let bitmap = if !icon_info.hbmColor.is_null() {
icon_info.hbmColor
} else {
icon_info.hbmMask
};
let hdc = unsafe { CreateCompatibleDC(std::ptr::null_mut()) };
if hdc.is_null() {
unsafe {
if !icon_info.hbmColor.is_null() {
DeleteObject(icon_info.hbmColor);
}
if !icon_info.hbmMask.is_null() {
DeleteObject(icon_info.hbmMask);
}
}
return Err(FocusTrackerError::platform("failed to create DC"));
}
let _dc_guard = DcGuard(hdc);
let old_bitmap = unsafe { SelectObject(hdc, bitmap) };
let _bmp_guard = IconBitmapGuard {
hdc,
old_bitmap,
hbm_color: icon_info.hbmColor,
hbm_mask: icon_info.hbmMask,
};
let mut bmi: BITMAPINFO = unsafe { std::mem::zeroed() };
bmi.bmiHeader.biSize = std::mem::size_of::<BITMAPINFOHEADER>() as u32;
if unsafe {
GetDIBits(
hdc,
bitmap,
0,
0,
std::ptr::null_mut(),
&mut bmi,
DIB_RGB_COLORS,
)
} == 0
{
return Err(FocusTrackerError::platform("failed to get bitmap info"));
}
let width = bmi.bmiHeader.biWidth as u32;
let height = bmi.bmiHeader.biHeight.unsigned_abs();
if width == 0 || height == 0 {
return Err(FocusTrackerError::platform("invalid icon dimensions"));
}
bmi.bmiHeader.biBitCount = 32;
bmi.bmiHeader.biCompression = BI_RGB;
bmi.bmiHeader.biHeight = -(height as i32);
let pixel_count = (width * height) as usize;
let mut pixels: Vec<u8> = vec![0; pixel_count * 4];
if unsafe {
GetDIBits(
hdc,
bitmap,
0,
height,
pixels.as_mut_ptr() as *mut _,
&mut bmi,
DIB_RGB_COLORS,
)
} == 0
{
return Err(FocusTrackerError::platform("failed to get bitmap bits"));
}
for i in (0..pixels.len()).step_by(4) {
pixels.swap(i, i + 2);
}
let mut image = image::RgbaImage::from_raw(width, height, pixels)
.ok_or_else(|| FocusTrackerError::platform("failed to create RgbaImage from pixel data"))?;
if let Some(target_size) = icon_config.size
&& (image.width() != target_size || image.height() != target_size)
{
image = image::imageops::resize(&image, target_size, target_size, icon_config.filter_type);
}
Ok(image)
}