klyff 0.1.3

use klyff_msdf::MsdfGenerator;
use skrifa::MetadataProvider;
pub use skrifa::outline::DrawError;

use super::types::{GlyphCacheKey, GlyphCacheValue};
use crate::Rect;

use super::TextureAtlas;

/// Unused texels reserved between adjacent glyph allocations, these pixels are never written to.
/// Avoids artifacts around the border of each glyph rect.
const ATLAS_GUTTER: usize = 1;

// Input

/// Identifier of glyphs that are cached in the atlas.
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
pub struct GlyphKey {
    font: fontdb::ID,
    font_size_bytes: [u8; 4],
    id: u16,
}
impl GlyphKey {
    pub fn new(font: fontdb::ID, font_size: f32, glyph_id: u16) -> Self {
        Self {
            font,
            font_size_bytes: font_size.to_le_bytes(),
            id: glyph_id,
        }
    }

    pub fn from_glyph(glyph: &cosmic_text::LayoutGlyph) -> Self {
        Self::new(glyph.font_id, glyph.font_size, glyph.glyph_id)
    }

    pub fn font_id(&self) -> fontdb::ID {
        self.font
    }

    pub fn font_size(&self) -> f32 {
        f32::from_le_bytes(self.font_size_bytes)
    }

    pub fn glyph_id(&self) -> u16 {
        self.id
    }
}

/// Contains and identify a unique font.
pub struct FontData {
    font: std::sync::Arc<cosmic_text::Font>,
    face_index: u32,
}

impl FontData {
    pub fn new(font: std::sync::Arc<cosmic_text::Font>, face_index: u32) -> Self {
        Self { font, face_index }
    }

    pub fn from_glyph(
        font_system: &mut cosmic_text::FontSystem,
        glyph: &cosmic_text::LayoutGlyph,
    ) -> Option<Self> {
        let font = font_system.get_font(glyph.font_id, glyph.font_weight)?;
        let face_index = font_system.db().face(glyph.font_id)?.index;
        Some(Self { font, face_index })
    }

    pub fn skrifa(&self) -> Option<skrifa::FontRef<'_>> {
        skrifa::FontRef::from_index(self.font.data(), self.face_index).ok()
    }

    pub fn swash(&self) -> Option<swash::FontRef<'_>> {
        swash::FontRef::from_index(self.font.data(), self.face_index as usize)
    }
}

// Output

/// The cached glyph info, directing how to render the glyph.
#[derive(Default, Debug, Clone, Copy)]
pub enum CachedGlyph {
    /// Contains nothing, no need to render at all.
    #[default]
    EmptyGlyph,
    /// Glyph was rasterzied as colored bitmaps.
    Rasterized {
        region: AtlasRegion,
        /// Glyph bounding box in em units, relative to the pen origin / baseline.
        /// X: left = pen origin. Y: positive is above baseline (font convention, Y-up).
        bounds_em: Rect,
    },
    /// Glyph has a distance field.
    DistanceField {
        region: AtlasRegion,
        em_size: f32,
        /// Glyph bounding box in em units, relative to the pen origin / baseline.
        /// X: left = pen origin. Y: positive is above baseline (font convention, Y-up).
        bounds_em: Rect,
    },
}

#[derive(Debug, Clone, Copy)]
#[non_exhaustive]
/// Region of the atlas where a glyph is stored.
///
/// Each glyph's region contains some padding.
pub struct AtlasRegion {
    pub layer: u32,
    pub min_x: u32,
    pub min_y: u32,
    pub width: u32,
    pub height: u32,
    pub padding_x: u32,
    pub padding_y: u32,
}

impl AtlasRegion {
    pub fn inner_x(&self) -> u32 {
        self.min_x + self.padding_x
    }
    pub fn inner_y(&self) -> u32 {
        self.min_y + self.padding_y
    }
    pub fn inner_width(&self) -> u32 {
        self.width - self.padding_x * 2
    }
    pub fn inner_height(&self) -> u32 {
        self.height - self.padding_y * 2
    }
}

