spottedcat 1.0.0

use crate::graphics::texture::TextureUploadRegion;
use crate::{Image, Pt};
use std::sync::Arc;

/// Internal structure porting the binary tree packing algorithm from packing.go
#[derive(Default, Clone)]
struct Node {
    left: Option<Box<Node>>,
    right: Option<Box<Node>>,
    x: i32,
    y: i32,
    w: i32,
    h: i32,
    split: bool,
    is_end: bool,
}

#[cfg(test)]
mod tests {
    use super::DynamicAtlas;
    use crate::{Context, Pt};

    #[test]
    fn atlas_append_on_ready_texture_queues_partial_upload_without_invalidating_generation() {
        let mut ctx = Context::new();
        let mut atlas = DynamicAtlas::new(256);
        let rgba = [255, 0, 0, 255];

        let first = atlas
            .add_region(
                &mut ctx.registry,
                1.0,
                Pt::from(1.0),
                Pt::from(1.0),
                1,
                1,
                &rgba,
            )
            .expect("first atlas insert should succeed");
        let texture_id = first.texture_id();
        let generation = ctx.registry.gpu_generation;
        let entry = ctx.registry.textures[texture_id as usize]
            .as_mut()
            .expect("atlas texture should be registered");
        entry.pending_uploads.clear();
        entry.runtime.generation = generation;

        atlas
            .add_region(
                &mut ctx.registry,
                1.0,
                Pt::from(1.0),
                Pt::from(1.0),
                1,
                1,
                &rgba,
            )
            .expect("second atlas insert should succeed");

        let entry = ctx.registry.textures[texture_id as usize]
            .as_ref()
            .expect("atlas texture should still be registered");
        assert_eq!(entry.runtime.generation, generation);
        assert_eq!(entry.pending_uploads.len(), 1);
        assert_eq!(entry.pending_uploads[0].width, 3);
        assert_eq!(entry.pending_uploads[0].height, 3);
    }
}

impl Node {
    fn insert_node(&mut self, rect_w: i32, rect_h: i32) -> Option<(i32, i32)> {
        // If we are already split, recurse into children
        if self.split {
            if let Some(left) = self.left.as_mut() {
                if let Some(pos) = left.insert_node(rect_w, rect_h) {
                    return Some(pos);
                }
            }
            if let Some(right) = self.right.as_mut() {
                return right.insert_node(rect_w, rect_h);
            }
            return None;
        }

        // If this node is used or too small, fail
        if self.is_end || self.w < rect_w || self.h < rect_h {
            return None;
        }

        // If perfect fit, use it
        if self.w == rect_w && self.h == rect_h {
            self.is_end = true;
            return Some((self.x, self.y));
        }

        // Otherwise split this node
        self.split = true;

        let dw = self.w - rect_w;
        let dh = self.h - rect_h;

        if dw > dh {
            // Split horizontally (left and right columns)
            self.left = Some(Box::new(Node {
                x: self.x,
                y: self.y,
                w: rect_w,
                h: self.h,
                ..Default::default()
            }));
            self.right = Some(Box::new(Node {
                x: self.x + rect_w,
                y: self.y,
                w: dw,
                h: self.h,
                ..Default::default()
            }));
        } else {
            // Split vertically (top and bottom rows)
            self.left = Some(Box::new(Node {
                x: self.x,
                y: self.y,
                w: self.w,
                h: rect_h,
                ..Default::default()
            }));
            self.right = Some(Box::new(Node {
                x: self.x,
                y: self.y + rect_h,
                w: self.w,
                h: dh,
                ..Default::default()
            }));
        }

        // Always put the item into the newly created left node
        self.left.as_mut().unwrap().insert_node(rect_w, rect_h)
    }
}

pub(crate) struct Packer {
    root: Node,
}

impl Packer {
    pub fn new(w: i32, h: i32) -> Self {
        Self {
            root: Node {
                x: 0,
                y: 0,
                w,
                h,
                ..Node::default()
            },
        }
    }

    pub fn insert(&mut self, w: i32, h: i32) -> Option<(i32, i32)> {
        self.root.insert_node(w, h)
    }
}

pub(crate) struct AtlasPage {
    pub texture_id: u32,
    pub packer: Packer,
    pub buffer: Vec<u8>,
    pub pixel_width: u32,
    pub pixel_height: u32,
}

/// A dynamic texture atlas that can grow across multiple pages.
pub(crate) struct DynamicAtlas {
    pub pages: Vec<AtlasPage>,
    pub max_dim: u32,
}

impl DynamicAtlas {
    pub fn new(max_dim: u32) -> Self {
        Self {
            pages: Vec::new(),
            max_dim,
        }
    }

