zeroclock/src-tauri/src/os_detection.rs

use serde::Serialize;
use std::collections::HashMap;
use windows::core::{PCWSTR, PWSTR};
use windows::Win32::Foundation::{BOOL, HWND, LPARAM};
use windows::Win32::Graphics::Gdi::{
    CreateCompatibleDC, DeleteDC, DeleteObject, GetDIBits, GetObjectW, BITMAP, BITMAPINFO,
    BITMAPINFOHEADER, DIB_RGB_COLORS, HBITMAP,
};
use windows::Win32::Storage::FileSystem::FILE_FLAGS_AND_ATTRIBUTES;
use windows::Win32::System::Threading::{
    OpenProcess, QueryFullProcessImageNameW, PROCESS_NAME_FORMAT, PROCESS_QUERY_LIMITED_INFORMATION,
};
use windows::Win32::UI::Input::KeyboardAndMouse::{GetLastInputInfo, LASTINPUTINFO};
use windows::Win32::UI::Shell::{SHGetFileInfoW, SHFILEINFOW, SHGFI_ICON, SHGFI_SMALLICON};
use windows::Win32::UI::WindowsAndMessaging::{
    DestroyIcon, EnumWindows, GetIconInfo, GetWindowTextLengthW, GetWindowTextW,
    GetWindowThreadProcessId, IsIconic, IsWindowVisible, ICONINFO,
};

#[derive(Debug, Serialize, Clone)]
pub struct WindowInfo {
    pub exe_name: String,
    pub exe_path: String,
    pub title: String,
    pub display_name: String,
    pub icon: Option<String>,
}

pub fn get_system_idle_seconds() -> u64 {
    unsafe {
        let mut info = LASTINPUTINFO {
            cbSize: std::mem::size_of::<LASTINPUTINFO>() as u32,
            dwTime: 0,
        };
        if GetLastInputInfo(&mut info).as_bool() {
            let tick_count = windows::Win32::System::SystemInformation::GetTickCount();
            let idle_ms = tick_count.wrapping_sub(info.dwTime);
            (idle_ms / 1000) as u64
        } else {
            0
        }
    }
}

fn get_process_exe_path(pid: u32) -> Option<String> {
    unsafe {
        let handle = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, false, pid).ok()?;
        let mut buf = [0u16; 1024];
        let mut size = buf.len() as u32;
        QueryFullProcessImageNameW(handle, PROCESS_NAME_FORMAT(0), PWSTR(buf.as_mut_ptr()), &mut size).ok()?;
        let _ = windows::Win32::Foundation::CloseHandle(handle);
        let path = String::from_utf16_lossy(&buf[..size as usize]);
        Some(path)
    }
}

fn get_window_title(hwnd: HWND) -> String {
    unsafe {
        let len = GetWindowTextLengthW(hwnd);
        if len == 0 {
            return String::new();
        }
        let mut buf = vec![0u16; (len + 1) as usize];
        let copied = GetWindowTextW(hwnd, &mut buf);
        String::from_utf16_lossy(&buf[..copied as usize])
    }
}

fn exe_name_from_path(path: &str) -> String {
    path.rsplit('\\').next().unwrap_or(path).to_string()
}

fn display_name_from_exe(exe_name: &str) -> String {
    exe_name
        .strip_suffix(".exe")
        .or_else(|| exe_name.strip_suffix(".EXE"))
        .unwrap_or(exe_name)
        .to_string()
}

struct EnumState {
    windows: Vec<WindowInfo>,
    include_minimized: bool,
}

