repose-render-wgpu 0.30.6

//! Unified offscreen helper (native/android/wasm).

use anyhow::Result;
use wgpu::{PollType, TextureFormat};

use crate::WgpuSceneRenderer;
use repose_core::Scene;
#[cfg(not(target_arch = "wasm32"))]
use std::sync::{Mutex, OnceLock};

#[cfg(not(target_arch = "wasm32"))]
static SHARED_DEVICE: OnceLock<Mutex<Option<(wgpu::Device, wgpu::Queue)>>> = OnceLock::new();

#[cfg(not(target_arch = "wasm32"))]
fn shared_slot() -> &'static Mutex<Option<(wgpu::Device, wgpu::Queue)>> {
    SHARED_DEVICE.get_or_init(|| Mutex::new(None))
}

#[cfg(target_arch = "wasm32")]
thread_local! {
    static SHARED_DEVICE_WASM: std::cell::RefCell<Option<(wgpu::Device, wgpu::Queue)>> = std::cell::RefCell::new(None);
}

/// Publish live Device/Queue for shared offscreen.
///
/// WASM: `thread_local`, so a worker thread never sees the main-thread
/// device. Do NOT send the `Device`/`Queue` to a worker and render there..
pub fn set_shared_device(device: wgpu::Device, queue: wgpu::Queue) {
    #[cfg(not(target_arch = "wasm32"))]
    {
        if let Ok(mut g) = shared_slot().lock() {
            *g = Some((device, queue));
        }
    }
    #[cfg(target_arch = "wasm32")]
    {
        SHARED_DEVICE_WASM.with(|c| *c.borrow_mut() = Some((device, queue)));
    }
}
pub fn shared_device() -> Option<(wgpu::Device, wgpu::Queue)> {
    #[cfg(not(target_arch = "wasm32"))]
    {
        shared_slot().lock().ok().and_then(|g| g.clone())
    }
    #[cfg(target_arch = "wasm32")]
    {
        SHARED_DEVICE_WASM.with(|c| c.borrow().clone())
    }
}
pub fn clear_shared_device() {
    #[cfg(not(target_arch = "wasm32"))]
    {
        if let Ok(mut g) = shared_slot().lock() {
            *g = None;
        }
    }
    #[cfg(target_arch = "wasm32")]
    {
        SHARED_DEVICE_WASM.with(|c| *c.borrow_mut() = None);
    }
}

/// Headless scene renderer: renders a [`Scene`] into CPU-readable RGBA8.
///
/// # Contracts
///
/// - Output is `width * height * 4` bytes, row-major, **premultiplied-alpha
///   sRGB** (`Rgba8UnormSrgb`): opaque content reads back exact sRGB bytes;
///   translucent pixels carry `rgb * a` (divide by `a` for straight alpha).
/// - `clear = None` fills with `scene.clear_color`; `Some(c)` overrides it
///   with linear-space `[r, g, b, a]` doubles.
/// - All sizes clamp to a minimum of 1 (no zero-sized textures/buffers).
/// - `render_rgba` blocks until the GPU work completes and the pixels are
///   mapped; it is safe to call every frame (no per-frame allocations
///   beyond the returned `Vec`).
/// - Without a usable GPU adapter, [`new`](Self::new)/[`new_blocking`](Self::new_blocking)
///   fail with an error (never a panic, never silent fallback).
pub struct OffscreenRenderer {
    renderer: WgpuSceneRenderer,
    texture: wgpu::Texture,
    view: wgpu::TextureView,
    readback: wgpu::Buffer,
    width: u32,
    height: u32,
    padded_bytes_per_row: u32,
}

