ambiligth/src/screenshot.rs

pub type ScreenResult<T> = Result<T, anyhow::Error>;

#[derive(Clone, Copy)]
pub struct Pixel {
    pub a: u8,
    pub r: u8,
    pub g: u8,
    pub b: u8,
}

pub struct Screenshot {
    pub data: Vec<Pixel>,
}

impl Screenshot {
    pub fn average_color(&self, brightness_boost: f32, saturation_boost: f32) -> rgb::Rgb<u8> {
        let mut sum_l = 0.0f32;
        let mut sum_a = 0.0f32;
        let mut sum_b = 0.0f32;

        for pixel in &self.data {
            if pixel.a == 0 {
                continue;
            }

            let r = pixel.r as f32 / 255.0;
            let g = pixel.g as f32 / 255.0;
            let b = pixel.b as f32 / 255.0;

            let linear_rgb = oklab::LinearRgb::new(r, g, b);
            let oklab_color = oklab::linear_srgb_to_oklab(linear_rgb);
            sum_l += oklab_color.l;
            sum_a += oklab_color.a;
            sum_b += oklab_color.b;
        }

        let count = self.data.len() as f32;
        let avg_l = sum_l / count;
        let avg_a = sum_a / count;
        let avg_b = sum_b / count;

        let target_l = brightness_boost;
        let boosted_l = (avg_l + target_l) / 2.0;

        let boosted_a = avg_a * saturation_boost;
        let boosted_b = avg_b * saturation_boost;

        let oklab_avg = oklab::Oklab {
            l: boosted_l,
            a: boosted_a,
            b: boosted_b,
        };
        let linear_rgb = oklab::oklab_to_linear_srgb(oklab_avg);

        rgb::Rgb {
            r: (linear_rgb.r * 255.0).clamp(0.0, 255.0) as u8,
            g: (linear_rgb.g * 255.0).clamp(0.0, 255.0) as u8,
            b: (linear_rgb.b * 255.0).clamp(0.0, 255.0) as u8,
        }
    }

    pub fn color_diff(
        &self,
        color: &rgb::Rgb<u8>,
        brightness_boost: f32,
        saturation_boost: f32,
    ) -> f32 {
        let avg_color = self.average_color(brightness_boost, saturation_boost);

        let r1 = avg_color.r as f32 / 255.0;
        let g1 = avg_color.g as f32 / 255.0;
        let b1 = avg_color.b as f32 / 255.0;

        let r2 = color.r as f32 / 255.0;
        let g2 = color.g as f32 / 255.0;
        let b2 = color.b as f32 / 255.0;

        let oklab1 = oklab::linear_srgb_to_oklab(oklab::LinearRgb::new(r1, g1, b1));
        let oklab2 = oklab::linear_srgb_to_oklab(oklab::LinearRgb::new(r2, g2, b2));

        let delta_l = (oklab1.l - oklab2.l).abs();
        let delta_a = (oklab1.a - oklab2.a).abs();
        let delta_b = (oklab1.b - oklab2.b).abs();

        let delta_e = (delta_l * delta_l + delta_a * delta_a + delta_b * delta_b).sqrt();

        (delta_e * 1000.0).clamp(0.0, 100.0) as f32
    }
}

#[cfg(any(
    target_os = "linux",
    not(any(target_os = "macos", target_os = "windows"))
))]
mod ffi {
    use super::{Pixel, ScreenResult, Screenshot};
    use grim_rs::Grim;

    pub fn get_screenshot_inner(_screen: usize) -> ScreenResult<Screenshot> {
        let mut grim =
            Grim::new().map_err(|e| anyhow::anyhow!("Failed to initialize grim: {}", e))?;
        let result = grim
            .capture_all()
            .map_err(|e| anyhow::anyhow!("Failed to capture screenshot: {}", e))?;

        let raw_data = result.data();

        let pixels: Vec<Pixel> = raw_data
            .chunks_exact(4)
            .map(|chunk| Pixel {
                r: chunk[0],
                g: chunk[1],
                b: chunk[2],
                a: chunk[3],
            })
            .collect();

        Ok(Screenshot { data: pixels })
    }
}

#[cfg(target_os = "macos")]
mod ffi {
    #![allow(non_upper_case_globals, dead_code)]
    use super::{Pixel, ScreenResult, Screenshot};
    use libc;
    use std::slice;

    type CFIndex = libc::c_long;
    type CFDataRef = *const u8;
    type CGError = libc::int32_t;
    type CGDirectDisplayID = libc::uint32_t;
    type CGDisplayCount = libc::uint32_t;
    type CGImageRef = *mut u8;
    type CGDataProviderRef = *mut u8;
    const CGDisplayNoErr: CGError = 0;

