forgewright 0.2.0

//! DXGI Desktop Duplication capture for forgewright.
//! Captures GPU-rendered content from DirectX/Vulkan applications.
//! Falls back to GDI BitBlt when DXGI is unavailable.
//!
//! Windows-only module — real Desktop Duplication.

#[cfg(target_os = "windows")]
use std::path::Path;
#[cfg(target_os = "windows")]
use serde_json::{json, Value};

#[cfg(target_os = "windows")]
use windows::core::Interface;
#[cfg(target_os = "windows")]
use windows::Win32::Graphics::Direct3D::D3D_DRIVER_TYPE_HARDWARE;
#[cfg(target_os = "windows")]
use windows::Win32::Graphics::Direct3D::D3D_FEATURE_LEVEL_11_0;
#[cfg(target_os = "windows")]
use windows::Win32::Graphics::Direct3D11::{
    D3D11CreateDevice, D3D11_CPU_ACCESS_READ, D3D11_CREATE_DEVICE_FLAG,
    D3D11_MAP_READ, D3D11_SDK_VERSION, D3D11_TEXTURE2D_DESC, D3D11_USAGE_STAGING,
    ID3D11Device, ID3D11DeviceContext, ID3D11Texture2D,
};
#[cfg(target_os = "windows")]
use windows::Win32::Graphics::Dxgi::{
    IDXGIAdapter, IDXGIDevice, IDXGIOutput, IDXGIOutput1, IDXGIOutputDuplication,
    IDXGIResource,
};
#[cfg(target_os = "windows")]
use windows::Win32::Graphics::Dxgi::Common::{
    DXGI_FORMAT_B8G8R8A8_UNORM, DXGI_SAMPLE_DESC,
};

// ═══ ERROR TYPE ══════════════════════════════════════════════════════════════

/// Error types for DXGI capture operations.
#[derive(Debug)]
pub enum DxgiError {
    /// DXGI Desktop Duplication not available on this system
    Unavailable,
    /// Another consumer already holds the output duplication interface
    AccessLost,
    /// Capture timed out waiting for a new frame
    Timeout,
    /// Win32/COM error from the windows crate
    #[cfg(target_os = "windows")]
    Win32(windows::core::Error),
    /// Other error with description
    Other(String),
}

impl std::fmt::Display for DxgiError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            DxgiError::Unavailable => write!(f, "DXGI Desktop Duplication not available"),
            DxgiError::AccessLost  => write!(f, "DXGI output duplication access lost"),
            DxgiError::Timeout     => write!(f, "DXGI frame capture timed out"),
            #[cfg(target_os = "windows")]
            DxgiError::Win32(e)    => write!(f, "DXGI Win32 error: {}", e),
            DxgiError::Other(s)    => write!(f, "DXGI error: {}", s),
        }
    }
}

impl std::error::Error for DxgiError {}

#[cfg(target_os = "windows")]
impl From<windows::core::Error> for DxgiError {
    fn from(e: windows::core::Error) -> Self {
        let code = e.code();
        // DXGI_ERROR_NOT_CURRENTLY_AVAILABLE
        if code.0 as u32 == 0x887A0022u32 {
            return DxgiError::Unavailable;
        }
        // DXGI_ERROR_ACCESS_LOST
        if code.0 as u32 == 0x887A0026u32 {
            return DxgiError::AccessLost;
        }
        // DXGI_ERROR_WAIT_TIMEOUT
        if code.0 as u32 == 0x887A0027u32 {
            return DxgiError::Timeout;
        }
        // E_ACCESSDENIED (common in RDP/service sessions)
        if code.0 as u32 == 0x80070005u32 {
            return DxgiError::Unavailable;
        }
        DxgiError::Win32(e)
    }
}

// ═══ RAW CAPTURE ═════════════════════════════════════════════════════════════

/// Raw pixel data from a DXGI capture (BGRA, 8bpc).
/// `data` contains tightly packed pixels with NO stride padding.
pub struct RawCapture {
    pub width: u32,
    pub height: u32,
    /// Stride is always `width * 4` (no padding in output).
    pub stride: u32,
    pub data: Vec<u8>,
}

// ═══ CAPTURE RESULT ══════════════════════════════════════════════════════════

/// Result of a successful capture operation.
pub struct CaptureResult {
    /// Path where the captured image was saved
    pub path: String,
    /// "gdi" if DXGI was unavailable and GDI fallback was used
    pub fallback: Option<String>,
    /// Width of the captured image in pixels
    pub width: u32,
    /// Height of the captured image in pixels
    pub height: u32,
}

