roast2d_internal 0.3.4

use std::cell::RefCell;

use crate::prelude::*;
use glam::Vec3Swizzles;

use super::layout::{LayoutBuilder, LayoutResult};

#[derive(Clone, Copy, PartialEq)]
pub enum LayoutDirection {
    Vertical,
    Horizontal,
}

pub enum UiItem {
    Text { text: Text, size: Vec2 },
    UiBox { ui_box: UiBox },
    Sprite { sprite: Sprite, size: Vec2 },
}

impl UiItem {
    pub fn size(&self) -> Vec2 {
        match self {
            UiItem::Text { size, .. } => *size,
            UiItem::UiBox { ui_box } => ui_box.box_size(),
            UiItem::Sprite { size, .. } => *size,
        }
    }

    pub fn with_text(g: &mut Engine, text: Text) -> Result<Self> {
        let (_, size) = g.render_text(&text)?;
        Ok(UiItem::Text {
            text,
            size: size.as_vec2(),
        })
    }

    pub fn with_ui_box(ui_box: UiBox) -> Self {
        UiItem::UiBox { ui_box }
    }

    pub fn with_sprite(sprite: Sprite, size: Vec2) -> Self {
        UiItem::Sprite { sprite, size }
    }
}

struct LayoutCache {
    box_pos: Vec3,
    layout_result: LayoutResult,
}

pub struct UiBox {
    pub items: Vec<UiItem>,
    pub width: f32,
    pub height: f32,
    pub max_width: f32,
    pub padding: Vec2,
    pub gap: f32,
    pub debug: bool,
    pub direction: LayoutDirection,
    layout_cache: RefCell<Option<LayoutCache>>,
}

impl UiBox {
    pub fn new(max_width: f32) -> Self {
        Self {
            items: Vec::new(),
            width: 0.0,
            height: 0.0,
            max_width,
            padding: Vec2::ZERO,
            gap: 0.0,
            debug: false,
            direction: LayoutDirection::Vertical,
            layout_cache: RefCell::new(None),
        }
    }

    pub fn with_gap(mut self, gap: f32) -> Self {
        self.gap = gap;
        self.layout_cache.replace(None);
        self
    }

    pub fn with_direction(mut self, direction: LayoutDirection) -> Self {
        self.direction = direction;
        self.layout_cache.replace(None);
        self
    }

    pub fn push(&mut self, item: UiItem) {
        let size = item.size();
        match self.direction {
            LayoutDirection::Vertical => {
                if size.x > self.width {
                    self.width = size.x.min(self.max_width);
                }
                self.height += size.y;
                if !self.items.is_empty() {
                    self.height += self.gap;
                }
            }
            LayoutDirection::Horizontal => {
                if size.y > self.height {
                    self.height = size.y;
                }
                self.width += size.x;
                if !self.items.is_empty() {
                    self.width += self.gap;
                }
            }
        }
        self.items.push(item);
        self.layout_cache.replace(None);
    }

    pub fn box_size(&self) -> Vec2 {
        Vec2::new(self.width, self.height) + self.padding * 2.0
    }

    fn calculate_layout(&self, box_pos: Vec3) -> LayoutResult {
        // Use LayoutBuilder to calculate item positions
        let bounds_center = box_pos.xy();
        let bounds_size = self.box_size();
        let mut layout_builder = LayoutBuilder::new(Rect {
            min: bounds_center - bounds_size * 0.5,
            max: bounds_center + bounds_size * 0.5,
        });

        layout_builder = match self.direction {
            LayoutDirection::Vertical => layout_builder.vertical(self.gap),
            LayoutDirection::Horizontal => layout_builder.horizontal(self.gap),
        };

        for item in &self.items {
            layout_builder = layout_builder.add_item(item.size());
        }

        layout_builder.calculate()
    }

    fn fetch_layout(&self, box_pos: Vec3) -> std::cell::Ref<'_, Option<LayoutCache>> {
        // check box_pos and clear cache if it's not the same
        let has_cache = self
            .layout_cache
            .borrow()
            .as_ref()
            .is_some_and(|cache| cache.box_pos == box_pos);
        if !has_cache {
            self.layout_cache.replace(Some(LayoutCache {
                box_pos,
                layout_result: self.calculate_layout(box_pos),
            }));
        }

        self.layout_cache.borrow()
    }

    pub fn draw(&self, g: &mut Engine, box_pos: Vec3) -> Result<()> {
        if self.debug {
            g.draw_rect(
                BLUE.with_a(0.3),
                self.box_size(),
                None,
                Transform::new(box_pos),
            );
        }
        let z_index = box_pos.z;

        let cache = self.fetch_layout(box_pos);
        let layout_result = &cache.as_ref().unwrap().layout_result;

        // Draw each item at its calculated position
        for (item, layout_item) in self.items.iter().zip(layout_result.items.iter()) {
            let item_pos = layout_item.position;
            match item {
                UiItem::Text { text, size } => {
                    if self.debug {
                        g.draw_rect(
                            RED.with_a(0.3),
                            *size,
                            None,
                            Transform::new(item_pos.extend(z_index)),
                        );
                    }

                    g.draw_text(text, None, Transform::new(item_pos.extend(z_index)));
                }
                UiItem::UiBox { ui_box } => {
                    ui_box.draw(g, item_pos.extend(z_index))?;
                }
                UiItem::Sprite { sprite, size } => {
                    if self.debug {
                        g.draw_rect(
                            GREEN.with_a(0.3),
                            *size,
                            None,
                            Transform::new(item_pos.extend(z_index)),
                        );
                    }

                    let mut sp = sprite.clone();
                    sp.size = (*size).as_uvec2();
                    g.draw(&sp, Transform::new(item_pos.extend(z_index)));
                }
            }
        }

        Ok(())
    }
}