klyff 0.1.3

use bytemuck::Zeroable as _;

use super::{DEFAULT_BUFFER_SIZE, DecodedGlyph, EncoderContext, MeshEncoder, upload_buf};
use crate::{GlyphColoring, Rect};

/// Per-vertex inputs consumed by the shader: which material this vertex uses and
/// the vertex's position in the material's gradient UV space.
#[derive(bytemuck::Pod, bytemuck::Zeroable, Clone, Copy)]
#[repr(C)]
pub(crate) struct MaterialVertexAttribs {
    pub material_index: u32,
    pub gradient_uv: glam::Vec2,
}

/// How to draw a glyph in a pass.
///
/// In a pipeline where multiple layers draw the same glyph but each draws a separate
/// band of SDF distance, [`MaterialEncoder`] encodes data to draw one such layer, and this is the
/// struct to describe how to draw each glyph in one layer.
///
/// Data is packed such that multiple glyphs with the same [`GlyphMaterial`] share
/// the same memory in GPU.
///
/// # Distance units
///
/// All distance fields (`render_dist_range`, `max_alpha_dist_range`) are specified in
/// **screen pixels**. Positive values extend outward from the glyph edge, negative values
/// extend inward. For example:
/// - A 3px outer stroke uses `render_dist_range: (-3.0, 0.0)`
/// - A 2px inner stroke uses `render_dist_range: (0.0, 2.0)`
/// - A 5px outer glow fading from the edge uses `render_dist_range: (-5.0, 0.0)`
///
/// The encoder automatically converts these to the internal SDF scale based on the glyph's
/// `em_size` and atlas resolution.
#[derive(Clone, Copy, PartialEq)]
pub struct GlyphMaterial {
    /// SDF distance range (in screen pixels) over which the pass renders.
    /// Outside this range, alpha is zero. Positive = outside glyph, negative = inside.
    pub render_dist_range: (f32, f32),
    /// SDF distance range (in screen pixels) over which alpha reaches 1.0.
    /// Used for glows/shadows that have a solid-inside portion.
    pub max_alpha_dist_range: (f32, f32),
    pub color: GlyphColoring,
    /// Offset amount compared to the base position. Useful for effects like shadows.
    pub offset: glam::Vec2,
    /// Screen-space rectangle used to normalise gradient positions. Vertex positions are mapped
    /// so that vertex at the rect's top-left corner is `(0, 0)` and its bottom-right corner is
    /// `(1, 1)`.
    ///
    /// When `None`, gradient UV is computed from the glyph's `offset_in_metadata_block` and
    /// `metadata_block_size`.
    pub gradient_region: Option<Rect>,
    /// Roundness of the shape when extruded.
    ///
    /// Implemented as blend factor between pseudo-sdf and true sdf. At 0, the pseudo-sdf is used
    /// which preserves corner, at 1 the true sdf is used which is has rounder color. Only
    /// points outside of the glyph body is relevant, since inside the glyph body, the true sdf
    /// pseudo-sdf is identical.
    pub roundness: f32,
}

/// GPU-layout of [`GlyphMaterial`]. Matches the `Material` struct in `msdf.wgsl`
#[repr(C)]
#[derive(bytemuck::Pod, bytemuck::Zeroable, Clone, Copy)]
struct GlyphMaterialGpu {
    color: [f32; 4],
    gradient_color: [f32; 4],
    gradient_direction: [f32; 2],
    render_dist_range: [f32; 2],
    max_alpha_dist_range: [f32; 2],
    offset: [f32; 2],
    roundness: f32,
    _pad: [f32; 3],
}