impl OffscreenRenderer {
    /// Create a renderer for `width`x`height` frames with `msaa` samples
    /// (clamped to what the adapter supports via `pick_surface_msaa`).
    /// Dimensions clamp to ≥ 1. Errors when no GPU adapter is available.
    pub async fn new(width: u32, height: u32, msaa: u32) -> Result<Self> {
        let width = width.max(1);
        let height = height.max(1);
        let instance = if cfg!(target_arch = "wasm32") {
            let mut desc = wgpu::InstanceDescriptor::new_without_display_handle();
            desc.backends = wgpu::Backends::BROWSER_WEBGPU | wgpu::Backends::GL;
            wgpu::util::new_instance_with_webgpu_detection(desc).await
        } else {
            wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle())
        };
        let adapter = instance
            .request_adapter(&wgpu::RequestAdapterOptions {
                power_preference: wgpu::PowerPreference::HighPerformance,
                compatible_surface: None,
                ..Default::default()
            })
            .await?;
        let format = TextureFormat::Rgba8UnormSrgb;
        let msaa = crate::pick_surface_msaa(&adapter, format, msaa);
        let (device, queue) = adapter
            .request_device(&wgpu::DeviceDescriptor {
                label: Some("repose-offscreen"),
                required_features: wgpu::Features::empty(),
                required_limits: wgpu::Limits::default(),
                memory_hints: wgpu::MemoryHints::Performance,
                experimental_features: wgpu::ExperimentalFeatures::default(),
                trace: wgpu::Trace::Off,
            })
            .await?;
        let renderer = WgpuSceneRenderer::from_device(device, queue, format, msaa);
        Self::from_renderer(renderer, width, height)
    }

    /// Blocking [`new`](Self::new). Must not run on the wasm main thread
    /// without block support (debug-asserted); use `render_rgba_async` there.
    pub fn new_blocking(width: u32, height: u32, msaa: u32) -> Result<Self> {
        pollster::block_on(Self::new(width, height, msaa))
    }

    /// Shared-device: reuse Device/Queue, no Adapter. Dimensions clamp to
    /// ≥ 1; `msaa` clamps to ≥ 1 (no adapter-based picking without one —
    /// prefer [`from_device_with_adapter`](Self::from_device_with_adapter)
    /// when an adapter is handy).
    pub fn from_device(
        device: wgpu::Device,
        queue: wgpu::Queue,
        width: u32,
        height: u32,
        msaa: u32,
    ) -> Result<Self> {
        let width = width.max(1);
        let height = height.max(1);
        let format = TextureFormat::Rgba8UnormSrgb;
        let renderer = WgpuSceneRenderer::from_device(device, queue, format, msaa.max(1));
        Self::from_renderer(renderer, width, height)
    }

    /// Shared-device with adapter-based MSAA picking. Dimensions clamp to
    /// ≥ 1.
    pub fn from_device_with_adapter(
        device: wgpu::Device,
        queue: wgpu::Queue,
        adapter: &wgpu::Adapter,
        width: u32,
        height: u32,
        msaa: u32,
    ) -> Result<Self> {
        let width = width.max(1);
        let height = height.max(1);
        let format = TextureFormat::Rgba8UnormSrgb;
        let msaa = crate::pick_surface_msaa(adapter, format, msaa);
        let renderer = WgpuSceneRenderer::from_device(device, queue, format, msaa);
        Self::from_renderer(renderer, width, height)
    }

    /// Wrap an existing scene renderer with a fresh target + readback
    /// buffer. Dimensions clamp to ≥ 1.
    pub fn from_renderer(mut renderer: WgpuSceneRenderer, width: u32, height: u32) -> Result<Self> {
        let width = width.max(1);
        let height = height.max(1);
        renderer.resize(width, height);
        let texture = renderer.device.create_texture(&wgpu::TextureDescriptor {
            label: Some("repose-offscreen-tex"),
            size: wgpu::Extent3d {
                width,
                height,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: TextureFormat::Rgba8UnormSrgb,
            usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
            view_formats: &[],
        });
        let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
        let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
        let unpadded = width * 4;
        let padded = unpadded.div_ceil(align) * align;
        let buf_size = (padded * height) as u64;
        let readback = renderer.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("repose-offscreen-readback"),
            size: buf_size,
            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
            mapped_at_creation: false,
        });
        Ok(Self {
            renderer,
            texture,
            view,
            readback,
            width,
            height,
            padded_bytes_per_row: padded,
        })
    }

    fn encode_rgba(&mut self, scene: &Scene, clear: Option<[f64; 4]>) -> wgpu::CommandBuffer {
        let mut encoder =
            self.renderer
                .device
                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                    label: Some("repose-offscreen-encoder"),
                });
        self.renderer
            .render_scene_to_encoder(scene, &mut encoder, &self.view, clear);
        encoder.copy_texture_to_buffer(
            wgpu::TexelCopyTextureInfo {
                texture: &self.texture,
                mip_level: 0,
                origin: wgpu::Origin3d::ZERO,
                aspect: wgpu::TextureAspect::All,
            },
            wgpu::TexelCopyBufferInfo {
                buffer: &self.readback,
                layout: wgpu::TexelCopyBufferLayout {
                    offset: 0,
                    bytes_per_row: Some(self.padded_bytes_per_row),
                    rows_per_image: Some(self.height),
                },
            },
            wgpu::Extent3d {
                width: self.width,
                height: self.height,
                depth_or_array_layers: 1,
            },
        );
        encoder.finish()
    }

    /// Render `scene` and read back `width * height * 4` premultiplied sRGB
    /// bytes (see the [struct contracts](Self)). `clear = None` uses
    /// `scene.clear_color`; an empty scene therefore reads back the clear
    /// color exactly. Blocks until mapping completes.
    pub fn render_rgba(&mut self, scene: &Scene, clear: Option<[f64; 4]>) -> Result<Vec<u8>> {
        // Blocking map on the wasm main thread (Window) can never complete..
        #[cfg(all(target_family = "wasm", target_os = "unknown"))]
        debug_assert!(
            web_workers::web::has_block_support(),
            "render_rgba (blocking) called on wasm main thread; use render_rgba_async"
        );
        let cmd = self.encode_rgba(scene, clear);
        self.renderer.queue.submit(Some(cmd));
        let slice = self.readback.slice(..);
        let (tx, rx) = web_workers::sync::mpsc::channel();
        slice.map_async(wgpu::MapMode::Read, move |r| {
            let _ = tx.send_sync(r);
        });
        self.renderer.device.poll(PollType::wait_indefinitely())?;
        rx.recv_sync()??;
        let mapped = slice.get_mapped_range()?;
        let out = strip_padding(&mapped, self.width, self.height, self.padded_bytes_per_row);
        drop(mapped);
        self.readback.unmap();
        Ok(out)
    }

    #[cfg(target_arch = "wasm32")]
    pub async fn render_rgba_async(
        &mut self,
        scene: &Scene,
        clear: Option<[f64; 4]>,
    ) -> Result<Vec<u8>> {
        let cmd = self.encode_rgba(scene, clear);
        self.renderer.queue.submit(Some(cmd));
        let slice = self.readback.slice(..);
        let (tx, rx) = web_workers::sync::mpsc::channel();
        slice.map_async(wgpu::MapMode::Read, {
            let tx = tx.clone();
            move |r| {
                let _ = tx.send_sync(r);
            }
        });
        let start = web_workers::sync::Instant::now();
        loop {
            let _ = self.renderer.device.poll(PollType::Poll);
            if let Ok(r) = rx.try_recv() {
                r?;
                break;
            }
            if start.elapsed().as_secs() > 30 {
                anyhow::bail!("offscreen readback timed out after 30s (map callback never fired)");
            }
            web_workers::web::yield_now_async(web_workers::web::YieldTime::UserVisible).await;
        }
        let mapped = slice.get_mapped_range()?;
        let out = strip_padding(&*mapped, self.width, self.height, self.padded_bytes_per_row);
        drop(mapped);
        self.readback.unmap();
        Ok(out)
    }

    pub fn renderer_mut(&mut self) -> &mut WgpuSceneRenderer {
        &mut self.renderer
    }
    pub fn renderer(&self) -> &WgpuSceneRenderer {
        &self.renderer
    }

    /// Current target width in pixels.
    pub fn width(&self) -> u32 {
        self.width
    }

    /// Current target height in pixels.
    pub fn height(&self) -> u32 {
        self.height
    }

    /// Resize the target + readback buffer. A no-op when the size already
    /// matches (no reallocation). Dimensions clamp to ≥ 1. After this,
    /// `render_rgba` returns `width * height * 4` bytes at the new size.
    pub fn ensure_size(&mut self, width: u32, height: u32) -> Result<()> {
        let width = width.max(1);
        let height = height.max(1);
        if self.width == width && self.height == height {
            return Ok(());
        }
        self.renderer.resize(width, height);
        let texture = self
            .renderer
            .device
            .create_texture(&wgpu::TextureDescriptor {
                label: Some("repose-offscreen-tex"),
                size: wgpu::Extent3d {
                    width,
                    height,
                    depth_or_array_layers: 1,
                },
                mip_level_count: 1,
                sample_count: 1,
                dimension: wgpu::TextureDimension::D2,
                format: TextureFormat::Rgba8UnormSrgb,
                usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
                view_formats: &[],
            });
        let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
        let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
        let padded = (width * 4).div_ceil(align) * align;
        let readback = self.renderer.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("repose-offscreen-readback"),
            size: (padded * height) as u64,
            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
            mapped_at_creation: false,
        });
        self.texture = texture;
        self.view = view;
        self.readback = readback;
        self.width = width;
        self.height = height;
        self.padded_bytes_per_row = padded;
        Ok(())
    }

    pub async fn render_rgba_unified(
        &mut self,
        scene: &Scene,
        clear: Option<[f64; 4]>,
    ) -> Result<Vec<u8>> {
        #[cfg(all(target_family = "wasm", target_os = "unknown"))]
        {
            if web_workers::web::has_block_support() {
                return self.render_rgba(scene, clear);
            } else {
                return self.render_rgba_async(scene, clear).await;
            }
        }
        #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
        {
            self.render_rgba(scene, clear)
        }
    }
}