    #[link(name = "CoreGraphics", kind = "framework")]
    extern "C" {
        fn CGGetActiveDisplayList(
            max_displays: libc::uint32_t,
            active_displays: *mut CGDirectDisplayID,
            display_count: *mut CGDisplayCount,
        ) -> CGError;
        fn CGDisplayCreateImage(displayID: CGDirectDisplayID) -> CGImageRef;
        fn CGImageRelease(image: CGImageRef);
        fn CGImageGetBitsPerPixel(image: CGImageRef) -> libc::size_t;
        fn CGImageGetDataProvider(image: CGImageRef) -> CGDataProviderRef;
        fn CGImageGetHeight(image: CGImageRef) -> libc::size_t;
        fn CGImageGetWidth(image: CGImageRef) -> libc::size_t;
        fn CGDataProviderCopyData(provider: CGDataProviderRef) -> CFDataRef;
    }

    #[link(name = "CoreFoundation", kind = "framework")]
    extern "C" {
        fn CFDataGetLength(theData: CFDataRef) -> CFIndex;
        fn CFDataGetBytePtr(theData: CFDataRef) -> *const u8;
        fn CFRelease(cf: *const libc::c_void);
    }

    pub fn get_screenshot_inner(screen: usize) -> ScreenResult<Screenshot> {
        unsafe {
            let mut count: CGDisplayCount = 0;
            let mut err = CGGetActiveDisplayList(0, std::ptr::null_mut(), &mut count);
            if err != CGDisplayNoErr {
                return Err(anyhow::anyhow!("Error getting number of displays."));
            }

            let mut disps: Vec<CGDisplayCount> = Vec::with_capacity(count as usize);
            disps.set_len(count as usize);
            err = CGGetActiveDisplayList(
                disps.len() as libc::uint32_t,
                disps.as_mut_ptr(),
                &mut count,
            );
            if err != CGDisplayNoErr {
                return Err(anyhow::anyhow!("Error getting list of displays."));
            }

            let disp_id = disps[screen];
            let cg_img = CGDisplayCreateImage(disp_id);
            let width = CGImageGetWidth(cg_img) as usize;
            let height = CGImageGetHeight(cg_img) as usize;
            let pixel_bits = CGImageGetBitsPerPixel(cg_img) as usize;

            if pixel_bits % 8 != 0 {
                return Err(anyhow::anyhow!("Pixels aren't integral bytes."));
            }

            let cf_data = CGDataProviderCopyData(CGImageGetDataProvider(cg_img));
            let raw_len = CFDataGetLength(cf_data) as usize;

            let data = slice::from_raw_parts(CFDataGetBytePtr(cf_data), raw_len);
            let pixels: Vec<Pixel> = data
                .chunks_exact(4)
                .map(|chunk| Pixel {
                    b: chunk[0],
                    g: chunk[1],
                    r: chunk[2],
                    a: chunk[3],
                })
                .collect();

            CGImageRelease(cg_img);
            CFRelease(cf_data as *const libc::c_void);

            Ok(Screenshot { data: pixels })
        }
    }
}

#[cfg(target_os = "windows")]
mod ffi {
    #![allow(non_snake_case, dead_code)]
    use super::{Pixel, ScreenResult, Screenshot};
    use std::ffi::c_void;
    use std::mem::size_of;
    use std::os::raw::{c_int, c_long, c_uint};

    type PVOID = *mut c_void;
    type LPVOID = *mut c_void;
    type WORD = u16;
    type DWORD = u32;
    type BOOL = c_int;
    type BYTE = u8;
    type UINT = c_uint;
    type LONG = c_long;
    type LPARAM = c_long;

    #[repr(C)]
    struct RECT {
        left: LONG,
        top: LONG,
        right: LONG,
        bottom: LONG,
    }
    type LPCRECT = *const RECT;
    type LPRECT = *mut RECT;

    type HANDLE = PVOID;
    type HMONITOR = HANDLE;
    type HWND = HANDLE;
    type HDC = HANDLE;

    type HBITMAP = HANDLE;
    type HGDIOBJ = HANDLE;
    type LPBITMAPINFO = PVOID;

    const NULL: *mut c_void = 0usize as *mut c_void;
    const HGDI_ERROR: *mut c_void = -1isize as *mut c_void;
    const SM_CXSCREEN: c_int = 0;
    const SM_CYSCREEN: c_int = 1;

    const SRCCOPY: u32 = 0x00CC0020;
    const CAPTUREBLT: u32 = 0x40000000;
    const DIB_RGB_COLORS: UINT = 0;
    const BI_RGB: DWORD = 0;

    #[repr(C)]
    struct BITMAPINFOHEADER {
        biSize: DWORD,
        biWidth: LONG,
        biHeight: LONG,
        biPlanes: WORD,
        biBitCount: WORD,
        biCompression: DWORD,
        biSizeImage: DWORD,
        biXPelsPerMeter: LONG,
        biYPelsPerMeter: LONG,
        biClrUsed: DWORD,
        biClrImportant: DWORD,
    }

    #[repr(C)]
    struct RGBQUAD {
        rgbBlue: BYTE,
        rgbGreen: BYTE,
        rgbRed: BYTE,
        rgbReserved: BYTE,
    }