impl GlyphMaterialGpu {
    fn from_material(m: &GlyphMaterial) -> Self {
        let (color, gradient_color, gradient_direction) = match m.color {
            GlyphColoring::Solid(color) => (color.0.to_array(), color.0.to_array(), [0.0; 2]),
            GlyphColoring::Gradient {
                primary,
                secondary,
                direction,
            } => (
                primary.0.to_array(),
                secondary.0.to_array(),
                direction.to_array(),
            ),
        };
        Self {
            color,
            gradient_color,
            gradient_direction,
            render_dist_range: [m.render_dist_range.0, m.render_dist_range.1],
            max_alpha_dist_range: [m.max_alpha_dist_range.0, m.max_alpha_dist_range.1],
            offset: m.offset.to_array(),
            roundness: m.roundness,
            _pad: [0.0; 3],
        }
    }
}

/// Encodes [`GlyphMaterial`] into GPU buffer, and material attributes for each vertex.
///
/// This helper emits material data into two buffers, which together dictates how a glyph is drawn
/// in the fragment shader:
/// - A storage buffer storing material instances.
/// - A vertex buffer, assigning material attributes to each vertex.
///
/// Materials are packed such that multiple instances of the same [`GlyphMaterial`] are
/// deduplicated and written as one instance. Each vertex is assigned a material index pointing
/// to its [`GlyphMaterial`] in the storage buffer.
///
/// The vertex buffer is parallel to the pure-geometry vertex buffer emitted by [`MeshEncoder`],
/// you can assign both vertex buffers in two slots in a draw call. This allows drawing the same
/// text in multiple layers, only needing to swap out the material vertex buffer for each layer and
/// keeping the geometry vertex buffer unchanged between draw calls. The high level API,
/// [`crate::TextRenderer`], uses this pattern to draw each effect (body fill, outer outline, inner
/// outline, outer glow, ...) with a separate [`MaterialEncoder`].
pub struct MaterialEncoder {
    content_buffer: wgpu::Buffer,
    content_write: Vec<u8>,
    index_lookup: Vec<GlyphMaterial>,

    material_vertex_buffer: wgpu::Buffer,
    material_vertex_write: Vec<u8>,

    bind_group_layout: wgpu::BindGroupLayout,
    bind_group: wgpu::BindGroup,
}

impl MaterialEncoder {
    /// Create a new material encoder with GPU buffers.
    pub fn new(device: &wgpu::Device) -> Self {
        let bind_group_layout = Self::storage_bind_group_layout(device);
        let content_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("klyff glyph material storage buffer"),
            size: DEFAULT_BUFFER_SIZE,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let material_vertex_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("klyff material vertex buffer"),
            size: DEFAULT_BUFFER_SIZE,
            usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let bind_group = make_material_bind_group(device, &bind_group_layout, &content_buffer);
        Self {
            content_buffer,
            content_write: Vec::new(),
            index_lookup: Vec::new(),
            material_vertex_buffer,
            material_vertex_write: Vec::new(),
            bind_group_layout,
            bind_group,
        }
    }

