vtsampler 0.1.1

//! D3D11 shared textures ↔ wgpu (DX12 interop).
//!
//! Bridge textures are created with `D3D11_RESOURCE_MISC_SHARED` and opened in wgpu via
//! `OpenSharedHandle`. Your [`wgpu::Device`] must use the **DX12** backend.
//!
//! # Typical flow
//!
//! 1. Create a [`VtD3d11Bridge`] (or obtain one from [`VtD3d11Pool`]).
//! 2. [`VtD3d11Bridge::copy_from`] from capture / decoder texture.
//! 3. [`crate::VTImage::from_d3d11_bridge`] + [`crate::VTSampler::process`].
//! 4. [`VtD3d11Bridge::copy_to`] into an encoder-facing D3D11 texture if needed.

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use windows::core::Interface;
use windows::Win32::Foundation::HANDLE;
use windows::Win32::Graphics::{
    Direct3D11::{
        D3D11_BIND_SHADER_RESOURCE, D3D11_RESOURCE_MISC_SHARED, D3D11_TEXTURE2D_DESC,
        D3D11_USAGE_DEFAULT, ID3D11Device, ID3D11DeviceContext, ID3D11Texture2D,
    },
    Direct3D12::ID3D12Resource,
    Dxgi::Common::{
        DXGI_FORMAT, DXGI_FORMAT_B8G8R8A8_UNORM, DXGI_FORMAT_NV12, DXGI_FORMAT_R8G8B8A8_UNORM,
    },
    Dxgi::IDXGIResource,
};

use wgpu::{
    Device, Extent3d, Texture, TextureDescriptor, TextureDimension, TextureFormat, TextureUsages,
    hal::api::Dx12,
};

use crate::{VTFormat, bridge::BridgeError};

/// Borrowed D3D11 device and immediate context for `CopySubresourceRegion`.
pub struct VtD3d11Device<'a> {
    /// D3D11 device that owns shared textures.
    pub device: &'a ID3D11Device,
    /// Immediate context used for GPU copies.
    pub context: &'a ID3D11DeviceContext,
}

impl<'a> VtD3d11Device<'a> {
    /// Wraps existing D3D11 device pointers (lifetime tied to the caller).
    pub fn new(device: &'a ID3D11Device, context: &'a ID3D11DeviceContext) -> Self {
        Self { device, context }
    }
}

/// Shared D3D11 texture and its wgpu import of the same resource.
pub struct VtD3d11Bridge {
    /// D3D11 texture with `D3D11_RESOURCE_MISC_SHARED`.
    pub d3d11: ID3D11Texture2D,
    /// wgpu texture created via DX12 shared handle (same memory as `d3d11`).
    pub wgpu: Texture,
}

impl VtD3d11Bridge {
    /// Creates a new shared D3D11 texture and imports it into wgpu.
    ///
    /// # Errors
    ///
    /// Returns [`crate::BridgeError::NotFoundDxBackend`] if wgpu is not on DX12.
    pub fn new(
        d3d: &VtD3d11Device<'_>,
        wgpu_device: &Device,
        width: u32,
        height: u32,
        dxgi_format: DXGI_FORMAT,
        wgpu_format: TextureFormat,
        usage: TextureUsages,
    ) -> Result<Self, BridgeError> {
        let mut desc = D3D11_TEXTURE2D_DESC::default();
        desc.Width = width;
        desc.Height = height;
        desc.MipLevels = 1;
        desc.ArraySize = 1;
        desc.SampleDesc.Count = 1;
        desc.Format = dxgi_format;
        desc.Usage = D3D11_USAGE_DEFAULT;
        desc.BindFlags = D3D11_BIND_SHADER_RESOURCE.0 as u32;
        desc.MiscFlags = D3D11_RESOURCE_MISC_SHARED.0 as u32;

        let mut texture = None;
        unsafe {
            d3d.device
                .CreateTexture2D(&desc, None, Some(&mut texture))?;
        }
        let d3d11 = texture.unwrap();

        let wgpu = import_d3d11_shared(wgpu_device, &d3d11, wgpu_format, width, height, usage)?;

        Ok(Self { d3d11, wgpu })
    }

    /// GPU copy from `src` into this bridge's D3D11 texture (visible to wgpu after import).
    pub fn copy_from(
        &self,
        d3d: &VtD3d11Device<'_>,
        src: &ID3D11Texture2D,
        array_index: u32,
    ) -> Result<(), BridgeError> {
        unsafe {
            d3d.context.CopySubresourceRegion(
                &self.d3d11, 0, 0, 0, 0, src, array_index, None,
            );
        }
        Ok(())
    }

    /// GPU copy from this bridge's D3D11 texture into `dst`.
    pub fn copy_to(
        &self,
        d3d: &VtD3d11Device<'_>,
        dst: &ID3D11Texture2D,
        array_index: u32,
    ) -> Result<(), BridgeError> {
        unsafe {
            d3d.context.CopySubresourceRegion(
                dst, array_index, 0, 0, 0, &self.d3d11, 0, None,
            );
        }
        Ok(())
    }
}

/// Cache of [`VtD3d11Bridge`] instances keyed by size, DXGI format, and wgpu usage.
///
/// Used internally by [`crate::VTSampler`] for [`crate::VTImage::from_d3d11`]; you may also
/// use it directly to avoid recreating shared textures every frame.
pub struct VtD3d11Pool {
    entries: Mutex<HashMap<PoolKey, Arc<VtD3d11Bridge>>>,
}