fn strip_padding(mapped: &[u8], width: u32, height: u32, padded: u32) -> Vec<u8> {
    let unpadded = width * 4;
    let mut out = Vec::with_capacity((width * height * 4) as usize);
    for y in 0..height {
        let start = (y * padded) as usize;
        let end = start + unpadded as usize;
        out.extend_from_slice(&mapped[start..end]);
    }
    out
}

pub fn map_buffer_blocking(slice: &wgpu::BufferSlice<'_>, device: &wgpu::Device) -> Result<()> {
    let (tx, rx) = web_workers::sync::mpsc::channel();
    slice.map_async(wgpu::MapMode::Read, move |r| {
        let _ = tx.send_sync(r);
    });
    device.poll(PollType::wait_indefinitely())?;
    rx.recv_sync()??;
    Ok(())
}

#[cfg(target_arch = "wasm32")]
pub async fn map_buffer_async(slice: &wgpu::BufferSlice<'_>, device: &wgpu::Device) -> Result<()> {
    let (tx, rx) = web_workers::sync::mpsc::channel();
    slice.map_async(wgpu::MapMode::Read, {
        let tx = tx.clone();
        move |r| {
            let _ = tx.send_sync(r);
        }
    });
    let start = web_workers::sync::Instant::now();
    loop {
        let _ = device.poll(PollType::Poll);
        if let Ok(r) = rx.try_recv() {
            r?;
            break;
        }
        if start.elapsed().as_secs() > 30 {
            anyhow::bail!("map_buffer timed out after 30s (map callback never fired)");
        }
        web_workers::web::yield_now_async(web_workers::web::YieldTime::UserVisible).await;
    }
    Ok(())
}

pub async fn map_buffer_unified(
    slice: &wgpu::BufferSlice<'_>,
    device: &wgpu::Device,
) -> Result<()> {
    #[cfg(all(target_family = "wasm", target_os = "unknown"))]
    {
        if web_workers::web::has_block_support() {
            return map_buffer_blocking(slice, device);
        } else {
            return map_buffer_async(slice, device).await;
        }
    }
    #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
    {
        map_buffer_blocking(slice, device)
    }
}