// ═══ DXGI CAPTURER ═══════════════════════════════════════════════════════════

/// DXGI Desktop Duplication capturer.
///
/// Holds the D3D11 device and output duplication interface.
/// Reusable across multiple captures (avoids re-creating the device each time).
///
/// COM objects are released automatically when dropped (windows crate RAII).
/// The explicit `Drop` impl ensures correct cleanup ordering:
/// duplication first, then staging, context, device.
/// Fields are ordered so Rust's default drop order (declaration order)
/// releases resources in the correct dependency sequence:
/// duplication → staging → context → device.
#[cfg(target_os = "windows")]
pub struct DxgiCapturer {
    duplication: IDXGIOutputDuplication,
    staging: ID3D11Texture2D,
    context: ID3D11DeviceContext,
    device: ID3D11Device,
    pub width: u32,
    pub height: u32,
}

#[cfg(target_os = "windows")]
impl DxgiCapturer {
    /// Initialize DXGI capture for the primary output.
    ///
    /// Creates a D3D11 device, walks the DXGI chain to obtain an
    /// IDXGIOutputDuplication interface, and allocates a staging texture
    /// for CPU readback.
    ///
    /// Returns `DxgiError::Unavailable` if Desktop Duplication is not supported
    /// (Windows 7, RDP session, no GPU, etc.).
    pub fn new() -> Result<Self, DxgiError> {
        unsafe {
            // Step 1: Create D3D11 device (default GPU)
            let mut device: Option<ID3D11Device> = None;
            let mut context: Option<ID3D11DeviceContext> = None;
            let feature_levels = [D3D_FEATURE_LEVEL_11_0];

            D3D11CreateDevice(
                None, // default adapter
                D3D_DRIVER_TYPE_HARDWARE,
                None, // no software rasterizer
                D3D11_CREATE_DEVICE_FLAG(0),
                Some(&feature_levels),
                D3D11_SDK_VERSION,
                Some(&mut device),
                None, // don't need actual feature level out
                Some(&mut context),
            )?;

            let device = device.ok_or(DxgiError::Unavailable)?;
            let context = context.ok_or(DxgiError::Unavailable)?;

            // Step 2: Walk DXGI chain — Device → Adapter → Output → Output1
            let dxgi_device: IDXGIDevice = device.cast()?;
            let adapter: IDXGIAdapter = dxgi_device.GetAdapter()?;

            // EnumOutputs(0) = primary monitor
            let output: IDXGIOutput = adapter.EnumOutputs(0).map_err(|_| {
                DxgiError::Unavailable
            })?;

            let output1: IDXGIOutput1 = output.cast().map_err(|_| DxgiError::Unavailable)?;

            // Step 3: Get output dimensions from DXGI_OUTPUT_DESC
            let desc = output.GetDesc()?;
            let width = (desc.DesktopCoordinates.right - desc.DesktopCoordinates.left) as u32;
            let height = (desc.DesktopCoordinates.bottom - desc.DesktopCoordinates.top) as u32;

            if width == 0 || height == 0 {
                return Err(DxgiError::Unavailable);
            }

            // Step 4: Create duplication interface
            let duplication = output1.DuplicateOutput(&device)?;

            // Step 5: Create staging texture for CPU readback
            let tex_desc = D3D11_TEXTURE2D_DESC {
                Width: width,
                Height: height,
                MipLevels: 1,
                ArraySize: 1,
                Format: DXGI_FORMAT_B8G8R8A8_UNORM,
                SampleDesc: DXGI_SAMPLE_DESC { Count: 1, Quality: 0 },
                Usage: D3D11_USAGE_STAGING,
                BindFlags: 0,
                CPUAccessFlags: D3D11_CPU_ACCESS_READ.0 as u32,
                MiscFlags: 0,
            };

            let mut staging: Option<ID3D11Texture2D> = None;
            device.CreateTexture2D(&tex_desc, None, Some(&mut staging))?;
            let staging = staging.ok_or(DxgiError::Other(
                "failed to create staging texture".into(),
            ))?;

            Ok(DxgiCapturer {
                duplication,
                staging,
                context,
                device,
                width,
                height,
            })
        }
    }

