klyff 0.1.2

use std::collections::HashSet;

use klyff_msdf::MtsdfGpuWriter;

use super::cache::Cache;
use super::size::{AtlasSize, grow_size};
use crate::Rect;

/// Dynamic texture atlas that contains computed glyph bitmaps or distance fields.
///
/// Each glyph is stored either a distance field, or rasterized bitmap in the atlas.
pub struct TextureAtlas {
    pub(super) bind_group: wgpu::BindGroup,

    pub(super) texture: wgpu::Texture,
    pub(super) sampler: wgpu::Sampler,
    pub(super) allocators: Vec<etagere::BucketedAtlasAllocator>,
    pub(super) cache: Cache<GlyphCacheKey, GlyphCacheValue>,
    pub(super) empty_glyphs: HashSet<(fontdb::ID, skrifa::GlyphId)>,
    pub(super) gpu_write: MtsdfGpuWriter,

    pub(super) current_size: wgpu::Extent3d,
    pub(super) max_size: wgpu::Extent3d,
    pub(super) ppem: f32,
    pub(super) padding_x: usize,
    pub(super) padding_y: usize,
    pub(super) allow_grow: bool,
    pub(super) retain_glyphs_frame_duration: u64,
    pub(super) frozen: bool,

    pub(super) scale_context: swash::scale::ScaleContext,
}

/// Identifier of a cached glyph.
///
/// The font size field is [`None`] if the glyph is stored as a distance field (since it applies to
/// any font size), and [`Some`] if the glyph is stored as a rasterized bitmap (since it only
/// applies to one specific font size).
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
pub(super) struct GlyphCacheKey {
    pub(super) font: fontdb::ID,
    pub(super) font_size_bytes: Option<[u8; 4]>,
    pub(super) id: skrifa::GlyphId,
}

/// Information about cached glyph.
///
/// Contains necessary information to render the glyph, and link to the allocated region in the
/// atlas managed by etagere.
#[derive(Debug, Clone, Copy)]
pub(super) struct GlyphCacheValue {
    /// Region in the atlas where the glyph was written to.
    pub(super) region: super::AtlasRegion,
    /// Size of 1em of the glyph's font. Used for reversing the msdf to coverage during rendering.
    pub(super) em_size: f32,
    /// Bounding box in em units, relative to the pen origin / baseline
    /// (min x = pen origin, Y: positive is above baseline).
    pub(super) bounds_em: Rect,
    /// etagere's allocation id.
    pub(super) alloc_id: etagere::AllocId,
}

/// Descriptor for how to initialize a [`TextureAtlas`].
///
/// The default option creates an atlas with the maximum allowed texture size, and grow into a
/// texture array if it runs out.
pub struct TextureAtlasDescriptor {
    /// Initial size of the atlas. Will be clamped to maximum allowed texture size by the device.
    pub initial_size: AtlasSize,
    /// Maximum size the atlas can grow to. Will be clamped to maximum texture size allowed by the device.
    /// If set to None then the maximum size allowed by the device will be used.
    pub max_size: Option<AtlasSize>,
    /// Whether to allow growing the atlas size dynamically during runtime.
    pub allow_grow: bool,
    /// The minimum number of frame to retain an unused glyph before it's removed from atlas.
    /// Glyphs are only removed when the atlas runs out of space, this configuration guarantees
    /// all glyphs are alive for some duration.
    pub retain_glyphs_frame_duration: u64,
    /// Specify a custom scaling for glyphs. If set to [`None`], [`super::DEFAULT_PPEM`] will be
    /// used.
    pub pixel_per_em: Option<f32>,
    /// Specify a custom scaling for glyphs. If set to [`None`], [`super::DEFAULT_PADDING`] will be
    /// used.
    pub glyph_padding: Option<(usize, usize)>,
}

impl Default for TextureAtlasDescriptor {
    fn default() -> Self {
        Self {
            initial_size: AtlasSize::MaxSize { layer_count: 1 },
            max_size: None,
            allow_grow: true,
            retain_glyphs_frame_duration: 0,
            pixel_per_em: None,
            glyph_padding: None,
        }
    }
}