#[derive(Debug)]
pub enum WriteGlyphError {
    OutOfAtlasSpace,
    InvalidFontData,
    GetOutlineError(DrawError),
}

impl std::fmt::Display for WriteGlyphError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            WriteGlyphError::OutOfAtlasSpace => write!(f, "Atlas has ran out of space"),
            WriteGlyphError::InvalidFontData => write!(f, "Font data could not be parsed"),
            WriteGlyphError::GetOutlineError(draw_error) => {
                write!(f, "Could not read outline of glyph: {draw_error}")
            }
        }
    }
}
impl std::error::Error for WriteGlyphError {}

impl TextureAtlas {
    /// Must be called once per frame at the beginning of frame.
    ///
    /// This method is necessary to allow the atlas keeps track of when each glyph was last used so
    /// it can remove unused glyph.
    pub fn start_of_frame(&mut self) {
        self.cache.new_tick();
    }

    /// Retrieve glyph, or generate and insert it into the atlas.
    ///
    /// # Parameters
    /// - `key`: Identifier of the glyph.
    /// - `font`: Font data. Construct one from cosmic_text's [`cosmic_text::FontSystem`].
    /// - `generator`: Cached MSDF generator.
    /// - `padding_request`: Override atlas padding `(x, y)` for distance-field glyphs.
    ///   `None` uses the atlas default.  A glyph already cached with larger padding can be
    ///   reused; a request with larger padding than the cached version evicts the old entry
    ///   and regenerates the glyph at the new size.
    /// - `device`: wgpu device.
    /// - `queue`: wgpu queue.
    /// - `encoder`: wgpu command encoder, necessary for copying texture to texture.
    ///
    /// # Returns
    /// [`CachedGlyph`] directing how to render the glyph (e.g whether glyph was inserted as a
    /// rasterized bitmap or mtsdf).
    ///
    /// # Errors
    /// Returns an error if atlas runs out of space, or if reading glyph data from font isn't
    /// possible.
    pub fn retrieve_or_generate_glyph(
        &mut self,
        key: GlyphKey,
        font: &FontData,
        generator: &mut MsdfGenerator,
        padding_request: Option<(usize, usize)>,
        wgpu_resources: (&wgpu::Device, &wgpu::Queue, &mut wgpu::CommandEncoder),
    ) -> Result<CachedGlyph, WriteGlyphError> {
        let (device, queue, encoder) = wgpu_resources;
        profiling::scope!("TextureAtlas::retrieve_or_generate_glyph");
        if self.empty_glyphs.contains(&(key.font, key.id.into())) {
            return Ok(CachedGlyph::EmptyGlyph);
        }

        let req_px = padding_request.unwrap_or((self.padding_x, self.padding_y));

        let sized_key = GlyphCacheKey {
            font_size_bytes: Some(key.font_size_bytes),
            font: key.font,
            id: key.id.into(),
        };
        let unsized_key = GlyphCacheKey {
            font_size_bytes: None,
            font: key.font,
            id: key.id.into(),
        };

        if let Some(v) = self.cache.get(&unsized_key) {
            // Reuse if the cached padding is at least as large as requested. When frozen, reuse
            // unconditionally rather than evicting and regenerating at a larger padding.
            if self.frozen
                || (v.region.padding_x as usize >= req_px.0
                    && v.region.padding_y as usize >= req_px.1)
            {
                return Ok(CachedGlyph::DistanceField {
                    region: v.region,
                    em_size: v.em_size,
                    bounds_em: v.bounds_em,
                });
            }
            // Requested padding is larger - evict and regenerate below.
            let alloc_id = v.alloc_id;
            let layer = v.region.layer as usize;
            self.cache.remove(&unsized_key);
            self.allocators[layer].deallocate(alloc_id);
        }
        if let Some(v) = self.cache.get(&sized_key) {
            return Ok(CachedGlyph::Rasterized {
                region: v.region,
                bounds_em: v.bounds_em,
            });
        }

        // Frozen atlases never generate new glyphs; an uncached glyph renders as empty.
        if self.frozen {
            return Ok(CachedGlyph::EmptyGlyph);
        }

        let Some(font_ref) = font.skrifa() else {
            return Err(WriteGlyphError::InvalidFontData);
        };

        let em_size = font_ref
            .metrics(
                skrifa::prelude::Size::unscaled(),
                skrifa::prelude::LocationRef::default(),
            )
            .units_per_em as f32;

        let read_glyph_result = {
            profiling::scope!("MsdfGenerator::read_glyph (call site)");
            generator.read_glyph(&klyff_msdf::SkrifaOutlineProvider::new(
                font_ref,
                key.id.into(),
            ))
        };
        match read_glyph_result {
            Ok(outline) => {
                let width_em = outline.width_em();
                let height_em = outline.height_em();
                let bounds_em = Rect {
                    min: glam::vec2(outline.min_x_em(), outline.min_y_em()),
                    max: glam::vec2(
                        outline.min_x_em() + width_em,
                        outline.min_y_em() + height_em,
                    ),
                };
                let inner_width = (width_em * self.ppem).round() as usize;
                let inner_height = (height_em * self.ppem).round() as usize;
                let width = inner_width + req_px.0 * 2;
                let height = inner_height + req_px.1 * 2;

                let (region, alloc_id) = self
                    .try_allocate(device, encoder, (width, height), req_px)
                    .ok_or(WriteGlyphError::OutOfAtlasSpace)?;
                self.cache.insert(
                    unsized_key,
                    GlyphCacheValue {
                        region,
                        em_size,
                        bounds_em,
                        alloc_id,
                    },
                );

                self.gpu_write.add_glyph(
                    outline,
                    klyff_msdf::GpuAtlasRegion {
                        layer: region.layer,
                        min_x: region.min_x,
                        min_y: region.min_y,
                        width: region.width,
                        height: region.height,
                        padding_x: region.padding_x,
                        padding_y: region.padding_y,
                    },
                );

                Ok(CachedGlyph::DistanceField {
                    region,
                    em_size,
                    bounds_em,
                })
            }
            Err(klyff_msdf::ReadGlyphOutlineError::ShouldRasterize) => {
                let Some(swash_ref) = font.swash() else {
                    log::warn!("Glyph {key:?} can not be rasterized: invalid swash font data.");
                    self.empty_glyphs.insert((key.font, key.id.into()));
                    return Ok(CachedGlyph::EmptyGlyph);
                };
                profiling::scope!("swash rasterize");
                let context = &mut self.scale_context;
                let mut scaler = context.builder(swash_ref).size(key.font_size()).build();
                let Some(image) = swash::scale::Render::new(&[
                    swash::scale::Source::ColorOutline(0),
                    swash::scale::Source::ColorBitmap(swash::scale::StrikeWith::BestFit),
                ])
                .render(&mut scaler, key.id) else {
                    log::warn!("Glyph {key:?} can not be rasterized.");
                    self.empty_glyphs.insert((key.font, key.id.into()));
                    return Ok(CachedGlyph::EmptyGlyph);
                };
                let width = image.placement.width as usize;
                let height = image.placement.height as usize;
                let rasterize_size = key.font_size();
                let bounds_em = Rect {
                    min: glam::vec2(
                        image.placement.left as f32 / rasterize_size,
                        (image.placement.top - height as i32) as f32 / rasterize_size,
                    ),
                    max: glam::vec2(
                        (image.placement.left + width as i32) as f32 / rasterize_size,
                        image.placement.top as f32 / rasterize_size,
                    ),
                };

                let (region, alloc_id) = self
                    .try_allocate(device, encoder, (width, height), (0, 0))
                    .ok_or(WriteGlyphError::OutOfAtlasSpace)?;
                self.cache.insert(
                    sized_key,
                    GlyphCacheValue {
                        region,
                        em_size,
                        bounds_em,
                        alloc_id,
                    },
                );

                log::trace!(
                    "Atlas: Allocating glyph ID={} at x={}, y={}, w={}, h={}, layer={}",
                    key.id,
                    region.min_x,
                    region.min_y,
                    width as u32,
                    height as u32,
                    region.layer
                );
                profiling::scope!("rasterized bitmap write_texture");
                queue.write_texture(
                    wgpu::TexelCopyTextureInfo {
                        texture: &self.texture,
                        mip_level: 0,
                        origin: wgpu::Origin3d {
                            x: region.min_x,
                            y: region.min_y,
                            z: region.layer,
                        },
                        aspect: wgpu::TextureAspect::All,
                    },
                    &image.data,
                    wgpu::TexelCopyBufferLayout {
                        offset: 0,
                        bytes_per_row: Some(width as u32 * 4),
                        rows_per_image: Some(height as u32),
                    },
                    wgpu::Extent3d {
                        width: width as u32,
                        height: height as u32,
                        depth_or_array_layers: 1,
                    },
                );

                Ok(CachedGlyph::Rasterized { region, bounds_em })
            }
            Err(klyff_msdf::ReadGlyphOutlineError::NotFound) => {
                log::warn!("Glyph {key:?} is not found in font");
                self.empty_glyphs.insert((key.font, key.id.into()));
                Ok(CachedGlyph::EmptyGlyph)
            }
            Err(klyff_msdf::ReadGlyphOutlineError::EmptyGlyph) => {
                self.empty_glyphs.insert((key.font, key.id.into()));
                Ok(CachedGlyph::EmptyGlyph)
            }
            Err(klyff_msdf::ReadGlyphOutlineError::Other(e)) => {
                Err(WriteGlyphError::GetOutlineError(e))
            }
        }
    }