    /// Capture a single frame. Acquires the next desktop frame from DXGI,
    /// copies it to the staging texture, maps for CPU read, and returns
    /// raw BGRA pixels with stride padding removed.
    ///
    /// `timeout_ms`: how long to wait for a new frame (100ms typical).
    ///
    /// Returns `DxgiError::Timeout` when no frame available within timeout.
    /// Returns `DxgiError::AccessLost` when the duplication interface is
    /// invalidated (display mode change, resolution change, etc.).
    pub fn capture_frame(&mut self, timeout_ms: u32) -> Result<RawCapture, DxgiError> {
        unsafe {
            // Step 1: Acquire next frame from Desktop Duplication
            let mut frame_info = std::mem::zeroed();
            let mut resource: Option<IDXGIResource> = None;

            self.duplication.AcquireNextFrame(
                timeout_ms,
                &mut frame_info,
                &mut resource,
            )?;

            let resource = resource.ok_or(DxgiError::Other(
                "AcquireNextFrame returned null resource".into(),
            ))?;

            // Step 2: Get the desktop texture from the resource
            let desktop_texture: ID3D11Texture2D = resource.cast()?;

            // Step 3: Copy desktop texture to staging texture
            self.context.CopyResource(&self.staging, &desktop_texture);

            // Step 4: Map staging texture for CPU read
            let mut mapped = std::mem::zeroed();
            self.context.Map(
                &self.staging,
                0,
                D3D11_MAP_READ,
                0,
                Some(&mut mapped),
            )?;

            // Step 5: Copy pixel data, discarding row pitch padding
            let row_pitch = mapped.RowPitch;
            let expected_pitch = self.width * 4;
            let total_bytes = (self.width * self.height * 4) as usize;
            let mut data = Vec::with_capacity(total_bytes);

            let src_base = mapped.pData as *const u8;
            for y in 0..self.height {
                let src = src_base.add((y * row_pitch) as usize);
                let slice = std::slice::from_raw_parts(src, expected_pitch as usize);
                data.extend_from_slice(slice);
            }

            // Step 6: Unmap staging texture
            self.context.Unmap(&self.staging, 0);

            // Step 7: Release the acquired frame
            self.duplication.ReleaseFrame()?;

            Ok(RawCapture {
                width: self.width,
                height: self.height,
                stride: expected_pitch,
                data,
            })
        }
    }

    /// Capture and save directly to a 24-bit BGR bottom-up BMP file.
    ///
    /// This is a convenience wrapper around `capture_frame()` that writes
    /// the result in the format that `bridge.rs` `decode_bmp` expects.
    pub fn capture_to_bmp(&mut self, path: &Path) -> Result<CaptureResult, DxgiError> {
        let raw = self.capture_frame(100)?;
        write_bgr_bmp(path, &raw)?;
        Ok(CaptureResult {
            path: path.to_string_lossy().to_string(),
            fallback: None,
            width: raw.width,
            height: raw.height,
        })
    }
}

/// RAII Drop for DxgiCapturer.
///
/// The struct fields are ordered so Rust's default drop sequence
/// (declaration order) releases resources correctly:
///   1. duplication — releases the output duplication lock
///   2. staging     — frees the CPU-readable texture
///   3. context     — releases the device context
///   4. device      — releases the D3D11 device (root object, last)
///
/// The windows crate COM wrappers are reference-counted and safe to
/// drop in any order, but this ordering is defensive best-practice.
#[cfg(target_os = "windows")]
impl Drop for DxgiCapturer {
    fn drop(&mut self) {
        // Rust drops fields in declaration order, which matches our
        // desired cleanup sequence. Nothing extra needed here.
        // This explicit Drop impl documents the intent and satisfies
        // the RAII requirement from the spec.
    }
}

// ═══ BMP WRITER ══════════════════════════════════════════════════════════════