impl TextureAtlas {
    /// Creates a new texture atlas.
    ///
    /// This maintains a large texture containing to cache rendered glyphs. Create one and reuse
    /// for all text rendering.
    pub fn new(device: &wgpu::Device, descriptor: TextureAtlasDescriptor) -> Self {
        let device_max_size = device.limits().max_texture_dimension_2d;
        let device_max_layers = device.limits().max_texture_array_layers;
        let size =
            atlas_size_to_extent(descriptor.initial_size, device_max_size, device_max_layers);
        let allocators = (0..size.depth_or_array_layers)
            .map(|_| {
                etagere::BucketedAtlasAllocator::new(etagere::size2(
                    size.width as i32,
                    size.height as i32,
                ))
            })
            .collect();

        let usage = texture_usage(descriptor.allow_grow);
        let (texture, sampler) = create_texture(device, size, usage);
        let bind_group = Self::create_bind_group(device, &texture, &sampler);

        let max_size = descriptor
            .max_size
            .map(|s| atlas_size_to_extent(s, device_max_size, device_max_layers))
            .unwrap_or(wgpu::Extent3d {
                width: device_max_size,
                height: device_max_size,
                depth_or_array_layers: device_max_layers,
            });

        Self {
            texture,
            sampler,
            allocators,
            empty_glyphs: Default::default(),
            cache: Cache::new(),
            current_size: size,
            max_size,
            ppem: descriptor.pixel_per_em.unwrap_or(super::DEFAULT_PPEM),
            padding_x: descriptor
                .glyph_padding
                .map(|x| x.0)
                .unwrap_or(super::DEFAULT_PADDING),
            padding_y: descriptor
                .glyph_padding
                .map(|x| x.1)
                .unwrap_or(super::DEFAULT_PADDING),
            allow_grow: descriptor.allow_grow,
            retain_glyphs_frame_duration: descriptor.retain_glyphs_frame_duration,
            frozen: false,
            bind_group,
            scale_context: swash::scale::ScaleContext::new(),
            gpu_write: MtsdfGpuWriter::new(device),
        }
    }

    /// Returns the bind group layout that all atlas bind groups must conform to.
    ///
    /// This layout is the same regardless of atlas size and can therefore be
    /// created once and shared across rendering pipelines.
    pub fn bind_group_layout(device: &wgpu::Device) -> wgpu::BindGroupLayout {
        device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("klyff atlas bind group layout"),
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Texture {
                        multisampled: false,
                        view_dimension: wgpu::TextureViewDimension::D2Array,
                        sample_type: wgpu::TextureSampleType::Float { filterable: true },
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
                    count: None,
                },
            ],
        })
    }

    pub(super) fn create_bind_group(
        device: &wgpu::Device,
        texture: &wgpu::Texture,
        sampler: &wgpu::Sampler,
    ) -> wgpu::BindGroup {
        let view = texture.create_view(&wgpu::TextureViewDescriptor {
            label: Some("klyff atlas texture view"),
            dimension: Some(wgpu::TextureViewDimension::D2Array),
            ..Default::default()
        });
        device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("klyff atlas bind group"),
            layout: &Self::bind_group_layout(device),
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::TextureView(&view),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::Sampler(sampler),
                },
            ],
        })
    }

    pub(super) fn grow_size(&mut self, device: &wgpu::Device, encoder: &mut wgpu::CommandEncoder) {
        profiling::scope!("TextureAtlas::grow_size");
        let usage = self.texture.usage();
        let old_size = self.current_size;
        self.current_size = grow_size(self.current_size, self.max_size);
        log::trace!(
            "Growing texture size from {old_size:?} to {:?}",
            self.current_size
        );
        let (new_texture, new_sampler) = create_texture(device, self.current_size, usage);
        let new_bind_group = Self::create_bind_group(device, &new_texture, &new_sampler);

        encoder.copy_texture_to_texture(
            self.texture.as_image_copy(),
            new_texture.as_image_copy(),
            self.texture.size(),
        );

        self.texture = new_texture;
        self.sampler = new_sampler;
        self.bind_group = new_bind_group;
        for allocator in self.allocators.iter_mut() {
            allocator.grow(etagere::size2(
                self.current_size.width as i32,
                self.current_size.height as i32,
            ));
        }
    }

    /// Returns the atlas pixels-per-em used when generating distance fields.
    ///
    /// Divide `region.padding_x` by this value to get the maximum outer effect
    /// radius in em units for glyphs stored in this atlas.
    pub fn ppem(&self) -> f32 {
        self.ppem
    }

    /// Gets the current atlas size.
    pub fn atlas_size(&self) -> (u32, u32, u32) {
        (
            self.texture.width(),
            self.texture.height(),
            self.texture.depth_or_array_layers(),
        )
    }

    /// Gets the bind group containing atlas texture and sample for rendering.
    pub fn bind_group(&self) -> &wgpu::BindGroup {
        &self.bind_group
    }

    /// Freezes the atlas so it stops generating new glyphs.
    ///
    /// While frozen, [`retrieve_or_generate_glyph`](Self::retrieve_or_generate_glyph) serves
    /// glyphs already in the atlas but never generates, allocates, evicts, or grows: a request
    /// for an uncached glyph returns [`CachedGlyph::EmptyGlyph`](super::CachedGlyph::EmptyGlyph)
    /// instead. Useful together with a baked atlas (see [`from_baked`](Self::from_baked)) to
    /// guarantee no runtime glyph generation.
    pub fn freeze(&mut self) {
        self.frozen = true;
    }

    /// Unfreezes the atlas, re-enabling runtime glyph generation. Reverses [`freeze`](Self::freeze).
    pub fn unfreeze(&mut self) {
        self.frozen = false;
    }

    /// Returns whether the atlas is currently frozen. See [`freeze`](Self::freeze).
    pub fn is_frozen(&self) -> bool {
        self.frozen
    }

    /// Render all glyphs queued via `retrieve_or_generate_glyph` since the last call into
    /// their atlas regions. Must be called once per frame after encoding all text and
    /// before any pass that samples the atlas.
    pub fn write_glyphs(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        encoder: &mut wgpu::CommandEncoder,
    ) {
        profiling::scope!("TextureAtlas::write_glyphs");
        self.gpu_write
            .write_glyphs(device, queue, encoder, &self.texture);
    }
}