#[derive(Hash, Eq, PartialEq)]
struct PoolKey {
    width: u32,
    height: u32,
    dxgi: i32,
    usage: u32,
}

impl VtD3d11Pool {
    /// Creates an empty pool.
    pub fn new() -> Self {
        Self {
            entries: Mutex::new(HashMap::new()),
        }
    }

    /// Returns an existing bridge or creates one for the given parameters.
    pub fn acquire(
        &self,
        d3d: &VtD3d11Device<'_>,
        wgpu: &Device,
        width: u32,
        height: u32,
        format: VTFormat,
        usage: TextureUsages,
    ) -> Result<Arc<VtD3d11Bridge>, BridgeError> {
        let (dxgi, wgpu_format) = vt_format_to_dxgi_wgpu(format)?;
        let key = PoolKey {
            width,
            height,
            dxgi: dxgi.0,
            usage: usage.bits(),
        };
        let mut map = self.entries.lock().expect("d3d11 pool lock");
        if let Some(bridge) = map.get(&key) {
            return Ok(bridge.clone());
        }
        let bridge = Arc::new(VtD3d11Bridge::new(
            d3d, wgpu, width, height, dxgi, wgpu_format, usage,
        )?);
        map.insert(key, bridge.clone());
        Ok(bridge)
    }
}

/// Maps a [`VTFormat`] to DXGI and wgpu texture formats for bridge creation.
///
/// [`VTFormat::YUV420P`] is not supported on D3D11 bridges.
pub fn vt_format_to_dxgi_wgpu(format: VTFormat) -> Result<(DXGI_FORMAT, TextureFormat), BridgeError> {
    match format {
        VTFormat::RGBA => Ok((DXGI_FORMAT_R8G8B8A8_UNORM, TextureFormat::Rgba8Unorm)),
        VTFormat::BGRA => Ok((DXGI_FORMAT_B8G8R8A8_UNORM, TextureFormat::Bgra8Unorm)),
        VTFormat::NV12 => Ok((DXGI_FORMAT_NV12, TextureFormat::NV12)),
        VTFormat::YUV420P => Err(BridgeError::UnsupportedFormat),
    }
}

/// Reads `DXGI_FORMAT` from an existing D3D11 texture descriptor.
pub fn dxgi_from_texture(texture: &ID3D11Texture2D) -> DXGI_FORMAT {
    let mut desc = D3D11_TEXTURE2D_DESC::default();
    unsafe {
        texture.GetDesc(&mut desc);
    }
    desc.Format
}

fn shared_handle(texture: &ID3D11Texture2D) -> Result<HANDLE, BridgeError> {
    let handle = unsafe { texture.cast::<IDXGIResource>()?.GetSharedHandle()? };
    if handle.is_invalid() {
        return Err(BridgeError::InvalidSharedHandle);
    }
    Ok(handle)
}

fn import_d3d11_shared(
    wgpu_device: &Device,
    d3d11: &ID3D11Texture2D,
    wgpu_format: TextureFormat,
    width: u32,
    height: u32,
    usage: TextureUsages,
) -> Result<Texture, BridgeError> {
    let tex_desc = TextureDescriptor {
        label: Some("vtsampler_d3d11_bridge"),
        size: Extent3d {
            width,
            height,
            depth_or_array_layers: 1,
        },
        mip_level_count: 1,
        sample_count: 1,
        dimension: TextureDimension::D2,
        format: wgpu_format,
        usage,
        view_formats: &[],
    };

    unsafe {
        Ok(wgpu_device.create_texture_from_hal::<Dx12>(
            <Dx12 as wgpu::hal::Api>::Device::texture_from_raw(
                wgpu_device
                    .as_hal::<Dx12, _, _>(|hdevice| {
                        let mut resource = None::<ID3D12Resource>;
                        hdevice
                            .ok_or(BridgeError::NotFoundDxBackend)?
                            .raw_device()
                            .OpenSharedHandle(shared_handle(d3d11)?, &mut resource)
                            .map(|_| resource.unwrap())
                            .map_err(BridgeError::Windows)
                    })
                    .ok_or(BridgeError::NotFoundDxBackend)??,
                wgpu_format,
                TextureDimension::D2,
                tex_desc.size,
                tex_desc.mip_level_count,
                tex_desc.sample_count,
            ),
            &tex_desc,
        ))
    }
}

/// Imports an **existing** shared D3D11 texture into wgpu (no `CopySubresourceRegion`).
///
/// The texture must already have been created with `D3D11_RESOURCE_MISC_SHARED`.
pub fn import_d3d11_texture(
    wgpu_device: &Device,
    d3d11: &ID3D11Texture2D,
    format: VTFormat,
    usage: TextureUsages,
) -> Result<Texture, BridgeError> {
    let mut desc = D3D11_TEXTURE2D_DESC::default();
    unsafe {
        d3d11.GetDesc(&mut desc);
    }
    let (_, wgpu_format) = vt_format_to_dxgi_wgpu(format)?;
    import_d3d11_shared(
        wgpu_device,
        d3d11,
        wgpu_format,
        desc.Width,
        desc.Height,
        usage,
    )
}