    fn try_allocate(
        &mut self,
        device: &wgpu::Device,
        encoder: &mut wgpu::CommandEncoder,
        (width, height): (usize, usize),
        (padding_x, padding_y): (usize, usize),
    ) -> Option<(AtlasRegion, etagere::AllocId)> {
        profiling::scope!("TextureAtlas::try_allocate");
        // Reserve a gutter of unwritten texels around each allocation so linear
        // sampling at a glyph's UV edge never reaches a neighbour. The region we
        // record (and that UVs are built from) covers only the inner `width`/`height`.
        let alloc_w = width + ATLAS_GUTTER;
        let alloc_h = height + ATLAS_GUTTER;
        loop {
            for (layer, allocator) in self.allocators.iter_mut().enumerate() {
                if let Some(allocation) =
                    allocator.allocate(etagere::size2(alloc_w as i32, alloc_h as i32))
                {
                    return Some((
                        AtlasRegion {
                            layer: layer as u32,
                            min_x: allocation.rectangle.min.x as u32,
                            min_y: allocation.rectangle.min.y as u32,
                            width: width as u32,
                            height: height as u32,
                            padding_x: padding_x as u32,
                            padding_y: padding_y as u32,
                        },
                        allocation.id,
                    ));
                }
            }

            let current_tick = self.cache.current_tick();
            let removable_lru = self
                .cache
                .peek_lru()
                .and_then(|(lru_key, lru_glyph, lru_tick)| {
                    (lru_tick + self.retain_glyphs_frame_duration < current_tick)
                        .then_some((*lru_key, *lru_glyph))
                });

            if let Some((lru_key, v)) = removable_lru {
                self.cache.remove(&lru_key);
                self.allocators[v.region.layer as usize].deallocate(v.alloc_id);
            } else if !self.allow_grow || self.current_size == self.max_size {
                return None;
            } else {
                self.grow_size(device, encoder);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use skrifa::MetadataProvider;

    use super::*;
    use crate::{TextureAtlas, TextureAtlasDescriptor};

    const OPEN_SANS: &[u8] = include_bytes!("../../../../fonts/OpenSans-Regular.ttf");

    async fn request_device() -> Option<(wgpu::Device, wgpu::Queue)> {
        let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor::default());
        let adapter = instance
            .request_adapter(&wgpu::RequestAdapterOptions {
                power_preference: wgpu::PowerPreference::None,
                compatible_surface: None,
                force_fallback_adapter: false,
            })
            .await
            .ok()?;
        adapter
            .request_device(&wgpu::DeviceDescriptor::default())
            .await
            .ok()
    }

    /// Generate a single glyph through the atlas, returning the resulting [`CachedGlyph`].
    fn request_glyph(
        atlas: &mut TextureAtlas,
        font: &FontData,
        generator: &mut klyff_msdf::MsdfGenerator,
        font_id: fontdb::ID,
        glyph_id: u16,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
    ) -> CachedGlyph {
        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
        atlas
            .retrieve_or_generate_glyph(
                GlyphKey::new(font_id, 32.0, glyph_id),
                font,
                generator,
                None,
                (device, queue, &mut encoder),
            )
            .unwrap()
    }

    #[test]
    fn frozen_atlas_serves_cached_but_not_new_glyphs() {
        let Some((device, queue)) = pollster::block_on(request_device()) else {
            eprintln!("no wgpu device available; skipping freeze test");
            return;
        };

        let mut db = fontdb::Database::new();
        db.load_font_data(OPEN_SANS.to_vec());
        let font_id = db.faces().next().unwrap().id;
        let face_index = db.face(font_id).unwrap().index;
        let mut font_system =
            cosmic_text::FontSystem::new_with_locale_and_db("en-US".to_string(), db);
        let font = FontData::new(
            font_system
                .get_font(font_id, fontdb::Weight::NORMAL)
                .unwrap(),
            face_index,
        );

        let charmap = skrifa::FontRef::new(OPEN_SANS).unwrap().charmap();
        let gid_a = charmap.map('a').unwrap().to_u32() as u16;
        let gid_b = charmap.map('b').unwrap().to_u32() as u16;

        let mut atlas = TextureAtlas::new(&device, TextureAtlasDescriptor::default());
        let mut generator = klyff_msdf::MsdfGenerator::new();

        // Cache glyph 'a' while unfrozen.
        atlas.start_of_frame();
        let cached_a = request_glyph(
            &mut atlas,
            &font,
            &mut generator,
            font_id,
            gid_a,
            &device,
            &queue,
        );
        assert!(matches!(cached_a, CachedGlyph::DistanceField { .. }));

        atlas.freeze();
        assert!(atlas.is_frozen());

        // A cached glyph is still served while frozen...
        let frozen_a = request_glyph(
            &mut atlas,
            &font,
            &mut generator,
            font_id,
            gid_a,
            &device,
            &queue,
        );
        assert!(matches!(frozen_a, CachedGlyph::DistanceField { .. }));

        // ...but a new, uncached glyph is not generated.
        let frozen_b = request_glyph(
            &mut atlas,
            &font,
            &mut generator,
            font_id,
            gid_b,
            &device,
            &queue,
        );
        assert!(matches!(frozen_b, CachedGlyph::EmptyGlyph));

        // Unfreezing re-enables generation for the previously missing glyph.
        atlas.unfreeze();
        assert!(!atlas.is_frozen());
        let unfrozen_b = request_glyph(
            &mut atlas,
            &font,
            &mut generator,
            font_id,
            gid_b,
            &device,
            &queue,
        );
        assert!(matches!(unfrozen_b, CachedGlyph::DistanceField { .. }));
    }
}