fn atlas_size_to_extent(
    size: AtlasSize,
    device_max_size: u32,
    device_max_layers: u32,
) -> wgpu::Extent3d {
    match size {
        AtlasSize::SingleLayer { width, height } => wgpu::Extent3d {
            width: width.clamp(super::MIN_TEXTURE_SIZE, device_max_size),
            height: height.clamp(super::MIN_TEXTURE_SIZE, device_max_size),
            depth_or_array_layers: 1,
        },
        AtlasSize::MaxSize { layer_count } => wgpu::Extent3d {
            width: device_max_size,
            height: device_max_size,
            depth_or_array_layers: layer_count.clamp(1, device_max_layers),
        },
    }
}

/// Texture usage flags for the atlas texture.
///
/// `COPY_SRC` is required to copy the texture when growing, and to read it back in `bake()`.
/// It is enabled unconditionally when the `serde` feature is on so baking works even for
/// non-growable atlases.
pub(super) fn texture_usage(allow_grow: bool) -> wgpu::TextureUsages {
    let mut usage = wgpu::TextureUsages::COPY_DST
        | wgpu::TextureUsages::TEXTURE_BINDING
        | wgpu::TextureUsages::RENDER_ATTACHMENT;
    if allow_grow || cfg!(feature = "serde") {
        usage |= wgpu::TextureUsages::COPY_SRC;
    }
    usage
}

pub(super) fn create_texture(
    device: &wgpu::Device,
    size: wgpu::Extent3d,
    usage: wgpu::TextureUsages,
) -> (wgpu::Texture, wgpu::Sampler) {
    let texture = device.create_texture(&wgpu::TextureDescriptor {
        label: Some("Klyff atlas"),
        size,
        mip_level_count: 1,
        sample_count: 1,
        dimension: wgpu::TextureDimension::D2,
        format: wgpu::TextureFormat::Rgba8Unorm,
        usage,
        view_formats: &[wgpu::TextureFormat::Rgba8Unorm],
    });

    let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
        label: Some("Klyff atlas sampler"),
        address_mode_u: wgpu::AddressMode::ClampToEdge,
        address_mode_v: wgpu::AddressMode::ClampToEdge,
        address_mode_w: wgpu::AddressMode::ClampToEdge,
        mag_filter: wgpu::FilterMode::Linear,
        min_filter: wgpu::FilterMode::Linear,
        ..Default::default()
    });
    (texture, sampler)
}