doe 1.1.89

//! Windows screen capture via GDI (compatible with `windows` crate 0.48).
//!
//! Enumerates displays with `EnumDisplayMonitors` and captures pixels with the
//! classic `CreateCompatibleDC` + `StretchBlt` + `GetDIBits` pipeline,
//! producing 32-bit BGRA data that is then converted to RGBA.

use super::image_utils::bgra_to_rgba;
use super::Screen;
use image::RgbaImage;
use std::mem;
use std::ops::Deref;
use std::ptr;
use windows::core::PCWSTR;
use windows::Win32::Foundation::{BOOL, LPARAM, RECT};
use windows::Win32::Graphics::Gdi::{
    CreateCompatibleBitmap, CreateCompatibleDC, CreateDCW, DeleteDC, DeleteObject,
    EnumDisplayMonitors, GetDIBits, GetMonitorInfoW, GetObjectW, SelectObject,
    SetStretchBltMode, StretchBlt, BITMAP, BITMAPINFO, BITMAPINFOHEADER, DIB_RGB_COLORS, HBITMAP,
    HDC, HGDIOBJ, HMONITOR, MONITORINFOEXW, RGBQUAD, SRCCOPY, STRETCH_HALFTONE,
};

// ------------------------------------------------------------------
// RAII GDI handle wrapper
// ------------------------------------------------------------------

struct DropBox<T, F: FnMut(T)> {
    value: Option<T>,
    drop_fn: Option<F>,
}

impl<T, F: FnMut(T)> DropBox<T, F> {
    fn new(value: T, drop_fn: F) -> Self {
        Self { value: Some(value), drop_fn: Some(drop_fn) }
    }
}

impl<T, F: FnMut(T)> Deref for DropBox<T, F> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        self.value.as_ref().expect("DropBox consumed")
    }
}

impl<T, F: FnMut(T)> Drop for DropBox<T, F> {
    fn drop(&mut self) {
        if let (Some(v), Some(mut f)) = (self.value.take(), self.drop_fn.take()) {
            f(v);
        }
    }
}

// ------------------------------------------------------------------
// Monitor enumeration
// ------------------------------------------------------------------

struct MonitorEntry {
    id:   u32,
    rect: RECT,
}

pub fn all() -> Result<Vec<Screen>, String> {
    let collected: *mut Vec<MonitorEntry> = Box::into_raw(Box::default());
    let lparam = LPARAM(collected as isize);

    // windows 0.48: EnumDisplayMonitors returns BOOL (not Result).
    let _ = unsafe {
        EnumDisplayMonitors(
            HDC::default(),
            None,
            Some(monitor_enum_proc),
            lparam,
        )
    };

    let entries = unsafe { *Box::from_raw(collected) };
    Ok(entries
        .into_iter()
        .map(|e| Screen {
            id: e.id,
            x: e.rect.left,
            y: e.rect.top,
            width: (e.rect.right - e.rect.left) as u32,
            height: (e.rect.bottom - e.rect.top) as u32,
            scale_factor: 1.0,
        })
        .collect())
}

extern "system" fn monitor_enum_proc(
    h_monitor: HMONITOR,
    _hdc: HDC,
    _rect: *mut RECT,
    state: LPARAM,
) -> BOOL {
    let state = unsafe { &mut *(state.0 as *mut Vec<MonitorEntry>) };
    // Always continue enumeration even if a single monitor fails.
    match monitor_info(h_monitor) {
        Some((id, rect)) => {
            state.push(MonitorEntry { id, rect });
        }
        None => {}  // skip this monitor, keep enumerating
    }
    BOOL(1)
}

/// Read `MONITORINFOEXW` and return the FNV-1a hash of the device name
/// together with the monitor's rectangle.
fn monitor_info(h_monitor: HMONITOR) -> Option<(u32, RECT)> {
    unsafe {
        let mut info: MONITORINFOEXW = mem::zeroed();
        info.monitorInfo.cbSize = mem::size_of::<MONITORINFOEXW>() as u32;
        let ptr = <*mut _>::cast(&mut info);
        // windows 0.48: GetMonitorInfoW returns BOOL; BOOL::ok() → Result<(), Error>
        // We ignore the error and just return None on failure.
        if GetMonitorInfoW(h_monitor, ptr).as_bool() {
            let name = device_name(&info);
            let id = fnv1a(name.as_bytes());
            Some((id, info.monitorInfo.rcMonitor))
        } else {
            None
        }
    }
}

