use crate::themes::{industrial_dark, industrial_fonts, industrial_light};
use crate::{HeadlessCaptureConfig, NotebookCore, NOTEBOOK_MIN_HEIGHT};
use dark_light::Mode;
use eframe::egui;
use egui_wgpu::wgpu;
use std::path::{Path, PathBuf};
use std::time::{Duration, Instant};
type HeadlessResult<T> = Result<T, Box<dyn std::error::Error + Send + Sync>>;
const HEADLESS_PNG_PPI: f32 = 254.0;
pub(super) fn run_headless(
mut core: NotebookCore,
config: HeadlessCaptureConfig,
) -> HeadlessResult<()> {
let mut runner = HeadlessWgpuRunner::new(config)?;
runner.capture_cards(&mut core)
}
struct HeadlessWgpuRunner {
output_dir: PathBuf,
card_width: f32,
ctx: egui::Context,
device: wgpu::Device,
queue: wgpu::Queue,
renderer: egui_wgpu::Renderer,
target_format: wgpu::TextureFormat,
pixels_per_point: f32,
target: Option<TargetBuffers>,
time_seconds: f64,
settle_timeout: Duration,
}
impl HeadlessWgpuRunner {
fn new(config: HeadlessCaptureConfig) -> HeadlessResult<Self> {
std::fs::create_dir_all(&config.output_dir)?;
let ctx = egui::Context::default();
ctx.set_fonts(industrial_fonts());
ctx.set_style_of(egui::Theme::Light, industrial_light());
ctx.set_style_of(egui::Theme::Dark, industrial_dark());
let theme = match dark_light::detect() {
Ok(Mode::Light) => egui::ThemePreference::Light,
Ok(Mode::Dark) => egui::ThemePreference::Dark,
Ok(Mode::Unspecified) | Err(_) => egui::ThemePreference::Dark,
};
ctx.set_theme(theme);
let instance = wgpu::Instance::default();
let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
power_preference: wgpu::PowerPreference::HighPerformance,
compatible_surface: None,
force_fallback_adapter: false,
}))?;
let device_desc = wgpu::DeviceDescriptor {
label: Some("gorbie_headless_device"),
required_features: wgpu::Features::empty(),
required_limits: wgpu::Limits::default(),
experimental_features: wgpu::ExperimentalFeatures::default(),
memory_hints: wgpu::MemoryHints::default(),
trace: wgpu::Trace::default(),
};
let (device, queue) = pollster::block_on(adapter.request_device(&device_desc))?;
let target_format = wgpu::TextureFormat::Rgba8UnormSrgb;
let renderer = egui_wgpu::Renderer::new(
&device,
target_format,
egui_wgpu::RendererOptions::default(),
);
Ok(Self {
output_dir: config.output_dir,
card_width: config.card_width,
ctx,
device,
queue,
renderer,
target_format,
pixels_per_point: config.pixels_per_point,
target: None,
time_seconds: 0.0,
settle_timeout: config.settle_timeout,
})
}
fn capture_cards(&mut self, core: &mut NotebookCore) -> HeadlessResult<()> {
let mut index = 0;
loop {
let (mut output, measured_height) =
self.run_frame_until_settled(core, index, NOTEBOOK_MIN_HEIGHT)?;
let Some(measured_height) = measured_height else {
break;
};
let mut textures_delta = egui::TexturesDelta::default();
textures_delta.append(std::mem::take(&mut output.textures_delta));
let desired_height = measured_height.max(1.0);
let (mut output, final_height) = if height_close(NOTEBOOK_MIN_HEIGHT, desired_height) {
(output, desired_height)
} else {
let (mut output, measured_height) =
self.run_frame_until_settled(core, index, desired_height)?;
let Some(measured_height) = measured_height else {
break;
};
textures_delta.append(std::mem::take(&mut output.textures_delta));
(output, measured_height.max(1.0))
};
output.textures_delta = textures_delta;
let image = self.render_output(output, egui::vec2(self.card_width, final_height))?;
self.save_capture(index, &image)?;
index += 1;
}
Ok(())
}
fn run_frame_until_settled(
&mut self,
core: &mut NotebookCore,
index: usize,
height: f32,
) -> HeadlessResult<(egui::FullOutput, Option<f32>)> {
let start = Instant::now();
let mut textures_delta = egui::TexturesDelta::default();
let (mut output, mut measured_height) = self.run_frame(core, index, height)?;
loop {
textures_delta.append(std::mem::take(&mut output.textures_delta));
if core.has_settled() || start.elapsed() >= self.settle_timeout {
output.textures_delta = textures_delta;
return Ok((output, measured_height));
}
let repaint_delay = min_repaint_delay(&output);
if repaint_delay > Duration::ZERO {
std::thread::sleep(repaint_delay.min(Duration::from_millis(16)));
}
let (next_output, next_height) = self.run_frame(core, index, height)?;
output = next_output;
measured_height = next_height;
}
}
fn run_frame(
&mut self,
core: &mut NotebookCore,
index: usize,
height: f32,
) -> HeadlessResult<(egui::FullOutput, Option<f32>)> {
let screen_rect =
egui::Rect::from_min_size(egui::Pos2::ZERO, egui::vec2(self.card_width, height));
let mut raw_input = egui::RawInput {
screen_rect: Some(screen_rect),
max_texture_side: Some(self.device.limits().max_texture_dimension_2d as usize),
time: Some(self.time_seconds),
..Default::default()
};
{
let viewport = raw_input
.viewports
.entry(raw_input.viewport_id)
.or_default();
viewport.native_pixels_per_point = Some(self.pixels_per_point);
viewport.inner_rect = Some(screen_rect);
viewport.outer_rect = Some(screen_rect);
}
self.time_seconds += f64::from(raw_input.predicted_dt);
let mut notebook = core.build_notebook();
let mut measured_height: Option<f32> = None;
let output = self.ctx.run(raw_input, |ctx| {
measured_height = core.draw_card(ctx, &mut notebook, index, self.card_width);
});
Ok((output, measured_height))
}
fn render_output(
&mut self,
output: egui::FullOutput,
size_points: egui::Vec2,
) -> HeadlessResult<RenderedImage> {
let egui::FullOutput {
textures_delta,
shapes,
pixels_per_point,
..
} = output;
let width = (size_points.x * pixels_per_point).round().max(1.0) as u32;
let height = (size_points.y * pixels_per_point).round().max(1.0) as u32;
self.ensure_target(width, height)?;
for (id, delta) in &textures_delta.set {
self.renderer
.update_texture(&self.device, &self.queue, *id, delta);
}
let clipped_primitives = self.ctx.tessellate(shapes, pixels_per_point);
let screen_descriptor = egui_wgpu::ScreenDescriptor {
size_in_pixels: [width, height],
pixels_per_point,
};
let mut encoder = self
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor {
label: Some("gorbie_headless_encoder"),
});
let mut callbacks = self.renderer.update_buffers(
&self.device,
&self.queue,
&mut encoder,
&clipped_primitives,
&screen_descriptor,
);
let clear = color32_to_wgpu(self.ctx.style().visuals.window_fill);
let target = self.target.as_ref().expect("target ensured");
{
let render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
label: Some("gorbie_headless_pass"),
color_attachments: &[Some(wgpu::RenderPassColorAttachment {
view: &target.view,
depth_slice: None,
resolve_target: None,
ops: wgpu::Operations {
load: wgpu::LoadOp::Clear(clear),
store: wgpu::StoreOp::Store,
},
})],
depth_stencil_attachment: None,
multiview_mask: None,
timestamp_writes: None,
occlusion_query_set: None,
});
let mut render_pass = render_pass.forget_lifetime();
self.renderer
.render(&mut render_pass, &clipped_primitives, &screen_descriptor);
}
target.copy_to_buffer(&mut encoder);
callbacks.push(encoder.finish());
self.queue.submit(callbacks);
let _ = self.device.poll(wgpu::PollType::wait_indefinitely());
let pixels = target.readback(&self.device)?;
for id in &textures_delta.free {
self.renderer.free_texture(id);
}
Ok(RenderedImage {
width,
height,
pixels,
})
}
fn ensure_target(&mut self, width: u32, height: u32) -> HeadlessResult<()> {
let needs_resize = self.target.as_ref().map_or(true, |target| {
target.dims.width != width || target.dims.height != height
});
if needs_resize {
self.target = Some(TargetBuffers::new(
&self.device,
width,
height,
self.target_format,
));
}
Ok(())
}
fn save_capture(&self, index: usize, image: &RenderedImage) -> HeadlessResult<()> {
let filename = format!("card_{:04}.png", index + 1);
let path = self.output_dir.join(filename);
write_png_rgba(&path, image.width, image.height, &image.pixels)?;
Ok(())
}
}
struct TargetBuffers {
texture: wgpu::Texture,
view: wgpu::TextureView,
buffer: wgpu::Buffer,
dims: BufferDimensions,
}
impl TargetBuffers {
fn new(device: &wgpu::Device, width: u32, height: u32, format: wgpu::TextureFormat) -> Self {
let dims = BufferDimensions::new(width, height);
let texture = device.create_texture(&wgpu::TextureDescriptor {
label: Some("gorbie_headless_target"),
size: wgpu::Extent3d {
width,
height,
depth_or_array_layers: 1,
},
mip_level_count: 1,
sample_count: 1,
dimension: wgpu::TextureDimension::D2,
format,
usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
view_formats: &[format],
});
let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
let buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("gorbie_headless_readback"),
size: dims.buffer_size(),
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
Self {
texture,
view,
buffer,
dims,
}
}
fn copy_to_buffer(&self, encoder: &mut wgpu::CommandEncoder) {
let size = wgpu::Extent3d {
width: self.dims.width,
height: self.dims.height,
depth_or_array_layers: 1,
};
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.buffer,
layout: wgpu::TexelCopyBufferLayout {
offset: 0,
bytes_per_row: Some(self.dims.padded_bytes_per_row),
rows_per_image: Some(self.dims.height),
},
},
size,
);
}
fn readback(&self, device: &wgpu::Device) -> HeadlessResult<Vec<u8>> {
let buffer_slice = self.buffer.slice(..);
let (sender, receiver) = std::sync::mpsc::channel();
buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
let _ = sender.send(result);
});
let _ = device.poll(wgpu::PollType::wait_indefinitely());
receiver
.recv()
.map_err(|_| std::io::Error::new(std::io::ErrorKind::Other, "map channel closed"))??;
let data = buffer_slice.get_mapped_range();
let mut pixels = Vec::with_capacity(
self.dims.unpadded_bytes_per_row as usize * self.dims.height as usize,
);
for chunk in data.chunks(self.dims.padded_bytes_per_row as usize) {
pixels.extend_from_slice(&chunk[..self.dims.unpadded_bytes_per_row as usize]);
}
drop(data);
self.buffer.unmap();
Ok(pixels)
}
}
#[derive(Clone, Copy)]
struct BufferDimensions {
width: u32,
height: u32,
unpadded_bytes_per_row: u32,
padded_bytes_per_row: u32,
}
impl BufferDimensions {
fn new(width: u32, height: u32) -> Self {
let bytes_per_pixel = 4;
let unpadded_bytes_per_row = width * bytes_per_pixel;
let alignment = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
let padded_bytes_per_row =
((unpadded_bytes_per_row + alignment - 1) / alignment) * alignment;
Self {
width,
height,
unpadded_bytes_per_row,
padded_bytes_per_row,
}
}
fn buffer_size(&self) -> u64 {
self.padded_bytes_per_row as u64 * self.height as u64
}
}
struct RenderedImage {
width: u32,
height: u32,
pixels: Vec<u8>,
}
fn write_png_rgba(path: &Path, width: u32, height: u32, data: &[u8]) -> std::io::Result<()> {
let file = std::fs::File::create(path)?;
let mut encoder = png::Encoder::new(file, width, height);
encoder.set_color(png::ColorType::Rgba);
encoder.set_depth(png::BitDepth::Eight);
let pixels_per_meter = (HEADLESS_PNG_PPI / 0.0254).round().max(1.0) as u32;
encoder.set_pixel_dims(Some(png::PixelDimensions {
xppu: pixels_per_meter,
yppu: pixels_per_meter,
unit: png::Unit::Meter,
}));
let mut writer = encoder
.write_header()
.map_err(|err| std::io::Error::new(std::io::ErrorKind::Other, err))?;
writer
.write_image_data(data)
.map_err(|err| std::io::Error::new(std::io::ErrorKind::Other, err))
}
fn color32_to_wgpu(color: egui::Color32) -> wgpu::Color {
wgpu::Color {
r: f64::from(color.r()) / 255.0,
g: f64::from(color.g()) / 255.0,
b: f64::from(color.b()) / 255.0,
a: f64::from(color.a()) / 255.0,
}
}
fn height_close(a: f32, b: f32) -> bool {
(a - b).abs() <= 0.5
}
fn min_repaint_delay(output: &egui::FullOutput) -> Duration {
output
.viewport_output
.values()
.map(|viewport| viewport.repaint_delay)
.min()
.unwrap_or(Duration::MAX)
}