klyff 0.1.2

use std::collections::HashSet;

use super::{DEFAULT_BUFFER_SIZE, DecodedGlyph, EncoderContext, upload_buf};
use crate::{CachedGlyph, CustomGlyph, CustomGlyphFont, FontData, GlyphKey, Rect, Text};

#[derive(bytemuck::Pod, bytemuck::Zeroable, Clone, Copy)]
#[repr(C)]
pub(crate) struct VertexDataRasterized {
    pub uvw: glam::Vec3,
    pub screen_pos: glam::Vec2,
}

#[derive(bytemuck::Pod, bytemuck::Zeroable, Clone, Copy)]
#[repr(C)]
pub(crate) struct VertexDataMsdf {
    pub uvw: glam::Vec3,
    pub screen_pos: glam::Vec2,
    pub font_size_px: f32,
}

/// Encodes [`Text`] into mesh geometry.
///
/// This helper emits the text's geomery into vertex / index buffers.
/// It also saves the glyphs it decoded which can be read by another pass to assign additional
/// attributes for each vertex. Typically you want to create a second vertex buffer, parallel to
/// this struct's pure-geomery vertex buffer, and assign both of them in two slots a draw call.
///
/// If you use the default provided shader, then [`super::MaterialEncoder`] implements this
/// pattern. A single [`MeshEncoder`] and its pure-geomery vertex / index buffer can be reused
/// across many [`super::MaterialEncoder`], reducing context switch overhead.
pub struct MeshEncoder {
    // MSDF mesh
    msdf_vertex_write: Vec<u8>,
    msdf_index_write: Vec<u16>,
    msdf_vertex_buf: wgpu::Buffer,
    msdf_index_buf: wgpu::Buffer,
    msdf_index_count: u32,

    // Rasterized mesh
    rasterized_vertex_write: Vec<u8>,
    rasterized_index_write: Vec<u16>,
    rasterized_vertex_buf: wgpu::Buffer,
    rasterized_index_buf: wgpu::Buffer,
    rasterized_index_count: u32,

    // Per-glyph info for effect encoders to iterate.
    decoded_glyphs_msdf: Vec<DecodedGlyph>,
    decoded_glyphs_rasterized: Vec<DecodedGlyph>,

    // Custom glyphs detected during the last `encode`, for application code to render.
    custom_glyph_font: Option<fontdb::ID>,
    custom_glyphs: Vec<CustomGlyph>,

    errored_glyph_ids: HashSet<u16>,
    errored_font_ids: HashSet<fontdb::ID>,
}

impl MeshEncoder {
    /// Create a new mesh encoder with GPU buffers.
    pub fn new(device: &wgpu::Device) -> Self {
        let make_buf = |label: &str, usage: wgpu::BufferUsages| {
            device.create_buffer(&wgpu::BufferDescriptor {
                label: Some(label),
                size: DEFAULT_BUFFER_SIZE,
                usage: usage | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            })
        };

        Self {
            msdf_vertex_write: Vec::new(),
            msdf_index_write: Vec::new(),
            msdf_vertex_buf: make_buf("klyff msdf vertex buffer", wgpu::BufferUsages::VERTEX),
            msdf_index_buf: make_buf("klyff msdf index buffer", wgpu::BufferUsages::INDEX),
            msdf_index_count: 0,
            rasterized_vertex_write: Vec::new(),
            rasterized_index_write: Vec::new(),
            rasterized_vertex_buf: make_buf(
                "klyff rasterized vertex buffer",
                wgpu::BufferUsages::VERTEX,
            ),
            rasterized_index_buf: make_buf(
                "klyff rasterized index buffer",
                wgpu::BufferUsages::INDEX,
            ),
            rasterized_index_count: 0,
            decoded_glyphs_msdf: vec![],
            decoded_glyphs_rasterized: vec![],
            custom_glyph_font: None,
            custom_glyphs: vec![],
            errored_glyph_ids: HashSet::new(),
            errored_font_ids: HashSet::new(),
        }
    }