/// Extract the monitor device name (e.g. `\\.\DISPLAY1`) as a Rust String.
fn device_name(info: &MONITORINFOEXW) -> String {
    let u16s: Vec<u16> = info.szDevice.iter().copied().take_while(|&c| c != 0).collect();
    String::from_utf16_lossy(&u16s).trim_end_matches('\0').to_string()
}

/// FNV-1a 32-bit hash (cheap, deterministic, sufficient for monitor ids).
fn fnv1a(bytes: &[u8]) -> u32 {
    let mut h: u32 = 0x811c9dc5;
    for &b in bytes {
        h ^= b as u32;
        h = h.wrapping_mul(0x01000193);
    }
    h
}

// ------------------------------------------------------------------
// Per-monitor lookup (re-enumerate to resolve id → MONITORINFOEXW)
// ------------------------------------------------------------------

fn monitor_by_id(id: u32) -> Result<MONITORINFOEXW, String> {
    let collected: *mut Vec<MONITORINFOEXW> = Box::into_raw(Box::default());
    let lparam = LPARAM(collected as isize);
    unsafe {
        EnumDisplayMonitors(HDC::default(), None, Some(collect_raw_proc), lparam);
    }
    let infos = unsafe { *Box::from_raw(collected) };
    infos
        .into_iter()
        .find(|i| fnv1a(device_name(i).as_bytes()) == id)
        .ok_or_else(|| format!("monitor with id {id} not found"))
}

extern "system" fn collect_raw_proc(
    h_monitor: HMONITOR,
    _hdc: HDC,
    _rect: *mut RECT,
    state: LPARAM,
) -> BOOL {
    let state = unsafe { &mut *(state.0 as *mut Vec<MONITORINFOEXW>) };
    unsafe {
        let mut info: MONITORINFOEXW = mem::zeroed();
        info.monitorInfo.cbSize = mem::size_of::<MONITORINFOEXW>() as u32;
        let ptr = <*mut _>::cast(&mut info);
        if GetMonitorInfoW(h_monitor, ptr).as_bool() {
            state.push(info);
        }
    }
    BOOL(1)
}

// ------------------------------------------------------------------
// Pixel capture
// ------------------------------------------------------------------

pub fn capture(screen: &Screen) -> Result<RgbaImage, String> {
    let w = (screen.width as f32 * screen.scale_factor) as i32;
    let h = (screen.height as f32 * screen.scale_factor) as i32;
    capture_rect(screen.id, 0, 0, w, h)
}

pub fn capture_area(
    screen: &Screen,
    x: i32,
    y: i32,
    width: u32,
    height: u32,
) -> Result<RgbaImage, String> {
    let w = (width as f32 * screen.scale_factor) as i32;
    let h = (height as f32 * screen.scale_factor) as i32;
    let ax = (x as f32 * screen.scale_factor) as i32;
    let ay = (y as f32 * screen.scale_factor) as i32;
    capture_rect(screen.id, ax, ay, w, h)
}