    pub fn add_region(
        &mut self,
        registry: &mut crate::context::ResourceRegistry,
        scale_factor: f64,
        logical_w: Pt,
        logical_h: Pt,
        w_px: u32,
        h_px: u32,
        rgba: &[u8],
    ) -> anyhow::Result<Image> {
        let padding = 2;
        let total_w = w_px + padding;
        let total_h = h_px + padding;

        if total_w > self.max_dim || total_h > self.max_dim {
            anyhow::bail!("Region too large for atlas: {}x{}", w_px, h_px);
        }

        // Try to fit in existing pages
        for page_idx in 0..self.pages.len() {
            let page = &mut self.pages[page_idx];
            if page.pixel_width >= total_w && page.pixel_height >= total_h {
                if let Some((x, y)) = page.packer.insert(total_w as i32, total_h as i32) {
                    return self.write_to_page(
                        registry,
                        scale_factor,
                        page_idx,
                        x as u32,
                        y as u32,
                        logical_w,
                        logical_h,
                        w_px,
                        h_px,
                        rgba,
                    );
                }
            }
        }

        // If no page works, create a new one
        let new_w = (total_w.next_power_of_two()).max(256).min(self.max_dim);
        let new_h = (total_h.next_power_of_two()).max(256).min(self.max_dim);

        let page_idx = self.create_page(registry, scale_factor, new_w, new_h);
        let page = &mut self.pages[page_idx];
        if let Some((x, y)) = page.packer.insert(total_w as i32, total_h as i32) {
            self.write_to_page(
                registry,
                scale_factor,
                page_idx,
                x as u32,
                y as u32,
                logical_w,
                logical_h,
                w_px,
                h_px,
                rgba,
            )
        } else {
            anyhow::bail!("Failed to insert region into new atlas page");
        }
    }

    fn create_page(
        &mut self,
        registry: &mut crate::context::ResourceRegistry,
        scale_factor: f64,
        w_px: u32,
        h_px: u32,
    ) -> usize {
        let buffer = vec![0u8; (w_px * h_px * 4) as usize];
        let scale_factor = scale_factor.max(1.0);
        let logical_w = Pt::from_physical_px(w_px as f64, scale_factor);
        let logical_h = Pt::from_physical_px(h_px as f64, scale_factor);

        // We need to register the texture manually in the registry
        let texture_id = registry.next_texture_id;
        registry.next_texture_id += 1;
        let image_id = registry.next_image_id;
        registry.next_image_id += 1;

        while registry.textures.len() <= texture_id as usize {
            registry.textures.push(None);
        }
        registry.textures[texture_id as usize] =
            Some(crate::graphics::texture::TextureEntry::new_dynamic_atlas(
                logical_w,
                logical_h,
                w_px,
                h_px,
                image_id,
                std::sync::Arc::from(buffer.as_slice()),
            ));

        let bounds = crate::image::Bounds::new(Pt(0.0), Pt(0.0), logical_w, logical_h);
        while registry.images.len() <= image_id as usize {
            registry.images.push(None);
        }
        registry.images[image_id as usize] = Some(crate::image::ImageEntry::new(
            texture_id,
            bounds,
            crate::image::PixelBounds {
                x: 0,
                y: 0,
                width: w_px,
                height: h_px,
            },
        ));

        registry.dirty_assets = true;

        let page = AtlasPage {
            texture_id,
            packer: Packer::new(w_px as i32, h_px as i32),
            buffer,
            pixel_width: w_px,
            pixel_height: h_px,
        };
        self.pages.push(page);
        self.pages.len() - 1
    }