    /// Register the font whose glyphs should be treated as custom glyphs.
    ///
    /// Pass the handle returned by [`crate::setup_custom_glyph_font`]. Glyphs shaped with this font
    /// are not rendered; they are collected and surfaced via [`Self::custom_glyphs`]. Pass `None`
    /// to disable custom-glyph detection.
    pub fn set_custom_glyph_font(&mut self, font: Option<CustomGlyphFont>) {
        self.custom_glyph_font = font.map(|f| f.id);
    }

    /// Custom glyphs detected during the last call to [`Self::encode`].
    ///
    /// Empty unless a custom-glyph font was registered with [`Self::set_custom_glyph_font`].
    pub fn custom_glyphs(&self) -> &[CustomGlyph] {
        &self.custom_glyphs
    }

    /// Gets the vertex buffer layout for MSDF rendering path.
    ///
    /// The shader code declaration is expected to be:
    /// ```wgsl
    /// struct VertexInput {
    ///     @location(0) uvw: vec3f,
    ///     @location(1) screen_pos: vec2f,
    ///     @location(2) font_size_px: f32,
    /// }
    /// ```
    pub fn msdf_vertex_buffer_layout() -> wgpu::VertexBufferLayout<'static> {
        const MSDF_MESH_ATTRIBUTES: [wgpu::VertexAttribute; 3] = [
            // uvw
            wgpu::VertexAttribute {
                format: wgpu::VertexFormat::Float32x3,
                offset: 0,
                shader_location: 0,
            },
            // screen_pos
            wgpu::VertexAttribute {
                format: wgpu::VertexFormat::Float32x2,
                offset: 12,
                shader_location: 1,
            },
            // font_size_px
            wgpu::VertexAttribute {
                format: wgpu::VertexFormat::Float32,
                offset: 20,
                shader_location: 2,
            },
        ];

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

    /// Gets the vertex buffer layout for MSDF rendering path.
    ///
    /// The shader code declaration is expected to be:
    /// ```wgsl
    /// struct VertexInput {
    ///     @location(0) uvw: vec3f,
    ///     @location(1) screen_pos: vec2f,
    /// }
    /// ```
    pub fn rasterized_vertex_buffer_layout() -> wgpu::VertexBufferLayout<'static> {
        const RASTERIZED_MESH_ATTRIBUTES: [wgpu::VertexAttribute; 2] = [
            wgpu::VertexAttribute {
                format: wgpu::VertexFormat::Float32x3,
                offset: 0,
                shader_location: 0,
            },
            wgpu::VertexAttribute {
                format: wgpu::VertexFormat::Float32x2,
                offset: 12,
                shader_location: 1,
            },
        ];

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

    /// Decoded glyphs produced during the last call to [`Self::encode`].
    pub fn decoded_glyphs_msdf(&self) -> &[DecodedGlyph] {
        &self.decoded_glyphs_msdf
    }

    /// Decoded glyphs produced during the last call to [`Self::encode`].
    pub fn decoded_glyphs_rasterized(&self) -> &[DecodedGlyph] {
        &self.decoded_glyphs_rasterized
    }

    /// The number of indices emitted in the last call to [`Self::encode`].
    ///
    /// You can use this value when submitting a draw call with [`wgpu::RenderPass`].
    pub fn msdf_index_count(&self) -> u32 {
        self.msdf_index_count
    }

    /// Returns the MSDF vertex buffer.
    pub fn msdf_vertex_buffer(&self) -> &wgpu::Buffer {
        &self.msdf_vertex_buf
    }

    /// Returns the MSDF index buffer.
    pub fn msdf_index_buffer(&self) -> &wgpu::Buffer {
        &self.msdf_index_buf
    }

    /// The number of indices emitted in the last call to [`Self::encode`].
    ///
    /// You can use this value when submitting a draw call with [`wgpu::RenderPass`].
    pub fn rasterized_index_count(&self) -> u32 {
        self.rasterized_index_count
    }

