vtsampler 0.1.1

//! `CVPixelBuffer` → Metal → wgpu upload.
//!
//! Requires a wgpu device on the **Metal** backend. Used by [`crate::VTImage::from_cv_pixel_buffer`]
//! and available directly via [`VtMetalCache`].

use std::ptr::{NonNull, null_mut};
use std::sync::{Arc, Mutex};

use metal::{MTLTextureType, Texture as MtlTexture};
use objc2::rc::Retained;
use objc2::runtime::ProtocolObject;
use objc2_core_foundation::kCFAllocatorDefault;
use objc2_core_video::{
    CVMetalTexture, CVMetalTextureCache, CVMetalTextureCacheCreate,
    CVMetalTextureCacheCreateTextureFromImage, CVMetalTextureCacheFlush, CVMetalTextureGetTexture,
    CVPixelBuffer, CVPixelBufferGetHeight, CVPixelBufferGetPixelFormatType, CVPixelBufferGetWidth,
    kCVPixelFormatType_32BGRA, kCVPixelFormatType_32RGBA,
    kCVPixelFormatType_420YpCbCr8BiPlanarFullRange,
    kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange, kCVReturnSuccess,
};
use objc2_metal::{MTLDevice as Objc2MTLDevice, MTLPixelFormat as Objc2MTLPixelFormat};

use wgpu::{
    CommandEncoder, Device, Extent3d, ImageCopyTexture, Origin3d, Texture, TextureDescriptor,
    TextureDimension, TextureFormat, TextureUsages, hal::{Api, CopyExtent, api::Metal},
};

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

/// Opaque Core Video pixel buffer handle (`CVPixelBuffer *`).
pub type CVPixelBufferRef = *mut CVPixelBuffer;

/// Maps `CVPixelBuffer` planes to Metal textures and copies into wgpu.
pub struct VtMetalCache {
    cache: Retained<CVMetalTextureCache>,
    wgpu_device: Arc<Device>,
}

impl VtMetalCache {
    /// Creates a texture cache tied to the same Metal device as `wgpu_device`.
    ///
    /// # Errors
    ///
    /// * [`crate::BridgeError::NotFoundMetalBackend`] — wgpu is not using Metal.
    /// * [`crate::BridgeError::CoreVideo`] — `CVMetalTextureCacheCreate` failed.
    pub fn new(wgpu_device: Arc<Device>) -> Result<Self, BridgeError> {
        let mut raw_mtl = None;
        wgpu_device.as_hal::<Metal, _, _>(|device| {
            if let Some(device) = device {
                raw_mtl = Some(device.raw_device().lock().clone());
            }
        });
        let raw_mtl = raw_mtl.ok_or(BridgeError::NotFoundMetalBackend)?;

        let device: Retained<ProtocolObject<dyn Objc2MTLDevice>> =
            unsafe { Retained::from_raw(raw_mtl.into_ptr().cast()).unwrap() };

        let mut cache = null_mut();
        let code = unsafe {
            CVMetalTextureCacheCreate(
                kCFAllocatorDefault,
                None,
                device.as_ref(),
                None,
                NonNull::new(&mut cache).unwrap(),
            )
        };
        if code != kCVReturnSuccess || cache.is_null() {
            return Err(BridgeError::CoreVideo(code));
        }

        Ok(Self {
            cache: unsafe { Retained::from_raw(cache).unwrap() },
            wgpu_device,
        })
    }

    /// Flushes the Core Video metal texture cache (call after batched uploads if needed).
    pub fn flush(&self) {
        unsafe {
            CVMetalTextureCacheFlush(&self.cache, 0);
        }
    }

    /// Copies pixel buffer contents into `dst_planes` within `encoder`.
    ///
    /// Plane count and formats must match `format` (1 plane for RGBA/BGRA, 2 for NV12).
    pub fn upload_to_planes(
        &self,
        encoder: &mut CommandEncoder,
        buffer: CVPixelBufferRef,
        format: VTFormat,
        width: u32,
        height: u32,
        dst_planes: &[Texture],
    ) -> Result<(), BridgeError> {
        let pixel_format = unsafe { CVPixelBufferGetPixelFormatType(&*buffer) };
        match format {
            VTFormat::BGRA | VTFormat::RGBA => {
                let mtl_format = match format {
                    VTFormat::BGRA => Objc2MTLPixelFormat::BGRA8Unorm,
                    VTFormat::RGBA => Objc2MTLPixelFormat::RGBA8Unorm,
                    _ => unreachable!(),
                };
                let mtl_tex = self.create_metal_texture(buffer, mtl_format, width, height, 0)?;
                self.copy_metal_to_wgpu(encoder, mtl_tex, format, &dst_planes[0], width, height)?;
            }
            VTFormat::NV12 => {
                if pixel_format != kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange
                    && pixel_format != kCVPixelFormatType_420YpCbCr8BiPlanarFullRange
                {
                    return Err(BridgeError::UnsupportedFormat);
                }
                let y_tex = self.create_metal_texture(
                    buffer,
                    Objc2MTLPixelFormat::R8Unorm,
                    width,
                    height,
                    0,
                )?;
                let (uvw, uvh) = plane_size(format, width, height, 1);
                let uv_tex = self.create_metal_texture(
                    buffer,
                    Objc2MTLPixelFormat::RG8Unorm,
                    uvw,
                    uvh,
                    1,
                )?;
                self.copy_metal_to_wgpu(encoder, y_tex, VTFormat::NV12, &dst_planes[0], width, height)?;
                self.copy_metal_to_wgpu(encoder, uv_tex, VTFormat::NV12, &dst_planes[1], uvw, uvh)?;
            }
            VTFormat::YUV420P => return Err(BridgeError::UnsupportedFormat),
        }
        self.flush();
        Ok(())
    }