    fn write_to_page(
        &mut self,
        registry: &mut crate::context::ResourceRegistry,
        scale_factor: f64,
        page_idx: usize,
        x: u32,
        y: u32,
        logical_w: Pt,
        logical_h: Pt,
        w: u32,
        h: u32,
        rgba: &[u8],
    ) -> anyhow::Result<Image> {
        let page = &mut self.pages[page_idx];
        let inner_x = x + 1;
        let inner_y = y + 1;

        for row in 0..h {
            let src_idx = (row * w * 4) as usize;
            let dst_idx = ((inner_y + row) * page.pixel_width + inner_x) as usize * 4;
            page.buffer[dst_idx..dst_idx + (w * 4) as usize]
                .copy_from_slice(&rgba[src_idx..src_idx + (w * 4) as usize]);
        }

        for row in 0..h {
            let src_idx = (row * w * 4) as usize;
            let dst_y = inner_y + row;
            let left_dst_idx = (dst_y * page.pixel_width + inner_x - 1) as usize * 4;
            let right_dst_idx = (dst_y * page.pixel_width + inner_x + w) as usize * 4;
            page.buffer[left_dst_idx..left_dst_idx + 4]
                .copy_from_slice(&rgba[src_idx..src_idx + 4]);
            let last_pixel = src_idx + ((w - 1) * 4) as usize;
            page.buffer[right_dst_idx..right_dst_idx + 4]
                .copy_from_slice(&rgba[last_pixel..last_pixel + 4]);
        }

        let top_dst_idx = ((inner_y - 1) * page.pixel_width + inner_x) as usize * 4;
        page.buffer[top_dst_idx..top_dst_idx + (w * 4) as usize]
            .copy_from_slice(&rgba[..(w * 4) as usize]);

        let bottom_src_idx = ((h - 1) * w * 4) as usize;
        let bottom_dst_idx = ((inner_y + h) * page.pixel_width + inner_x) as usize * 4;
        page.buffer[bottom_dst_idx..bottom_dst_idx + (w * 4) as usize]
            .copy_from_slice(&rgba[bottom_src_idx..bottom_src_idx + (w * 4) as usize]);

        let top_left = &rgba[..4];
        let top_right_idx = ((w - 1) * 4) as usize;
        let top_right = &rgba[top_right_idx..top_right_idx + 4];
        let bottom_left = &rgba[bottom_src_idx..bottom_src_idx + 4];
        let bottom_right_idx = bottom_src_idx + ((w - 1) * 4) as usize;
        let bottom_right = &rgba[bottom_right_idx..bottom_right_idx + 4];

        let top_left_dst_idx = ((inner_y - 1) * page.pixel_width + inner_x - 1) as usize * 4;
        page.buffer[top_left_dst_idx..top_left_dst_idx + 4].copy_from_slice(top_left);
        let top_right_dst_idx = ((inner_y - 1) * page.pixel_width + inner_x + w) as usize * 4;
        page.buffer[top_right_dst_idx..top_right_dst_idx + 4].copy_from_slice(top_right);
        let bottom_left_dst_idx = ((inner_y + h) * page.pixel_width + inner_x - 1) as usize * 4;
        page.buffer[bottom_left_dst_idx..bottom_left_dst_idx + 4].copy_from_slice(bottom_left);
        let bottom_right_dst_idx = ((inner_y + h) * page.pixel_width + inner_x + w) as usize * 4;
        page.buffer[bottom_right_dst_idx..bottom_right_dst_idx + 4].copy_from_slice(bottom_right);

        // Update the texture data in Registry
        let registry_generation = registry.gpu_generation;
        if let Some(entry) = registry
            .textures
            .get_mut(page.texture_id as usize)
            .and_then(|v| v.as_mut())
        {
            let upload_x = inner_x - 1;
            let upload_y = inner_y - 1;
            let upload_w = w + 2;
            let upload_h = h + 2;
            let mut rgba_region = Vec::with_capacity((upload_w * upload_h * 4) as usize);
            for row in 0..upload_h {
                let src_idx = ((upload_y + row) * page.pixel_width + upload_x) as usize * 4;
                let row_len = (upload_w * 4) as usize;
                rgba_region.extend_from_slice(&page.buffer[src_idx..src_idx + row_len]);
            }

            if !entry.is_ready(registry_generation) {
                entry.raw_data = Some(Arc::from(page.buffer.as_slice()));
            }
            entry.pending_uploads.push(TextureUploadRegion {
                x: upload_x,
                y: upload_y,
                width: upload_w,
                height: upload_h,
                rgba: rgba_region,
            });
            registry.dirty_assets = true;
        }

        let scale_factor = scale_factor.max(1.0);
        let logical_x = Pt::from_physical_px(inner_x as f64, scale_factor);
        let logical_y = Pt::from_physical_px(inner_y as f64, scale_factor);

        // Manual sub-image registration
        let view_id = registry.next_image_id;
        registry.next_image_id += 1;

        let entry = crate::image::ImageEntry::new(
            page.texture_id,
            crate::image::Bounds::new(logical_x, logical_y, logical_w, logical_h),
            crate::image::PixelBounds {
                x: inner_x,
                y: inner_y,
                width: w,
                height: h,
            },
        );

        while registry.images.len() <= view_id as usize {
            registry.images.push(None);
        }
        registry.images[view_id as usize] = Some(entry);
        registry.dirty_assets = true;

        Ok(Image {
            id: view_id,
            texture_id: page.texture_id,
            x: logical_x,
            y: logical_y,
            width: logical_w,
            height: logical_h,
            pixel_bounds: crate::image::PixelBounds {
                x: inner_x,
                y: inner_y,
                width: w,
                height: h,
            },
        })
    }

    pub(crate) fn sync_raw_data(&self, registry: &mut crate::context::ResourceRegistry) {
        for page in &self.pages {
            if let Some(entry) = registry
                .textures
                .get_mut(page.texture_id as usize)
                .and_then(|v| v.as_mut())
            {
                entry.raw_data = Some(Arc::from(page.buffer.as_slice()));
            }
        }
    }
}