helios-remote/crates/client/src/screenshot.rs

/// Screenshot capture — Windows GDI on Windows, stub on other platforms.
use base64::Engine;

#[cfg(windows)]
pub fn take_screenshot() -> Result<(String, u32, u32), String> {

    use windows::Win32::Graphics::Gdi::{
        BitBlt, CreateCompatibleBitmap, CreateCompatibleDC, DeleteDC, DeleteObject,
        GetDIBits, SelectObject, BITMAPINFO, BITMAPINFOHEADER,
        DIB_RGB_COLORS, SRCCOPY,
    };
    use windows::Win32::UI::WindowsAndMessaging::GetDesktopWindow;
    use windows::Win32::Graphics::Gdi::GetWindowDC;
    use windows::Win32::Graphics::Gdi::ReleaseDC;

    unsafe {
        let hwnd = GetDesktopWindow();
        let hdc_screen = GetWindowDC(hwnd);
        if hdc_screen.is_invalid() {
            return Err("GetWindowDC failed — cannot capture screen".to_string());
        }

        // Get screen dimensions
        use windows::Win32::Graphics::Gdi::{GetDeviceCaps, HORZRES, VERTRES};
        let width = GetDeviceCaps(hdc_screen, HORZRES) as u32;
        let height = GetDeviceCaps(hdc_screen, VERTRES) as u32;

        if width == 0 || height == 0 {
            ReleaseDC(hwnd, hdc_screen);
            return Err(format!("Invalid screen dimensions: {width}x{height}"));
        }

        // Create compatible DC and bitmap
        let hdc_mem = CreateCompatibleDC(hdc_screen);
        if hdc_mem.is_invalid() {
            ReleaseDC(hwnd, hdc_screen);
            return Err("CreateCompatibleDC failed".to_string());
        }

        let hbm = CreateCompatibleBitmap(hdc_screen, width as i32, height as i32);
        if hbm.is_invalid() {
            DeleteDC(hdc_mem);
            ReleaseDC(hwnd, hdc_screen);
            return Err("CreateCompatibleBitmap failed".to_string());
        }

        let old_obj = SelectObject(hdc_mem, hbm);

        // BitBlt the screen into our bitmap
        let blt_result = BitBlt(
            hdc_mem,
            0, 0,
            width as i32, height as i32,
            hdc_screen,
            0, 0,
            SRCCOPY,
        );

        if blt_result.is_err() {
            SelectObject(hdc_mem, old_obj);
            DeleteObject(hbm);
            DeleteDC(hdc_mem);
            ReleaseDC(hwnd, hdc_screen);
            return Err("BitBlt failed — could not copy screen pixels".to_string());
        }

        // Get raw pixel data via GetDIBits
        let mut bmi = BITMAPINFO {
            bmiHeader: BITMAPINFOHEADER {
                biSize: std::mem::size_of::<BITMAPINFOHEADER>() as u32,
                biWidth: width as i32,
                biHeight: -(height as i32), // negative = top-down
                biPlanes: 1,
                biBitCount: 32,
                biCompression: 0, // BI_RGB
                biSizeImage: 0,
                biXPelsPerMeter: 0,
                biYPelsPerMeter: 0,
                biClrUsed: 0,
                biClrImportant: 0,
            },
            bmiColors: [Default::default()],
        };

        let buf_size = (width * height * 4) as usize;
        let mut pixel_buf: Vec<u8> = vec![0u8; buf_size];

        let lines = GetDIBits(
            hdc_mem,
            hbm,
            0,
            height,
            Some(pixel_buf.as_mut_ptr() as *mut _),
            &mut bmi,
            DIB_RGB_COLORS,
        );

        SelectObject(hdc_mem, old_obj);
        DeleteObject(hbm);
        DeleteDC(hdc_mem);
        ReleaseDC(hwnd, hdc_screen);

        if lines == 0 {
            return Err(format!("GetDIBits failed — returned 0 scan lines (expected {height})"));
        }

        // Convert BGRA → RGBA
        for chunk in pixel_buf.chunks_exact_mut(4) {
            chunk.swap(0, 2); // B <-> R
        }

        // Encode as PNG
        let png_bytes = encode_png(&pixel_buf, width, height)?;
        let b64 = base64::engine::general_purpose::STANDARD.encode(&png_bytes);

        Ok((b64, width, height))
    }
}