/// Write a RawCapture (BGRA) as a 24-bit BGR bottom-up BMP file.
/// This matches the format that bridge.rs decode_bmp expects.
#[cfg(target_os = "windows")]
fn write_bgr_bmp(path: &Path, raw: &RawCapture) -> Result<(), DxgiError> {
    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent)
            .map_err(|e| DxgiError::Other(e.to_string()))?;
    }

    let w = raw.width;
    let h = raw.height;
    let row_bytes = ((w * 3 + 3) / 4) * 4; // padded to 4-byte boundary
    let pixel_data_size = row_bytes * h;
    let file_size = 54 + pixel_data_size;

    let mut buf = Vec::with_capacity(file_size as usize);

    // BMP file header (14 bytes)
    buf.extend_from_slice(b"BM");
    buf.extend_from_slice(&file_size.to_le_bytes());
    buf.extend_from_slice(&[0u8; 4]); // reserved
    buf.extend_from_slice(&54u32.to_le_bytes()); // pixel data offset

    // BITMAPINFOHEADER (40 bytes)
    buf.extend_from_slice(&40u32.to_le_bytes());
    buf.extend_from_slice(&(w as i32).to_le_bytes());
    buf.extend_from_slice(&(h as i32).to_le_bytes()); // positive = bottom-up
    buf.extend_from_slice(&1u16.to_le_bytes()); // planes
    buf.extend_from_slice(&24u16.to_le_bytes()); // bits per pixel
    buf.extend_from_slice(&0u32.to_le_bytes()); // compression
    buf.extend_from_slice(&pixel_data_size.to_le_bytes());
    buf.extend_from_slice(&2835u32.to_le_bytes()); // ppm_x
    buf.extend_from_slice(&2835u32.to_le_bytes()); // ppm_y
    buf.extend_from_slice(&0u32.to_le_bytes()); // colors used
    buf.extend_from_slice(&0u32.to_le_bytes()); // colors important

    // Pixel data: bottom-up BGR with row padding
    let pad = (row_bytes - w * 3) as usize;
    for row in (0..h as usize).rev() {
        for col in 0..w as usize {
            let px = (row * w as usize + col) * 4;
            buf.push(raw.data[px]);     // B
            buf.push(raw.data[px + 1]); // G
            buf.push(raw.data[px + 2]); // R
        }
        for _ in 0..pad {
            buf.push(0);
        }
    }

    std::fs::write(path, &buf).map_err(|e| DxgiError::Other(e.to_string()))?;
    Ok(())
}

// ═══ PUBLIC API (backwards compatible) ═══════════════════════════════════════

/// Attempt DXGI capture of the full desktop output, save to path.
/// Falls back to GDI BitBlt if DXGI is unavailable.
///
/// Returns CaptureResult with fallback="gdi" if GDI was used.
#[cfg(target_os = "windows")]
pub fn capture_screen(path: &Path) -> Result<CaptureResult, DxgiError> {
    match try_dxgi_capture(path) {
        Ok(result) => Ok(result),
        Err(DxgiError::Unavailable) | Err(DxgiError::Other(_)) => {
            gdi_screen_capture(path)
        }
        Err(e) => Err(e),
    }
}

/// Attempt DXGI Desktop Duplication capture via a one-shot DxgiCapturer.
#[cfg(target_os = "windows")]
fn try_dxgi_capture(path: &Path) -> Result<CaptureResult, DxgiError> {
    let mut capturer = DxgiCapturer::new()?;
    capturer.capture_to_bmp(path)
}