    fn create_metal_texture(
        &self,
        buffer: CVPixelBufferRef,
        mtl_format: Objc2MTLPixelFormat,
        width: u32,
        height: u32,
        plane_index: usize,
    ) -> Result<MtlTexture, BridgeError> {
        let mut cv_texture = null_mut();
        let code = unsafe {
            CVMetalTextureCacheCreateTextureFromImage(
                kCFAllocatorDefault,
                &self.cache,
                &*buffer,
                None,
                mtl_format,
                width as usize,
                height as usize,
                plane_index,
                NonNull::new(&mut cv_texture).unwrap(),
            )
        };
        if code != kCVReturnSuccess || cv_texture.is_null() {
            return Err(BridgeError::CoreVideo(code));
        }
        let cv_texture = unsafe { Retained::<CVMetalTexture>::from_raw(cv_texture).unwrap() };
        if let Some(texture) = unsafe { CVMetalTextureGetTexture(&cv_texture) } {
            Ok(unsafe {
                MtlTexture::from_ptr(Retained::into_raw(texture).cast()).to_owned()
            })
        } else {
            Err(BridgeError::CoreVideo(-1))
        }
    }

    fn copy_metal_to_wgpu(
        &self,
        encoder: &mut CommandEncoder,
        mtl_src: MtlTexture,
        format: VTFormat,
        dst: &Texture,
        width: u32,
        height: u32,
    ) -> Result<(), BridgeError> {
        let wgpu_format = format.plane_formats()[0];
        let src = unsafe {
            self.wgpu_device.create_texture_from_hal::<Metal>(
                <Metal as Api>::Device::texture_from_raw(
                    mtl_src,
                    wgpu_format,
                    MTLTextureType::D2,
                    1,
                    1,
                    CopyExtent {
                        width,
                        height,
                        depth: 1,
                    },
                ),
                &TextureDescriptor {
                    label: Some("vtsampler_metal_import"),
                    size: Extent3d {
                        width,
                        height,
                        depth_or_array_layers: 1,
                    },
                    mip_level_count: 1,
                    sample_count: 1,
                    dimension: TextureDimension::D2,
                    format: wgpu_format,
                    usage: TextureUsages::COPY_SRC,
                    view_formats: &[],
                },
            )
        };
        encoder.copy_texture_to_texture(
            ImageCopyTexture {
                texture: &src,
                mip_level: 0,
                origin: Origin3d::ZERO,
                aspect: wgpu::TextureAspect::All,
            },
            ImageCopyTexture {
                texture: dst,
                mip_level: 0,
                origin: Origin3d::ZERO,
                aspect: wgpu::TextureAspect::All,
            },
            Extent3d {
                width,
                height,
                depth_or_array_layers: 1,
            },
        );
        Ok(())
    }
}

/// Maps a `kCVPixelFormatType_*` constant to [`VTFormat`].
pub fn cv_pixel_format_to_vt(pixel_format: u32) -> Result<VTFormat, BridgeError> {
    match pixel_format {
        kCVPixelFormatType_32RGBA => Ok(VTFormat::RGBA),
        kCVPixelFormatType_32BGRA => Ok(VTFormat::BGRA),
        kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange
        | kCVPixelFormatType_420YpCbCr8BiPlanarFullRange => Ok(VTFormat::NV12),
        _ => Err(BridgeError::UnsupportedFormat),
    }
}

/// Returns `(width, height)` of a pixel buffer.
pub fn cv_pixel_buffer_size(buffer: CVPixelBufferRef) -> (u32, u32) {
    (
        unsafe { CVPixelBufferGetWidth(&*buffer) } as u32,
        unsafe { CVPixelBufferGetHeight(&*buffer) } as u32,
    )
}