GORBIE 0.15.0

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;

/// Shift every clip rect and mesh vertex in `primitives` vertically by
/// `delta` (in logical points). Used during tiled rendering to pull
/// content from a later band of the card into the next tile's
/// texture, so the renderer can keep using a small `screen_descriptor`.
/// Paint callbacks are left untouched — they aren't generally used by
/// the GORBIE widget set, and translating them would require deeper
/// integration with `egui_wgpu`.
fn shift_primitives_y(primitives: &mut [egui::ClippedPrimitive], delta: f32) {
    let shift = egui::vec2(0.0, delta);
    for prim in primitives {
        prim.clip_rect = prim.clip_rect.translate(shift);
        if let egui::epaint::Primitive::Mesh(mesh) = &mut prim.primitive {
            for v in &mut mesh.vertices {
                v.pos.y += delta;
            }
        }
    }
}

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,
        }))?;
        // Default `Limits::default()` caps `max_texture_dimension_2d`
        // at 8192, which clips tall notebook cards (e.g. compass with
        // many goals, messages with long bodies). Request the
        // adapter's full hardware ceiling — usually 16384 — so the
        // render succeeds for cards that just need a slightly larger
        // single-tile budget. Cards above the hardware limit still
        // need a tile-and-stitch backend.
        let device_desc = wgpu::DeviceDescriptor {
            label: Some("gorbie_headless_device"),
            required_features: wgpu::Features::empty(),
            required_limits: adapter.limits(),
            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 tiles =
                self.render_output(output, egui::vec2(self.card_width, final_height))?;
            for (tile_idx, image) in tiles.iter().enumerate() {
                self.save_capture(index, tile_idx, tiles.len(), 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_ui(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<Vec<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 full_height = (size_points.y * pixels_per_point).round().max(1.0) as u32;

        for (id, delta) in &textures_delta.set {
            self.renderer
                .update_texture(&self.device, &self.queue, *id, delta);
        }

        // Cards taller than `tile_max_height` are split into vertical
        // tiles. The shapes are tessellated once; between tiles, every
        // clip rect and mesh vertex is shifted upward by the previous
        // tile's height so the renderer sees them landing inside a
        // tile-sized texture with its own `screen_descriptor`. This
        // produces multiple PNGs per card — also handy for human
        // readers (and LLM ingesters) who'd rather not load a
        // 200-megapixel single image.
        //
        // The cap is well below the GPU's `max_texture_dimension_2d`
        // (typically 16384) so each tile stays browseable on its own.
        const TARGET_TILE_HEIGHT: u32 = 4000;
        let gpu_max = self.device.limits().max_texture_dimension_2d;
        let tile_max_height = TARGET_TILE_HEIGHT.min(gpu_max).min(full_height).max(1);
        let mut clipped_primitives = self.ctx.tessellate(shapes, pixels_per_point);

        let mut tiles = Vec::new();
        let mut y_cursor: u32 = 0;
        let mut tile_offset_points: f32 = 0.0;
        while y_cursor < full_height {
            let this_tile_height = (full_height - y_cursor).min(tile_max_height);

            // Shift primitives up by the previous tile's height (not on
            // the first iteration). `tile_offset_points` accumulates
            // across tiles so the second tile sees vertices already at
            // (orig - tile_max_height_points), the third at (orig - 2 *
            // tile_max_height_points), etc.
            if y_cursor > 0 {
                let shift = -(tile_max_height as f32) / pixels_per_point;
                shift_primitives_y(&mut clipped_primitives, shift);
                tile_offset_points += -shift;
            }
            let _ = tile_offset_points; // currently informational only

            self.ensure_target(width, this_tile_height)?;
            let screen_descriptor = egui_wgpu::ScreenDescriptor {
                size_in_pixels: [width, this_tile_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.global_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)?;
            tiles.push(RenderedImage {
                width,
                height: this_tile_height,
                pixels,
            });

            y_cursor += this_tile_height;
        }

        for id in &textures_delta.free {
            self.renderer.free_texture(id);
        }
        Ok(tiles)
    }

    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,
        tile_idx: usize,
        tile_count: usize,
        image: &RenderedImage,
    ) -> HeadlessResult<()> {
        // Single-tile cards stay as `card_NNNN.png`; multi-tile cards
        // get `card_NNNN_pMM.png` (1-based page index) so the original
        // naming is preserved for the common case.
        let filename = if tile_count <= 1 {
            format!("card_{:04}.png", index + 1)
        } else {
            format!("card_{:04}_p{:02}.png", index + 1, tile_idx + 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,
        );
        // Strip row padding required by wgpu's copy alignment.
        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)
}