guru-connect/agent/src/capture/dxgi.rs

//! DXGI Desktop Duplication screen capture
//!
//! Uses the Windows Desktop Duplication API (available on Windows 8+) for
//! high-performance, low-latency screen capture with hardware acceleration.
//!
//! Reference: RustDesk's scrap library implementation

use super::{CapturedFrame, Capturer, DirtyRect, Display};
use anyhow::{Context, Result};
use std::ptr;
use std::time::Instant;

use windows::core::Interface;
use windows::Win32::Graphics::Direct3D::D3D_DRIVER_TYPE_UNKNOWN;
use windows::Win32::Graphics::Direct3D11::{
    D3D11CreateDevice, ID3D11Device, ID3D11DeviceContext, ID3D11Texture2D,
    D3D11_MAPPED_SUBRESOURCE, D3D11_MAP_READ, D3D11_SDK_VERSION, D3D11_TEXTURE2D_DESC,
    D3D11_USAGE_STAGING,
};
use windows::Win32::Graphics::Dxgi::{
    CreateDXGIFactory1, IDXGIAdapter1, IDXGIFactory1, IDXGIOutput, IDXGIOutput1,
    IDXGIOutputDuplication, IDXGIResource, DXGI_ERROR_ACCESS_LOST, DXGI_ERROR_WAIT_TIMEOUT,
    DXGI_OUTDUPL_DESC, DXGI_OUTDUPL_FRAME_INFO, DXGI_RESOURCE_PRIORITY_MAXIMUM,
};

/// DXGI Desktop Duplication capturer
pub struct DxgiCapturer {
    display: Display,
    device: ID3D11Device,
    context: ID3D11DeviceContext,
    duplication: IDXGIOutputDuplication,
    staging_texture: Option<ID3D11Texture2D>,
    width: u32,
    height: u32,
    last_frame: Option<Vec<u8>>,
}

impl DxgiCapturer {
    /// Create a new DXGI capturer for the specified display
    pub fn new(display: Display) -> Result<Self> {
        let (device, context, duplication, desc) = Self::create_duplication(&display)?;

        Ok(Self {
            display,
            device,
            context,
            duplication,
            staging_texture: None,
            width: desc.ModeDesc.Width,
            height: desc.ModeDesc.Height,
            last_frame: None,
        })
    }

    /// Create D3D device and output duplication
    fn create_duplication(
        target_display: &Display,
    ) -> Result<(
        ID3D11Device,
        ID3D11DeviceContext,
        IDXGIOutputDuplication,
        DXGI_OUTDUPL_DESC,
    )> {
        unsafe {
            // Create DXGI factory
            let factory: IDXGIFactory1 =
                CreateDXGIFactory1().context("Failed to create DXGI factory")?;

            // Find the adapter and output for this display
            let (adapter, output) = Self::find_adapter_output(&factory, target_display)?;

            // Create D3D11 device
            let mut device: Option<ID3D11Device> = None;
            let mut context: Option<ID3D11DeviceContext> = None;

            D3D11CreateDevice(
                &adapter,
                D3D_DRIVER_TYPE_UNKNOWN,
                None,
                Default::default(),
                None,
                D3D11_SDK_VERSION,
                Some(&mut device),
                None,
                Some(&mut context),
            )
            .context("Failed to create D3D11 device")?;

            let device = device.context("D3D11 device is None")?;
            let context = context.context("D3D11 context is None")?;

            // Get IDXGIOutput1 interface
            let output1: IDXGIOutput1 = output
                .cast()
                .context("Failed to get IDXGIOutput1 interface")?;

            // Create output duplication
            let duplication = output1
                .DuplicateOutput(&device)
                .context("Failed to create output duplication")?;

            // Get duplication description
            let desc = duplication.GetDesc();

            tracing::info!(
                "Created DXGI duplication: {}x{}, display: {}",
                desc.ModeDesc.Width,
                desc.ModeDesc.Height,
                target_display.name
            );

            Ok((device, context, duplication, desc))
        }
    }

    /// Find the adapter and output for the specified display
    fn find_adapter_output(
        factory: &IDXGIFactory1,
        display: &Display,
    ) -> Result<(IDXGIAdapter1, IDXGIOutput)> {
        unsafe {
            let mut adapter_idx = 0u32;

            loop {
                // Enumerate adapters
                let adapter: IDXGIAdapter1 = match factory.EnumAdapters1(adapter_idx) {
                    Ok(a) => a,
                    Err(_) => break,
                };

                let mut output_idx = 0u32;

                loop {
                    // Enumerate outputs for this adapter
                    let output: IDXGIOutput = match adapter.EnumOutputs(output_idx) {
                        Ok(o) => o,
                        Err(_) => break,
                    };

                    // Check if this is the display we want
                    let desc = output.GetDesc()?;

                    let name = String::from_utf16_lossy(
                        &desc.DeviceName[..desc
                            .DeviceName
                            .iter()
                            .position(|&c| c == 0)
                            .unwrap_or(desc.DeviceName.len())],
                    );

                    if name == display.name || desc.Monitor.0 as isize == display.handle {
                        return Ok((adapter, output));
                    }

                    output_idx += 1;
                }

                adapter_idx += 1;
            }

            // If we didn't find the specific display, use the first one
            let adapter: IDXGIAdapter1 = factory.EnumAdapters1(0).context("No adapters found")?;
            let output: IDXGIOutput = adapter.EnumOutputs(0).context("No outputs found")?;

            Ok((adapter, output))
        }
    }