unsafe extern "system" fn enum_windows_callback(hwnd: HWND, lparam: LPARAM) -> BOOL {
    let state = &mut *(lparam.0 as *mut EnumState);

    if !IsWindowVisible(hwnd).as_bool() {
        return BOOL(1);
    }

    if !state.include_minimized && IsIconic(hwnd).as_bool() {
        return BOOL(1);
    }

    let title = get_window_title(hwnd);
    if title.is_empty() {
        return BOOL(1);
    }

    let mut pid: u32 = 0;
    GetWindowThreadProcessId(hwnd, Some(&mut pid));
    if pid == 0 {
        return BOOL(1);
    }

    if let Some(exe_path) = get_process_exe_path(pid) {
        let exe_name = exe_name_from_path(&exe_path);
        let display_name = display_name_from_exe(&exe_name);
        state.windows.push(WindowInfo {
            exe_name,
            exe_path,
            title,
            display_name,
            icon: None,
        });
    }

    BOOL(1)
}

pub fn enumerate_visible_windows() -> Vec<WindowInfo> {
    let mut state = EnumState {
        windows: Vec::new(),
        include_minimized: false,
    };
    unsafe {
        let _ = EnumWindows(
            Some(enum_windows_callback),
            LPARAM(&mut state as *mut EnumState as isize),
        );
    }
    state.windows
}

pub fn enumerate_running_processes() -> Vec<WindowInfo> {
    let mut state = EnumState {
        windows: Vec::new(),
        include_minimized: true,
    };
    unsafe {
        let _ = EnumWindows(
            Some(enum_windows_callback),
            LPARAM(&mut state as *mut EnumState as isize),
        );
    }
    // Deduplicate by exe_path (case-insensitive)
    let mut seen = HashMap::new();
    let mut result = Vec::new();
    for w in state.windows {
        let key = w.exe_path.to_lowercase();
        if !seen.contains_key(&key) {
            seen.insert(key, true);
            result.push(w);
        }
    }
    result.sort_by(|a, b| a.display_name.to_lowercase().cmp(&b.display_name.to_lowercase()));

    // Extract icons for the deduplicated list
    for w in &mut result {
        w.icon = extract_icon_data_url(&w.exe_path);
    }

    result
}

// --- Icon extraction ---

fn extract_icon_data_url(exe_path: &str) -> Option<String> {
    unsafe {
        let wide: Vec<u16> = exe_path.encode_utf16().chain(std::iter::once(0)).collect();

        let mut fi = SHFILEINFOW::default();
        let res = SHGetFileInfoW(
            PCWSTR(wide.as_ptr()),
            FILE_FLAGS_AND_ATTRIBUTES(0),
            Some(&mut fi),
            std::mem::size_of::<SHFILEINFOW>() as u32,
            SHGFI_ICON | SHGFI_SMALLICON,
        );
        if res == 0 || fi.hIcon.is_invalid() {
            return None;
        }

        let hicon = fi.hIcon;
        let mut ii = ICONINFO::default();
        if GetIconInfo(hicon, &mut ii).is_err() {
            let _ = DestroyIcon(hicon);
            return None;
        }

        let result = extract_icon_pixels(ii.hbmColor, ii.hbmMask).and_then(|(rgba, w, h)| {
            let png_bytes = encode_rgba_to_png(&rgba, w, h)?;
            Some(format!("data:image/png;base64,{}", base64_encode(&png_bytes)))
        });

        // Cleanup
        if !ii.hbmColor.is_invalid() {
            let _ = DeleteObject(ii.hbmColor);
        }
        if !ii.hbmMask.is_invalid() {
            let _ = DeleteObject(ii.hbmMask);
        }
        let _ = DestroyIcon(hicon);

        result
    }
}