/// GDI screen capture fallback using BitBlt.
#[cfg(target_os = "windows")]
fn gdi_screen_capture(path: &Path) -> Result<CaptureResult, DxgiError> {
    use windows::Win32::Graphics::Gdi::{
        BitBlt, CreateCompatibleBitmap, CreateCompatibleDC, DeleteDC, DeleteObject,
        GetDC, GetDIBits, ReleaseDC, SelectObject,
        BITMAPINFO, BITMAPINFOHEADER, BI_RGB, DIB_RGB_COLORS, SRCCOPY,
    };
    use windows::Win32::UI::WindowsAndMessaging::{GetSystemMetrics, SM_CXSCREEN, SM_CYSCREEN};
    use std::mem;

    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent).ok();
    }

    unsafe {
        let w = GetSystemMetrics(SM_CXSCREEN);
        let h = GetSystemMetrics(SM_CYSCREEN);
        if w <= 0 || h <= 0 {
            return Err(DxgiError::Other("zero screen dimensions".into()));
        }

        let hdc_screen = GetDC(None);
        let hdc_mem = CreateCompatibleDC(hdc_screen);
        let bmp = CreateCompatibleBitmap(hdc_screen, w, h);
        let old = SelectObject(hdc_mem, bmp);
        let _ = BitBlt(hdc_mem, 0, 0, w, h, hdc_screen, 0, 0, SRCCOPY);

        let row = ((w * 3 + 3) / 4) * 4;
        let size = (row * h) as usize;
        let mut px = vec![0u8; size];
        let mut bmi = BITMAPINFO {
            bmiHeader: BITMAPINFOHEADER {
                biSize: mem::size_of::<BITMAPINFOHEADER>() as u32,
                biWidth: w, biHeight: h, biPlanes: 1, biBitCount: 24,
                biCompression: BI_RGB.0 as u32, biSizeImage: size as u32,
                ..Default::default()
            },
            ..Default::default()
        };
        GetDIBits(hdc_mem, bmp, 0, h as u32, Some(px.as_mut_ptr() as _), &mut bmi, DIB_RGB_COLORS);

        SelectObject(hdc_mem, old);
        let _ = DeleteObject(bmp);
        let _ = DeleteDC(hdc_mem);
        ReleaseDC(None, hdc_screen);

        // Write BMP file
        let file_size = 54 + size as u32;
        let mut buf = Vec::with_capacity(file_size as usize);
        buf.extend_from_slice(b"BM");
        buf.extend_from_slice(&file_size.to_le_bytes());
        buf.extend_from_slice(&[0u8; 4]);
        buf.extend_from_slice(&54u32.to_le_bytes());
        buf.extend_from_slice(&40u32.to_le_bytes());
        buf.extend_from_slice(&w.to_le_bytes());
        buf.extend_from_slice(&h.to_le_bytes());
        buf.extend_from_slice(&1u16.to_le_bytes());
        buf.extend_from_slice(&24u16.to_le_bytes());
        buf.extend_from_slice(&0u32.to_le_bytes());
        buf.extend_from_slice(&(size as u32).to_le_bytes());
        buf.extend_from_slice(&2835u32.to_le_bytes());
        buf.extend_from_slice(&2835u32.to_le_bytes());
        buf.extend_from_slice(&0u32.to_le_bytes());
        buf.extend_from_slice(&0u32.to_le_bytes());
        buf.extend_from_slice(&px);
        std::fs::write(path, &buf)
            .map_err(|e| DxgiError::Other(e.to_string()))?;

        Ok(CaptureResult {
            path: path.to_string_lossy().to_string(),
            fallback: Some("gdi".to_string()),
            width: w as u32,
            height: h as u32,
        })
    }
}

/// Capture and return JSON result suitable for drive response.
#[cfg(target_os = "windows")]
pub fn capture_to_json(path: &Path) -> Value {
    match capture_screen(path) {
        Ok(result) => {
            let mut obj = json!({
                "status": "ok",
                "screenshot": result.path,
                "width": result.width,
                "height": result.height,
            });
            if let Some(fb) = result.fallback {
                obj["fallback"] = json!(fb);
            }
            obj
        }
        Err(DxgiError::AccessLost) => json!({
            "status": "error",
            "error": "DXGI output duplication access lost — another consumer is active",
            "code": "dxgi_access_lost",
        }),
        Err(e) => json!({
            "status": "error",
            "error": e.to_string(),
        }),
    }
}

// ═══ TESTS ═══════════════════════════════════════════════════════════════════

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

    /// **Property 5: BGRA data integrity**
    ///
    /// For any RawCapture produced by capture, verify
    /// `data.len() == width * height * 4` with no stride padding.
    ///
    /// Note: Uses synthetic RawCapture data since DXGI requires a live display.
    /// Generates arbitrary (width, height) pairs and verifies the invariant
    /// on constructed RawCapture structs.
    proptest! {
        #[test]
        fn prop_bgra_data_integrity(
            width in 1u32..=4096,
            height in 1u32..=4096,
            seed in any::<u8>(),
        ) {
            // Construct a RawCapture the same way capture_frame() does:
            // tightly packed BGRA with stride = width * 4, no padding.
            let expected_stride = width * 4;
            let total_bytes = (width as usize) * (height as usize) * 4;

            // Generate synthetic pixel data
            let data: Vec<u8> = (0..total_bytes)
                .map(|i| ((i & 0xFF) as u8).wrapping_add(seed))
                .collect();

            let raw = RawCapture {
                width,
                height,
                stride: expected_stride,
                data,
            };

            // Property: data length equals width * height * 4 exactly
            prop_assert_eq!(
                raw.data.len(),
                (raw.width as usize) * (raw.height as usize) * 4,
                "data.len() must equal width * height * 4"
            );

            // Property: stride equals width * 4 (no padding)
            prop_assert_eq!(
                raw.stride,
                raw.width * 4,
                "stride must equal width * 4"
            );

            // Property: no extra bytes beyond the expected count
            prop_assert_eq!(
                raw.data.len() % 4,
                0,
                "data length must be a multiple of 4 (BGRA)"
            );
        }
    }
}