    /// Create or get the staging texture for CPU access
    fn get_staging_texture(&mut self, src_texture: &ID3D11Texture2D) -> Result<&ID3D11Texture2D> {
        if self.staging_texture.is_none() {
            unsafe {
                let mut desc = D3D11_TEXTURE2D_DESC::default();
                src_texture.GetDesc(&mut desc);

                desc.Usage = D3D11_USAGE_STAGING;
                desc.BindFlags = Default::default();
                desc.CPUAccessFlags = 0x20000; // D3D11_CPU_ACCESS_READ
                desc.MiscFlags = Default::default();

                let mut staging: Option<ID3D11Texture2D> = None;
                self.device
                    .CreateTexture2D(&desc, None, Some(&mut staging))
                    .context("Failed to create staging texture")?;

                let staging = staging.context("Staging texture is None")?;

                // Set high priority
                let resource: IDXGIResource = staging.cast()?;
                resource.SetEvictionPriority(DXGI_RESOURCE_PRIORITY_MAXIMUM)?;

                self.staging_texture = Some(staging);
            }
        }

        Ok(self.staging_texture.as_ref().unwrap())
    }

    /// Acquire the next frame from the desktop
    fn acquire_frame(
        &mut self,
        timeout_ms: u32,
    ) -> Result<Option<(ID3D11Texture2D, DXGI_OUTDUPL_FRAME_INFO)>> {
        unsafe {
            let mut frame_info = DXGI_OUTDUPL_FRAME_INFO::default();
            let mut desktop_resource: Option<IDXGIResource> = None;

            let result = self.duplication.AcquireNextFrame(
                timeout_ms,
                &mut frame_info,
                &mut desktop_resource,
            );

            match result {
                Ok(_) => {
                    let resource = desktop_resource.context("Desktop resource is None")?;

                    // Check if there's actually a new frame
                    if frame_info.LastPresentTime == 0 {
                        self.duplication.ReleaseFrame().ok();
                        return Ok(None);
                    }

                    let texture: ID3D11Texture2D = resource
                        .cast()
                        .context("Failed to cast to ID3D11Texture2D")?;

                    Ok(Some((texture, frame_info)))
                }
                Err(e) if e.code() == DXGI_ERROR_WAIT_TIMEOUT => {
                    // No new frame available
                    Ok(None)
                }
                Err(e) if e.code() == DXGI_ERROR_ACCESS_LOST => {
                    // Desktop duplication was invalidated, need to recreate
                    tracing::warn!("Desktop duplication access lost, will need to recreate");
                    Err(anyhow::anyhow!("Access lost"))
                }
                Err(e) => Err(e).context("Failed to acquire frame"),
            }
        }
    }

    /// Copy frame data to CPU-accessible memory
    fn copy_frame_data(&mut self, texture: &ID3D11Texture2D) -> Result<Vec<u8>> {
        unsafe {
            // Get or create staging texture
            let staging = self.get_staging_texture(texture)?.clone();

            // Copy from GPU texture to staging texture
            self.context.CopyResource(&staging, texture);

            // Map the staging texture for CPU read
            let mut mapped = D3D11_MAPPED_SUBRESOURCE::default();
            self.context
                .Map(&staging, 0, D3D11_MAP_READ, 0, Some(&mut mapped))
                .context("Failed to map staging texture")?;

            // Copy pixel data
            let src_pitch = mapped.RowPitch as usize;
            let dst_pitch = (self.width * 4) as usize;
            let height = self.height as usize;

            let mut data = vec![0u8; dst_pitch * height];

            let src_ptr = mapped.pData as *const u8;
            for y in 0..height {
                let src_row = src_ptr.add(y * src_pitch);
                let dst_row = data.as_mut_ptr().add(y * dst_pitch);
                ptr::copy_nonoverlapping(src_row, dst_row, dst_pitch);
            }

            // Unmap
            self.context.Unmap(&staging, 0);

            Ok(data)
        }
    }

    /// Extract dirty rectangles from frame info
    fn extract_dirty_rects(&self, _frame_info: &DXGI_OUTDUPL_FRAME_INFO) -> Option<Vec<DirtyRect>> {
        // TODO: Implement dirty rectangle extraction using
        // IDXGIOutputDuplication::GetFrameDirtyRects and GetFrameMoveRects
        // For now, return None to indicate full frame update
        None
    }
}

impl Capturer for DxgiCapturer {
    fn capture(&mut self) -> Result<Option<CapturedFrame>> {
        // Try to acquire a frame with 100ms timeout
        let frame_result = self.acquire_frame(100)?;

        let (texture, frame_info) = match frame_result {
            Some((t, f)) => (t, f),
            None => return Ok(None), // No new frame
        };

        // Copy frame data to CPU memory
        let data = self.copy_frame_data(&texture)?;

        // Release the frame
        unsafe {
            self.duplication.ReleaseFrame().ok();
        }

        // Extract dirty rectangles if available
        let dirty_rects = self.extract_dirty_rects(&frame_info);

        Ok(Some(CapturedFrame {
            width: self.width,
            height: self.height,
            data,
            timestamp: Instant::now(),
            display_id: self.display.id,
            dirty_rects,
        }))
    }

    fn display(&self) -> &Display {
        &self.display
    }

    fn is_valid(&self) -> bool {
        // Could check if duplication is still valid
        true
    }
}

impl Drop for DxgiCapturer {
    fn drop(&mut self) {
        // Release any held frame
        unsafe {
            self.duplication.ReleaseFrame().ok();
        }
    }
}