    /// Gets the bind group layout for the material storage buffer's bind group.
    pub fn storage_bind_group_layout(device: &wgpu::Device) -> wgpu::BindGroupLayout {
        device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("klyff glyph material storage buffer layout"),
            entries: &[wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Storage { read_only: true },
                    has_dynamic_offset: false,
                    min_binding_size: None,
                },
                count: None,
            }],
        })
    }

    /// Gets the vertex buffer layout for per-vertex material inputs.
    pub fn material_vertex_buffer_layout() -> wgpu::VertexBufferLayout<'static> {
        const MSDF_MATERIAL_ATTRIBUTES: [wgpu::VertexAttribute; 2] = [
            // material_index
            wgpu::VertexAttribute {
                format: wgpu::VertexFormat::Uint32,
                offset: 0,
                shader_location: 3,
            },
            // gradient_uv
            wgpu::VertexAttribute {
                format: wgpu::VertexFormat::Float32x2,
                offset: 4,
                shader_location: 4,
            },
        ];

        wgpu::VertexBufferLayout {
            array_stride: std::mem::size_of::<MaterialVertexAttribs>() as u64,
            step_mode: wgpu::VertexStepMode::Vertex,
            attributes: &MSDF_MATERIAL_ATTRIBUTES,
        }
    }

    pub fn storage_bind_group(&self) -> &wgpu::BindGroup {
        &self.bind_group
    }

    /// Gets the buffer containing each vertex's material attributes.
    ///
    /// This buffer contains the index into storage buffer [`Self::storage_bind_group`] and
    /// gradient UV for each vertex. You can bind this as a vertex buffer on a different slot
    /// to [`MeshEncoder`]'s vertex buffer to access the full data for each vertex.
    pub fn material_vertex_buffer(&self) -> &wgpu::Buffer {
        &self.material_vertex_buffer
    }

    /// Encode material data for every MSDF glyph by pairing each decoded glyph with its [`GlyphMaterial`].
    pub fn encode(
        &mut self,
        ctx: EncoderContext<'_>,
        mesh_encoder: &MeshEncoder,
        mut glyph_material_map: impl FnMut(&DecodedGlyph) -> GlyphMaterial,
    ) {
        profiling::scope!("MaterialEncoder::encode");
        self.content_write.clear();
        self.material_vertex_write.clear();
        self.index_lookup.clear();

        for decoded_glyph in mesh_encoder.decoded_glyphs_msdf() {
            let glyph_material = glyph_material_map(decoded_glyph);
            let index = {
                profiling::scope!("MaterialEncoder::dedup_lookup");
                match self.index_lookup.iter().position(|m| m == &glyph_material) {
                    Some(i) => i as u32,
                    None => {
                        let i = self.index_lookup.len() as u32;
                        self.index_lookup.push(glyph_material);
                        self.content_write.extend(bytemuck::bytes_of(
                            &GlyphMaterialGpu::from_material(&glyph_material),
                        ));
                        i
                    }
                }
            };
            let rect = decoded_glyph.glyph_rect;
            let glyph_size = rect.max - rect.min;
            for corner_offset in [
                glam::vec2(0.0, 0.0),
                glam::vec2(0.0, glyph_size.y),
                glam::vec2(glyph_size.x, 0.0),
                glam::vec2(glyph_size.x, glyph_size.y),
            ] {
                let gradient_uv = if let Some(region) = glyph_material.gradient_region {
                    let corner = rect.min + corner_offset;
                    let size = region.max - region.min;
                    (corner - region.min) / size
                } else {
                    (decoded_glyph.offset_in_metadata_block + corner_offset)
                        / decoded_glyph.metadata_block_size
                };
                self.material_vertex_write
                    .extend(bytemuck::bytes_of(&MaterialVertexAttribs {
                        material_index: index,
                        gradient_uv,
                    }));
            }
        }

        // Storage buffer must be non-empty for the pipeline to bind it. Write a single
        // zeroed material so the shader can safely index `materials[0]`.
        if self.content_write.is_empty() {
            self.content_write
                .extend(bytemuck::bytes_of(&GlyphMaterialGpu::zeroed()));
        }

        profiling::scope!("MaterialEncoder::upload");
        let reallocated = upload_buf(
            ctx.device,
            ctx.queue,
            &mut self.content_buffer,
            &self.content_write,
            wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            "klyff glyph material storage buffer",
        );
        if reallocated {
            self.bind_group =
                make_material_bind_group(ctx.device, &self.bind_group_layout, &self.content_buffer);
        }

        upload_buf(
            ctx.device,
            ctx.queue,
            &mut self.material_vertex_buffer,
            &self.material_vertex_write,
            wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
            "klyff material vertex buffer",
        );
    }
}

fn make_material_bind_group(
    device: &wgpu::Device,
    layout: &wgpu::BindGroupLayout,
    content_buffer: &wgpu::Buffer,
) -> wgpu::BindGroup {
    device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("klyff glyph material bind group"),
        layout,
        entries: &[wgpu::BindGroupEntry {
            binding: 0,
            resource: content_buffer.as_entire_binding(),
        }],
    })
}