unsafe fn extract_icon_pixels(
    hbm_color: HBITMAP,
    hbm_mask: HBITMAP,
) -> Option<(Vec<u8>, u32, u32)> {
    if hbm_color.is_invalid() {
        return None;
    }

    let mut bm = BITMAP::default();
    if GetObjectW(
        hbm_color,
        std::mem::size_of::<BITMAP>() as i32,
        Some(&mut bm as *mut BITMAP as *mut std::ffi::c_void),
    ) == 0
    {
        return None;
    }

    let w = bm.bmWidth as u32;
    let h = bm.bmHeight as u32;
    if w == 0 || h == 0 {
        return None;
    }

    let hdc = CreateCompatibleDC(None);

    // Read color bitmap as 32-bit BGRA
    let mut bmi = make_bmi(w, h);
    let mut bgra = vec![0u8; (w * h * 4) as usize];
    let lines = GetDIBits(
        hdc,
        hbm_color,
        0,
        h,
        Some(bgra.as_mut_ptr() as *mut std::ffi::c_void),
        &mut bmi,
        DIB_RGB_COLORS,
    );
    if lines == 0 {
        let _ = DeleteDC(hdc);
        return None;
    }

    // Check if any pixel has a non-zero alpha
    let has_alpha = bgra.chunks_exact(4).any(|px| px[3] != 0);

    if !has_alpha && !hbm_mask.is_invalid() {
        // Read the mask bitmap as 32-bit to determine transparency
        let mut mask_bmi = make_bmi(w, h);
        let mut mask = vec![0u8; (w * h * 4) as usize];
        GetDIBits(
            hdc,
            hbm_mask,
            0,
            h,
            Some(mask.as_mut_ptr() as *mut std::ffi::c_void),
            &mut mask_bmi,
            DIB_RGB_COLORS,
        );
        // Mask: black (0,0,0) = opaque, white = transparent
        for i in (0..bgra.len()).step_by(4) {
            bgra[i + 3] = if mask[i] == 0 && mask[i + 1] == 0 && mask[i + 2] == 0 {
                255
            } else {
                0
            };
        }
    } else if !has_alpha {
        // No mask, assume fully opaque
        for px in bgra.chunks_exact_mut(4) {
            px[3] = 255;
        }
    }

    let _ = DeleteDC(hdc);

    // BGRA -> RGBA
    for px in bgra.chunks_exact_mut(4) {
        px.swap(0, 2);
    }

    Some((bgra, w, h))
}

fn make_bmi(w: u32, h: u32) -> BITMAPINFO {
    BITMAPINFO {
        bmiHeader: BITMAPINFOHEADER {
            biSize: std::mem::size_of::<BITMAPINFOHEADER>() as u32,
            biWidth: w as i32,
            biHeight: -(h as i32), // top-down
            biPlanes: 1,
            biBitCount: 32,
            biCompression: 0, // BI_RGB
            ..Default::default()
        },
        ..Default::default()
    }
}

fn encode_rgba_to_png(pixels: &[u8], width: u32, height: u32) -> Option<Vec<u8>> {
    let mut buf = Vec::new();
    {
        let mut encoder = png::Encoder::new(&mut buf, width, height);
        encoder.set_color(png::ColorType::Rgba);
        encoder.set_depth(png::BitDepth::Eight);
        let mut writer = encoder.write_header().ok()?;
        writer.write_image_data(pixels).ok()?;
        writer.finish().ok()?;
    }
    Some(buf)
}

fn base64_encode(data: &[u8]) -> String {
    const CHARS: &[u8; 64] =
        b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    let mut out = String::with_capacity((data.len() + 2) / 3 * 4);
    for chunk in data.chunks(3) {
        let b = [
            chunk[0],
            chunk.get(1).copied().unwrap_or(0),
            chunk.get(2).copied().unwrap_or(0),
        ];
        let n = ((b[0] as u32) << 16) | ((b[1] as u32) << 8) | (b[2] as u32);
        out.push(CHARS[((n >> 18) & 63) as usize] as char);
        out.push(CHARS[((n >> 12) & 63) as usize] as char);
        out.push(if chunk.len() > 1 {
            CHARS[((n >> 6) & 63) as usize] as char
        } else {
            '='
        });
        out.push(if chunk.len() > 2 {
            CHARS[(n & 63) as usize] as char
        } else {
            '='
        });
    }
    out
}