    /// Returns the rasterized glyph vertex buffer.
    pub fn rasterized_vertex_buffer(&self) -> &wgpu::Buffer {
        &self.rasterized_vertex_buf
    }

    /// Returns the rasterized glyph index buffer.
    pub fn rasterized_index_buffer(&self) -> &wgpu::Buffer {
        &self.rasterized_index_buf
    }

    /// Encode geometry for the supplied texts.
    pub fn encode<'txt>(
        &mut self,
        ctx: EncoderContext<'_>,
        msdf_generator: &mut klyff_msdf::MsdfGenerator,
        texts: impl IntoIterator<Item = Text<'txt>>,
        mut per_glyph_transform: impl FnMut(&DecodedGlyph, &mut [glam::Vec2; 4]),
    ) {
        profiling::scope!("MeshEncoder::encode");
        ctx.atlas.start_of_frame();
        self.msdf_vertex_write.clear();
        self.msdf_index_write.clear();
        self.rasterized_vertex_write.clear();
        self.rasterized_index_write.clear();
        self.decoded_glyphs_msdf.clear();
        self.decoded_glyphs_rasterized.clear();
        self.custom_glyphs.clear();

        let mut msdf_vertex_idx: u16 = 0;
        let mut rasterized_vertex_idx: u16 = 0;

        {
            profiling::scope!("MeshEncoder::layout_loop");
            for text in texts {
                let padding = text.custom_padding;
                let mut glyph_index = 0;
                let start_index_msdf = self.decoded_glyphs_msdf.len();
                let start_index_rasterized = self.decoded_glyphs_rasterized.len();
                let mut text_rect_tight = Rect {
                    min: glam::Vec2::MAX,
                    max: glam::Vec2::MIN,
                };
                let mut current_metadata_block_metadata: Option<usize> = None;
                let mut current_metadata_block_start_msdf = start_index_msdf;
                let mut current_metadata_block_start_rasterized = start_index_rasterized;
                let mut current_metadata_block_rect = Rect {
                    min: glam::Vec2::MAX,
                    max: glam::Vec2::MIN,
                };

                for layout_run in text.text_buffer.layout_runs() {
                    for glyph in layout_run.glyphs {
                        let font_size_px = glyph.font_size * text.scale;

                        // Custom glyphs are empty placeholders identified by their font. Record
                        // their layout box for the application to render and skip atlas processing.
                        if Some(glyph.font_id) == self.custom_glyph_font {
                            let pen_x = text.region.min.x + (glyph.x + glyph.x_offset) * text.scale;
                            let baseline_y = text.region.min.y
                                + (layout_run.line_y + glyph.y + glyph.y_offset) * text.scale;
                            let width = glyph.w * text.scale;
                            let height =
                                glyph.line_height_opt.unwrap_or(glyph.font_size) * text.scale;
                            let descent_offset = height * 0.15;

                            self.custom_glyphs.push(CustomGlyph {
                                id: glyph.metadata as u64,
                                text_id: text.id,
                                rect: Rect {
                                    min: glam::vec2(pen_x, baseline_y - height + descent_offset),
                                    max: glam::vec2(pen_x + width, baseline_y + descent_offset),
                                },
                                font_size_px,
                            });

                            // TODO: this skips important processing such as metadata block, text
                            // rect tight -> hidden bug
                            continue;
                        }

                        let Some(font_data) = FontData::from_glyph(ctx.font_system, glyph) else {
                            if !self.errored_font_ids.contains(&glyph.font_id) {
                                log::warn!("Font {:?} not found, skipping glyph", glyph.font_id);
                                self.errored_font_ids.insert(glyph.font_id);
                            }
                            continue;
                        };
                        let cached_glyph = match ctx.atlas.retrieve_or_generate_glyph(
                            GlyphKey::from_glyph(glyph),
                            &font_data,
                            msdf_generator,
                            padding,
                            (ctx.device, ctx.queue, ctx.cmd_encoder),
                        ) {
                            Ok(g) => g,
                            Err(e) => {
                                if !self.errored_glyph_ids.contains(&glyph.glyph_id) {
                                    log::warn!("Error processing glyph: {}", e);
                                    self.errored_glyph_ids.insert(glyph.glyph_id);
                                }
                                continue;
                            }
                        };

                        let (is_msdf, region, bounds_em) = match cached_glyph {
                            CachedGlyph::EmptyGlyph => continue,
                            CachedGlyph::Rasterized { region, bounds_em } => {
                                (false, region, bounds_em)
                            }
                            CachedGlyph::DistanceField {
                                region, bounds_em, ..
                            } => (true, region, bounds_em),
                        };

                        let glyph_size = bounds_em.max - bounds_em.min;
                        let glyph_w = glyph_size.x * glyph.font_size * text.scale;
                        let glyph_h = glyph_size.y * glyph.font_size * text.scale;
                        let (atlas_width, atlas_height, _) = ctx.atlas.atlas_size();

                        let pen_x = text.region.min.x + (glyph.x + glyph.x_offset) * text.scale;
                        let pen_y = text.region.min.y
                            + (layout_run.line_y + glyph.y + glyph.y_offset) * text.scale;
                        let glyph_x = pen_x + bounds_em.min.x * glyph.font_size * text.scale;
                        let glyph_y = pen_y - bounds_em.max.y * glyph.font_size * text.scale;

                        text_rect_tight.min.x = text_rect_tight.min.x.min(glyph_x);
                        text_rect_tight.min.y = text_rect_tight.min.y.min(glyph_y);
                        text_rect_tight.max.x = text_rect_tight.max.x.max(glyph_x + glyph_w);
                        text_rect_tight.max.y = text_rect_tight.max.y.max(glyph_y + glyph_h);

                        let (rect, uv) = if is_msdf {
                            let ppem = ctx.atlas.ppem();
                            let expand_x =
                                region.padding_x as f32 / ppem * glyph.font_size * text.scale;
                            let expand_y =
                                region.padding_y as f32 / ppem * glyph.font_size * text.scale;
                            (
                                Rect {
                                    min: glam::vec2(glyph_x - expand_x, glyph_y - expand_y),
                                    max: glam::vec2(
                                        glyph_x + glyph_w + expand_x,
                                        glyph_y + glyph_h + expand_y,
                                    ),
                                },
                                Rect {
                                    min: glam::vec2(
                                        region.min_x as f32 / atlas_width as f32,
                                        region.min_y as f32 / atlas_height as f32,
                                    ),
                                    max: glam::vec2(
                                        (region.min_x + region.width) as f32 / atlas_width as f32,
                                        (region.min_y + region.height) as f32 / atlas_height as f32,
                                    ),
                                },
                            )
                        } else {
                            (
                                Rect {
                                    min: glam::vec2(glyph_x, glyph_y),
                                    max: glam::vec2(glyph_x + glyph_w, glyph_y + glyph_h),
                                },
                                Rect {
                                    min: glam::vec2(
                                        region.inner_x() as f32 / atlas_width as f32,
                                        region.inner_y() as f32 / atlas_height as f32,
                                    ),
                                    max: glam::vec2(
                                        (region.inner_x() + region.inner_width()) as f32
                                            / atlas_width as f32,
                                        (region.inner_y() + region.inner_height()) as f32
                                            / atlas_height as f32,
                                    ),
                                },
                            )
                        };

                        let decoded_glyph = DecodedGlyph {
                            glyph_rect: rect,
                            text_id: text.id,
                            metadata: glyph.metadata,
                            glyph_index_in_text: glyph_index,
                            total_glyphs_in_text: 0,
                            text_rect: text.region,
                            font_size_px,
                            uvw: (uv, region.layer),
                            // Temporary value, will be assigned after iterating through all glyphs of
                            // text.
                            text_rect_tight: Rect {
                                min: glam::Vec2::ZERO,
                                max: glam::Vec2::ZERO,
                            },
                            // Temporary values, will be assigned after iterating through all
                            // glyphs of attr block.
                            metadata_block_size: glam::Vec2::ZERO,
                            offset_in_metadata_block: glam::Vec2::ZERO,
                        };

                        // When metadata changes, finalize the previous attr block
                        if let Some(curr_metadata) = current_metadata_block_metadata
                            && curr_metadata != glyph.metadata
                        {
                            let block_size =
                                current_metadata_block_rect.max - current_metadata_block_rect.min;

                            // Backfill metadata_block_size and offset_in_metadata_block for all glyphs in
                            // the completed block
                            for g in
                                &mut self.decoded_glyphs_msdf[current_metadata_block_start_msdf..]
                            {
                                g.metadata_block_size = block_size;
                                g.offset_in_metadata_block =
                                    g.glyph_rect.min - current_metadata_block_rect.min;
                            }
                            for g in &mut self.decoded_glyphs_rasterized
                                [current_metadata_block_start_rasterized..]
                            {
                                g.metadata_block_size = block_size;
                                g.offset_in_metadata_block =
                                    g.glyph_rect.min - current_metadata_block_rect.min;
                            }

                            // Reset for new block
                            current_metadata_block_start_msdf = self.decoded_glyphs_msdf.len();
                            current_metadata_block_start_rasterized =
                                self.decoded_glyphs_rasterized.len();
                            current_metadata_block_rect = Rect {
                                min: glam::Vec2::MAX,
                                max: glam::Vec2::MIN,
                            };
                        }
                        current_metadata_block_metadata = Some(glyph.metadata);

                        // Expand attr block rect to include this glyph
                        current_metadata_block_rect.min = current_metadata_block_rect
                            .min
                            .min(decoded_glyph.glyph_rect.min);
                        current_metadata_block_rect.max = current_metadata_block_rect
                            .max
                            .max(decoded_glyph.glyph_rect.max);

                        if is_msdf {
                            self.decoded_glyphs_msdf.push(decoded_glyph);
                        } else {
                            self.decoded_glyphs_rasterized.push(decoded_glyph);
                        }

                        glyph_index += 1;
                    }
                }

                // Finalize the last attr block for this text
                let block_size = current_metadata_block_rect.max - current_metadata_block_rect.min;
                for g in &mut self.decoded_glyphs_msdf[current_metadata_block_start_msdf..] {
                    g.metadata_block_size = block_size;
                    g.offset_in_metadata_block = g.glyph_rect.min - current_metadata_block_rect.min;
                }
                for g in
                    &mut self.decoded_glyphs_rasterized[current_metadata_block_start_rasterized..]
                {
                    g.metadata_block_size = block_size;
                    g.offset_in_metadata_block = g.glyph_rect.min - current_metadata_block_rect.min;
                }

                for glyph in &mut self.decoded_glyphs_msdf[start_index_msdf..] {
                    glyph.text_rect_tight = text_rect_tight;
                    glyph.total_glyphs_in_text = glyph_index;
                }
                for glyph in &mut self.decoded_glyphs_rasterized[start_index_rasterized..] {
                    glyph.text_rect_tight = text_rect_tight;
                    glyph.total_glyphs_in_text = glyph_index;
                }
            }
        }

        {
            profiling::scope!("MeshEncoder::emit_vertices");
            for glyph in self.decoded_glyphs_msdf.iter() {
                let (uv, layer) = &glyph.uvw;
                let rect = &glyph.glyph_rect;
                let uvws = [
                    glam::vec3(uv.min.x, uv.min.y, *layer as f32),
                    glam::vec3(uv.min.x, uv.max.y, *layer as f32),
                    glam::vec3(uv.max.x, uv.min.y, *layer as f32),
                    glam::vec3(uv.max.x, uv.max.y, *layer as f32),
                ];
                let mut positions = [
                    glam::vec2(rect.min.x, rect.min.y),
                    glam::vec2(rect.min.x, rect.max.y),
                    glam::vec2(rect.max.x, rect.min.y),
                    glam::vec2(rect.max.x, rect.max.y),
                ];
                per_glyph_transform(glyph, &mut positions);
                for (uvw, screen_pos) in std::iter::zip(uvws, positions) {
                    let vertex = VertexDataMsdf {
                        uvw,
                        screen_pos,
                        font_size_px: glyph.font_size_px,
                    };
                    self.msdf_vertex_write.extend(bytemuck::bytes_of(&vertex));
                }
                self.msdf_index_write.extend([
                    msdf_vertex_idx,
                    msdf_vertex_idx + 1,
                    msdf_vertex_idx + 2,
                    msdf_vertex_idx + 2,
                    msdf_vertex_idx + 1,
                    msdf_vertex_idx + 3,
                ]);
                msdf_vertex_idx += 4;
            }

            for glyph in self.decoded_glyphs_rasterized.iter() {
                let (uv, layer) = &glyph.uvw;
                let rect = &glyph.glyph_rect;
                let uvws = [
                    glam::vec3(uv.min.x, uv.min.y, *layer as f32),
                    glam::vec3(uv.min.x, uv.max.y, *layer as f32),
                    glam::vec3(uv.max.x, uv.min.y, *layer as f32),
                    glam::vec3(uv.max.x, uv.max.y, *layer as f32),
                ];
                let mut positions = [
                    glam::vec2(rect.min.x, rect.min.y),
                    glam::vec2(rect.min.x, rect.max.y),
                    glam::vec2(rect.max.x, rect.min.y),
                    glam::vec2(rect.max.x, rect.max.y),
                ];
                per_glyph_transform(glyph, &mut positions);
                for (uvw, screen_pos) in std::iter::zip(uvws, positions) {
                    let vertex = VertexDataRasterized { uvw, screen_pos };
                    self.rasterized_vertex_write
                        .extend(bytemuck::bytes_of(&vertex));
                }
                self.rasterized_index_write.extend([
                    rasterized_vertex_idx,
                    rasterized_vertex_idx + 1,
                    rasterized_vertex_idx + 2,
                    rasterized_vertex_idx + 2,
                    rasterized_vertex_idx + 1,
                    rasterized_vertex_idx + 3,
                ]);
                rasterized_vertex_idx += 4;
            }
        }

        ctx.atlas
            .write_glyphs(ctx.device, ctx.queue, ctx.cmd_encoder);

        profiling::scope!("MeshEncoder::upload");
        self.msdf_index_count = self.msdf_index_write.len() as u32;
        if !self.msdf_vertex_write.is_empty() {
            upload_buf(
                ctx.device,
                ctx.queue,
                &mut self.msdf_vertex_buf,
                &self.msdf_vertex_write,
                wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
                "klyff msdf vertex buffer",
            );
            upload_buf(
                ctx.device,
                ctx.queue,
                &mut self.msdf_index_buf,
                bytemuck::cast_slice(&self.msdf_index_write),
                wgpu::BufferUsages::INDEX | wgpu::BufferUsages::COPY_DST,
                "klyff msdf index buffer",
            );
        }

        self.rasterized_index_count = self.rasterized_index_write.len() as u32;
        if !self.rasterized_vertex_write.is_empty() {
            upload_buf(
                ctx.device,
                ctx.queue,
                &mut self.rasterized_vertex_buf,
                &self.rasterized_vertex_write,
                wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
                "klyff rasterized vertex buffer",
            );
            upload_buf(
                ctx.device,
                ctx.queue,
                &mut self.rasterized_index_buf,
                bytemuck::cast_slice(&self.rasterized_index_write),
                wgpu::BufferUsages::INDEX | wgpu::BufferUsages::COPY_DST,
                "klyff rasterized index buffer",
            );
        }
    }
}