fn capture_rect(
    display_id: u32,
    x: i32,
    y: i32,
    width: i32,
    height: i32,
) -> Result<RgbaImage, String> {
    if width <= 0 || height <= 0 {
        return Err(format!("invalid capture size {width}x{height}"));
    }
    let monitor_info = monitor_by_id(display_id)?;
    let device = device_name(&monitor_info);
    let mut wide: Vec<u16> = device.encode_utf16().collect();
    wide.push(0);

    // Source DC for the monitor.
    #[allow(non_camel_case_types)]
    type CreatedHDC = windows::Win32::Graphics::Gdi::CreatedHDC;

    let src_dc = unsafe {
        CreateDCW(
            PCWSTR(wide.as_ptr()),
            PCWSTR(wide.as_ptr()),
            PCWSTR(ptr::null()),
            None,
        )
    };
    // Check validity through the inner HANDLE.
    if HDC(src_dc.0).is_invalid() {
        return Err("CreateDCW failed".to_string());
    }
    let src_dc = DropBox::new(src_dc, |dc| unsafe { let _ = DeleteDC(dc); });

    // Memory DC + compatible bitmap.
    let mem_dc = unsafe { CreateCompatibleDC(*src_dc) };
    if HDC(mem_dc.0).is_invalid() {
        return Err("CreateCompatibleDC failed".to_string());
    }
    let mem_dc = DropBox::new(mem_dc, |dc| unsafe { let _ = DeleteDC(dc); });

    let bitmap = unsafe { CreateCompatibleBitmap(*src_dc, width, height) };
    if bitmap.is_invalid() {
        return Err("CreateCompatibleBitmap failed".to_string());
    }
    // Save the old bitmap so we can restore it before cleanup (GDI requirement).
    let old_bmp = unsafe { SelectObject(*mem_dc, bitmap) };

    // Wrap bitmap in DropBox AFTER SelectObject so DropBox cleanup can
    // restore the old bitmap first.
    struct BitmapGuard {
        mem_dc: HDC,
        bitmap: HBITMAP,
        old_bmp: HGDIOBJ,
    }
    impl Drop for BitmapGuard {
        fn drop(&mut self) {
            unsafe {
                // Restore original bitmap into DC before deleting anything.
                let _ = SelectObject(self.mem_dc, self.old_bmp);
                let _ = DeleteObject(self.bitmap);
            }
        }
    }

    let _guard = BitmapGuard { mem_dc: HDC(mem_dc.0), bitmap, old_bmp };

    unsafe {
        SetStretchBltMode(*mem_dc, STRETCH_HALFTONE);
        StretchBlt(*mem_dc, 0, 0, width, height, *src_dc, x, y, width, height, SRCCOPY);
    }

    // Read back as 32-bit (BGRA) DIB.
    let mut info = BITMAPINFO {
        bmiHeader: BITMAPINFOHEADER {
            biSize: mem::size_of::<BITMAPINFOHEADER>() as u32,
            biWidth: width,
            biHeight: height,
            biPlanes: 1,
            biBitCount: 32,
            biCompression: 0,
            biSizeImage: 0,
            biXPelsPerMeter: 0,
            biYPelsPerMeter: 0,
            biClrUsed: 0,
            biClrImportant: 0,
        },
        bmiColors: [RGBQUAD::default(); 1],
    };

    let mut data = vec![0u8; (width as usize) * (height as usize) * 4];
    let copied = unsafe {
        GetDIBits(
            *mem_dc,
            _guard.bitmap,
            0,
            height as u32,
            Some(data.as_mut_ptr() as *mut _),
            &mut info,
            DIB_RGB_COLORS,
        )
    };
    if copied == 0 {
        return Err("GetDIBits returned 0".to_string());
    }

    // GDI bitmaps are bottom-up. Flip row order before colour-swapping.
    let row_len = (width as usize) * 4;
    let mut rows: Vec<Vec<u8>> = data.chunks(row_len).map(|r| r.to_vec()).collect();
    rows.reverse();
    let flat: Vec<u8> = rows.into_iter().flatten().collect();

    // Recover the real bitmap dimensions (defensive).
    let mut bmp: BITMAP = unsafe { mem::zeroed() };
    unsafe {
        GetObjectW(
            _guard.bitmap,
            mem::size_of::<BITMAP>() as i32,
            Some(<*mut _>::cast(&mut bmp)),
        );
    }
    let real_w = bmp.bmWidth.max(width) as u32;
    let real_h = bmp.bmHeight.max(height) as u32;
    bgra_to_rgba(real_w, real_h, flat)
}

// ------------------------------------------------------------------
// Tests
// ------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn fnv1a_is_stable() {
        // "\\.\DISPLAY1" as bytes
        let d1: &[u8] = &[b'\\', b'\\', b'.', b'\\', b'D', b'I', b'S', b'P', b'L', b'A', b'Y', b'1'];
        // "\\.\DISPLAY2" as bytes
        let d2: &[u8] = &[b'\\', b'\\', b'.', b'\\', b'D', b'I', b'S', b'P', b'L', b'A', b'Y', b'2'];
        assert_eq!(fnv1a(d1), fnv1a(d1));
        assert_ne!(fnv1a(d1), fnv1a(d2));
    }

    #[test]
    fn device_name_basic() {
        let mut info: MONITORINFOEXW = unsafe { mem::zeroed() };
        let s: Vec<u16> = "\\\\.\\DISPLAY1".encode_utf16().collect();
        for (i, c) in s.iter().enumerate() {
            info.szDevice[i] = *c;
        }
        assert_eq!(device_name(&info), "\\\\.\\DISPLAY1");
    }
}