    #[repr(C)]
    struct BITMAPINFO {
        bmiHeader: BITMAPINFOHEADER,
        bmiColors: [RGBQUAD; 1],
    }

    #[link(name = "user32")]
    extern "system" {
        fn GetSystemMetrics(m: c_int) -> c_int;
        fn GetDesktopWindow() -> HWND;
        fn GetDC(hWnd: HWND) -> HDC;
    }

    #[link(name = "gdi32")]
    extern "system" {
        fn CreateCompatibleDC(hdc: HDC) -> HDC;
        fn CreateCompatibleBitmap(hdc: HDC, nWidth: c_int, nHeight: c_int) -> HBITMAP;
        fn SelectObject(hdc: HDC, hgdiobj: HGDIOBJ) -> HGDIOBJ;
        fn BitBlt(
            hdcDest: HDC,
            nXDest: c_int,
            nYDest: c_int,
            nWidth: c_int,
            nHeight: c_int,
            hdcSrc: HDC,
            nXSrc: c_int,
            nYSrc: c_int,
            dwRop: DWORD,
        ) -> BOOL;
        fn GetDIBits(
            hdc: HDC,
            hbmp: HBITMAP,
            uStartScan: UINT,
            cScanLines: UINT,
            lpvBits: LPVOID,
            lpbi: LPBITMAPINFO,
            uUsage: UINT,
        ) -> c_int;
        fn DeleteObject(hObject: HGDIOBJ) -> BOOL;
        fn ReleaseDC(hWnd: HWND, hDC: HDC) -> c_int;
        fn DeleteDC(hdc: HDC) -> BOOL;
    }

    pub fn get_screenshot_inner(_screen: usize) -> ScreenResult<Screenshot> {
        unsafe {
            let h_wnd_screen = GetDesktopWindow();
            let h_dc_screen = GetDC(h_wnd_screen);
            let width = GetSystemMetrics(SM_CXSCREEN);
            let height = GetSystemMetrics(SM_CYSCREEN);

            let h_dc = CreateCompatibleDC(h_dc_screen);
            if h_dc == NULL {
                return Err(anyhow::anyhow!("Can't get a Windows display."));
            }

            let h_bmp = CreateCompatibleBitmap(h_dc_screen, width, height);
            if h_bmp == NULL {
                return Err(anyhow::anyhow!("Can't create a Windows buffer"));
            }

            let res = SelectObject(h_dc, h_bmp);
            if res == NULL || res == HGDI_ERROR {
                return Err(anyhow::anyhow!("Can't select Windows buffer."));
            }

            let res = BitBlt(
                h_dc,
                0,
                0,
                width,
                height,
                h_dc_screen,
                0,
                0,
                SRCCOPY | CAPTUREBLT,
            );
            if res == 0 {
                return Err(anyhow::anyhow!("Failed to copy screen to Windows buffer"));
            }

            let pixel_width: usize = 4;
            let mut bmi = BITMAPINFO {
                bmiHeader: BITMAPINFOHEADER {
                    biSize: size_of::<BITMAPINFOHEADER>() as DWORD,
                    biWidth: width as LONG,
                    biHeight: height as LONG,
                    biPlanes: 1,
                    biBitCount: 8 * pixel_width as WORD,
                    biCompression: BI_RGB,
                    biSizeImage: (width * height * pixel_width as c_int) as DWORD,
                    biXPelsPerMeter: 0,
                    biYPelsPerMeter: 0,
                    biClrUsed: 0,
                    biClrImportant: 0,
                },
                bmiColors: [RGBQUAD {
                    rgbBlue: 0,
                    rgbGreen: 0,
                    rgbRed: 0,
                    rgbReserved: 0,
                }],
            };

            let size: usize = (width * height) as usize * pixel_width;
            let mut data: Vec<u8> = vec![0; size];

            GetDIBits(
                h_dc,
                h_bmp,
                0,
                height as DWORD,
                data.as_mut_ptr() as *mut c_void,
                &mut bmi as *mut BITMAPINFO as *mut c_void,
                DIB_RGB_COLORS,
            );

            ReleaseDC(h_wnd_screen, h_dc_screen);
            DeleteDC(h_dc);
            DeleteObject(h_bmp);

            let row_len = width as usize * pixel_width;
            let mut pixels: Vec<Pixel> = Vec::with_capacity((width * height) as usize);

            for row_i in (0..height as usize).rev() {
                for byte_i in (0..width as usize) {
                    let idx = row_i * row_len + byte_i * pixel_width;
                    pixels.push(Pixel {
                        b: data[idx],
                        g: data[idx + 1],
                        r: data[idx + 2],
                        a: data[idx + 3],
                    });
                }
            }

            Ok(Screenshot { data: pixels })
        }
    }
}

pub fn get_screenshot(screen: usize) -> ScreenResult<Screenshot> {
    ffi::get_screenshot_inner(screen)
}