/// Capture a specific window by cropping the full screen to its rect.
#[cfg(windows)]
pub fn take_window_screenshot(window_id: u64) -> Result<(String, u32, u32), String> {
    use windows::Win32::Foundation::{HWND, RECT};
    use windows::Win32::UI::WindowsAndMessaging::GetWindowRect;

    let hwnd = HWND(window_id as isize);
    let (x, y, w, h) = unsafe {
        let mut rect = RECT::default();
        GetWindowRect(hwnd, &mut rect).map_err(|e| format!("GetWindowRect failed: {e}"))?;
        let w = (rect.right - rect.left) as u32;
        let h = (rect.bottom - rect.top) as u32;
        if w == 0 || h == 0 { return Err(format!("Window has zero size: {w}x{h}")); }
        (rect.left, rect.top, w, h)
    };

    // Take full screenshot and crop to window rect
    let (full_b64, full_w, full_h) = take_screenshot()?;
    let full_bytes = base64::engine::general_purpose::STANDARD
        .decode(&full_b64).map_err(|e| format!("base64 decode: {e}"))?;

    // Decode PNG back to raw RGBA
    let cursor = std::io::Cursor::new(&full_bytes);
    let decoder = png::Decoder::new(cursor);
    let mut reader = decoder.read_info().map_err(|e| format!("PNG decode: {e}"))?;
    let mut img_buf = vec![0u8; reader.output_buffer_size()];
    reader.next_frame(&mut img_buf).map_err(|e| format!("PNG frame: {e}"))?;

    // Clamp window rect to screen bounds
    let x0 = (x.max(0) as u32).min(full_w);
    let y0 = (y.max(0) as u32).min(full_h);
    let x1 = ((x as u32 + w)).min(full_w);
    let y1 = ((y as u32 + h)).min(full_h);
    let cw = x1 - x0;
    let ch = y1 - y0;

    // Crop: 4 bytes per pixel (RGBA)
    let mut cropped = Vec::with_capacity((cw * ch * 4) as usize);
    for row in y0..y1 {
        let start = ((row * full_w + x0) * 4) as usize;
        let end = start + (cw * 4) as usize;
        cropped.extend_from_slice(&img_buf[start..end]);
    }

    let png_bytes = encode_png(&cropped, cw, ch)?;
    let b64 = base64::engine::general_purpose::STANDARD.encode(&png_bytes);
    Ok((b64, cw, ch))
}

#[cfg(not(windows))]
pub fn take_window_screenshot(_window_id: u64) -> Result<(String, u32, u32), String> {
    Err("Window screenshot only supported on Windows".to_string())
}

#[cfg(not(windows))]
pub fn take_screenshot() -> Result<(String, u32, u32), String> {
    // Stub for non-Windows builds
    // In a real scenario, could use X11/scrot on Linux
    let width = 1u32;
    let height = 1u32;
    let pixel_data = vec![0u8, 0u8, 0u8, 255u8]; // single black pixel RGBA
    let png_bytes = encode_png(&pixel_data, width, height)?;
    let b64 = base64::engine::general_purpose::STANDARD.encode(&png_bytes);
    Ok((b64, width, height))
}

fn encode_png(rgba: &[u8], width: u32, height: u32) -> Result<Vec<u8>, String> {
    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()
            .map_err(|e| format!("PNG header error: {e}"))?;
        writer
            .write_image_data(rgba)
            .map_err(|e| format!("PNG write error: {e}"))?;
    }
    Ok(buf)
}