//! Block layout engine: box generation, block/inline layout, and text wrapping.
//!
//! Builds a layout tree from a styled tree (DOM + computed styles) and positions
//! block-level elements vertically with proper inline formatting context.

use we_css::values::Color;
use we_dom::{Document, NodeData, NodeId};
use we_style::computed::{
    BorderStyle, ComputedStyle, Display, LengthOrAuto, StyledNode, TextAlign, TextDecoration,
};
use we_text::font::Font;

/// Edge sizes for box model (margin, padding, border).
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub struct EdgeSizes {
    pub top: f32,
    pub right: f32,
    pub bottom: f32,
    pub left: f32,
}

/// A positioned rectangle with content area dimensions.
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub struct Rect {
    pub x: f32,
    pub y: f32,
    pub width: f32,
    pub height: f32,
}

/// The type of layout box.
#[derive(Debug)]
pub enum BoxType {
    /// Block-level box from an element.
    Block(NodeId),
    /// Inline-level box from an element.
    Inline(NodeId),
    /// A run of text from a text node.
    TextRun { node: NodeId, text: String },
    /// Anonymous block wrapping inline content within a block container.
    Anonymous,
}

/// A single positioned text fragment with its own styling.
///
/// Multiple fragments can share the same y-coordinate when they are
/// on the same visual line (e.g. `<p>Hello <em>world</em></p>` produces
/// two fragments at the same y).
#[derive(Debug, Clone, PartialEq)]
pub struct TextLine {
    pub text: String,
    pub x: f32,
    pub y: f32,
    pub width: f32,
    pub font_size: f32,
    pub color: Color,
    pub text_decoration: TextDecoration,
    pub background_color: Color,
}

/// A box in the layout tree with dimensions and child boxes.
#[derive(Debug)]
pub struct LayoutBox {
    pub box_type: BoxType,
    pub rect: Rect,
    pub margin: EdgeSizes,
    pub padding: EdgeSizes,
    pub border: EdgeSizes,
    pub children: Vec<LayoutBox>,
    pub font_size: f32,
    /// Positioned text fragments (populated for boxes with inline content).
    pub lines: Vec<TextLine>,
    /// Text color.
    pub color: Color,
    /// Background color.
    pub background_color: Color,
    /// Text decoration (underline, etc.).
    pub text_decoration: TextDecoration,
    /// Border styles (top, right, bottom, left).
    pub border_styles: [BorderStyle; 4],
    /// Border colors (top, right, bottom, left).
    pub border_colors: [Color; 4],
    /// Text alignment for this box's inline content.
    pub text_align: TextAlign,
    /// Computed line height in px.
    pub line_height: f32,
}

impl LayoutBox {
    fn new(box_type: BoxType, style: &ComputedStyle) -> Self {
        LayoutBox {
            box_type,
            rect: Rect::default(),
            margin: EdgeSizes::default(),
            padding: EdgeSizes::default(),
            border: EdgeSizes::default(),
            children: Vec::new(),
            font_size: style.font_size,
            lines: Vec::new(),
            color: style.color,
            background_color: style.background_color,
            text_decoration: style.text_decoration,
            border_styles: [
                style.border_top_style,
                style.border_right_style,
                style.border_bottom_style,
                style.border_left_style,
            ],
            border_colors: [
                style.border_top_color,
                style.border_right_color,
                style.border_bottom_color,
                style.border_left_color,
            ],
            text_align: style.text_align,
            line_height: style.line_height,
        }
    }

    /// Total height including margin, border, and padding.
    pub fn margin_box_height(&self) -> f32 {
        self.margin.top
            + self.border.top
            + self.padding.top
            + self.rect.height
            + self.padding.bottom
            + self.border.bottom
            + self.margin.bottom
    }

    /// Iterate over all boxes in depth-first pre-order.
    pub fn iter(&self) -> LayoutBoxIter<'_> {
        LayoutBoxIter { stack: vec![self] }
    }
}

/// Depth-first pre-order iterator over layout boxes.
pub struct LayoutBoxIter<'a> {
    stack: Vec<&'a LayoutBox>,
}

impl<'a> Iterator for LayoutBoxIter<'a> {
    type Item = &'a LayoutBox;

    fn next(&mut self) -> Option<&'a LayoutBox> {
        let node = self.stack.pop()?;
        for child in node.children.iter().rev() {
            self.stack.push(child);
        }
        Some(node)
    }
}

/// The result of laying out a document.
#[derive(Debug)]
pub struct LayoutTree {
    pub root: LayoutBox,
    pub width: f32,
    pub height: f32,
}

impl LayoutTree {
    /// Iterate over all layout boxes in depth-first pre-order.
    pub fn iter(&self) -> LayoutBoxIter<'_> {
        self.root.iter()
    }
}

// ---------------------------------------------------------------------------
// Resolve LengthOrAuto to f32
// ---------------------------------------------------------------------------

fn resolve_length(value: LengthOrAuto) -> f32 {
    match value {
        LengthOrAuto::Length(px) => px,
        LengthOrAuto::Auto => 0.0,
    }
}

// ---------------------------------------------------------------------------
// Build layout tree from styled tree
// ---------------------------------------------------------------------------

fn build_box(styled: &StyledNode, doc: &Document) -> Option<LayoutBox> {
    let node = styled.node;
    let style = &styled.style;

    match doc.node_data(node) {
        NodeData::Document => {
            let mut children = Vec::new();
            for child in &styled.children {
                if let Some(child_box) = build_box(child, doc) {
                    children.push(child_box);
                }
            }
            if children.len() == 1 {
                children.into_iter().next()
            } else if children.is_empty() {
                None
            } else {
                let mut b = LayoutBox::new(BoxType::Anonymous, style);
                b.children = children;
                Some(b)
            }
        }
        NodeData::Element { .. } => {
            if style.display == Display::None {
                return None;
            }

            let margin = EdgeSizes {
                top: resolve_length(style.margin_top),
                right: resolve_length(style.margin_right),
                bottom: resolve_length(style.margin_bottom),
                left: resolve_length(style.margin_left),
            };
            let padding = EdgeSizes {
                top: style.padding_top,
                right: style.padding_right,
                bottom: style.padding_bottom,
                left: style.padding_left,
            };
            let border = EdgeSizes {
                top: if style.border_top_style != BorderStyle::None {
                    style.border_top_width
                } else {
                    0.0
                },
                right: if style.border_right_style != BorderStyle::None {
                    style.border_right_width
                } else {
                    0.0
                },
                bottom: if style.border_bottom_style != BorderStyle::None {
                    style.border_bottom_width
                } else {
                    0.0
                },
                left: if style.border_left_style != BorderStyle::None {
                    style.border_left_width
                } else {
                    0.0
                },
            };

            let mut children = Vec::new();
            for child in &styled.children {
                if let Some(child_box) = build_box(child, doc) {
                    children.push(child_box);
                }
            }

            let box_type = match style.display {
                Display::Block => BoxType::Block(node),
                Display::Inline => BoxType::Inline(node),
                Display::None => unreachable!(),
            };

            if style.display == Display::Block {
                children = normalize_children(children, style);
            }

            let mut b = LayoutBox::new(box_type, style);
            b.margin = margin;
            b.padding = padding;
            b.border = border;
            b.children = children;
            Some(b)
        }
        NodeData::Text { data } => {
            let collapsed = collapse_whitespace(data);
            if collapsed.is_empty() {
                return None;
            }
            Some(LayoutBox::new(
                BoxType::TextRun {
                    node,
                    text: collapsed,
                },
                style,
            ))
        }
        NodeData::Comment { .. } => None,
    }
}

/// Collapse runs of whitespace to a single space.
fn collapse_whitespace(s: &str) -> String {
    let mut result = String::new();
    let mut in_ws = false;
    for ch in s.chars() {
        if ch.is_whitespace() {
            if !in_ws {
                result.push(' ');
            }
            in_ws = true;
        } else {
            in_ws = false;
            result.push(ch);
        }
    }
    result
}

/// If a block container has a mix of block-level and inline-level children,
/// wrap consecutive inline runs in anonymous block boxes.
fn normalize_children(children: Vec<LayoutBox>, parent_style: &ComputedStyle) -> Vec<LayoutBox> {
    if children.is_empty() {
        return children;
    }

    let has_block = children.iter().any(is_block_level);
    if !has_block {
        return children;
    }

    let has_inline = children.iter().any(|c| !is_block_level(c));
    if !has_inline {
        return children;
    }

    let mut result = Vec::new();
    let mut inline_group: Vec<LayoutBox> = Vec::new();

    for child in children {
        if is_block_level(&child) {
            if !inline_group.is_empty() {
                let mut anon = LayoutBox::new(BoxType::Anonymous, parent_style);
                anon.children = std::mem::take(&mut inline_group);
                result.push(anon);
            }
            result.push(child);
        } else {
            inline_group.push(child);
        }
    }

    if !inline_group.is_empty() {
        let mut anon = LayoutBox::new(BoxType::Anonymous, parent_style);
        anon.children = inline_group;
        result.push(anon);
    }

    result
}

fn is_block_level(b: &LayoutBox) -> bool {
    matches!(b.box_type, BoxType::Block(_) | BoxType::Anonymous)
}

// ---------------------------------------------------------------------------
// Layout algorithm
// ---------------------------------------------------------------------------

/// Position and size a layout box within `available_width` at position (`x`, `y`).
fn compute_layout(
    b: &mut LayoutBox,
    x: f32,
    y: f32,
    available_width: f32,
    font: &Font,
    doc: &Document,
) {
    let content_x = x + b.margin.left + b.border.left + b.padding.left;
    let content_y = y + b.margin.top + b.border.top + b.padding.top;
    let content_width = (available_width
        - b.margin.left
        - b.margin.right
        - b.border.left
        - b.border.right
        - b.padding.left
        - b.padding.right)
        .max(0.0);

    b.rect.x = content_x;
    b.rect.y = content_y;
    b.rect.width = content_width;

    match &b.box_type {
        BoxType::Block(_) | BoxType::Anonymous => {
            if has_block_children(b) {
                layout_block_children(b, font, doc);
            } else {
                layout_inline_children(b, font, doc);
            }
        }
        BoxType::TextRun { .. } | BoxType::Inline(_) => {
            // Handled by the parent's inline layout.
        }
    }
}

fn has_block_children(b: &LayoutBox) -> bool {
    b.children.iter().any(is_block_level)
}

/// Lay out block-level children: stack them vertically.
fn layout_block_children(parent: &mut LayoutBox, font: &Font, doc: &Document) {
    let content_x = parent.rect.x;
    let content_width = parent.rect.width;
    let mut cursor_y = parent.rect.y;

    for child in &mut parent.children {
        compute_layout(child, content_x, cursor_y, content_width, font, doc);
        cursor_y += child.margin_box_height();
    }

    parent.rect.height = cursor_y - parent.rect.y;
}

// ---------------------------------------------------------------------------
// Inline formatting context
// ---------------------------------------------------------------------------

/// An inline item produced by flattening the inline tree.
enum InlineItemKind {
    /// A word of text with associated styling.
    Word {
        text: String,
        font_size: f32,
        color: Color,
        text_decoration: TextDecoration,
        background_color: Color,
    },
    /// Whitespace between words.
    Space { font_size: f32 },
    /// Forced line break (`<br>`).
    ForcedBreak,
    /// Start of an inline box (for margin/padding/border tracking).
    InlineStart {
        margin_left: f32,
        padding_left: f32,
        border_left: f32,
    },
    /// End of an inline box.
    InlineEnd {
        margin_right: f32,
        padding_right: f32,
        border_right: f32,
    },
}

/// A pending fragment on the current line.
struct PendingFragment {
    text: String,
    x: f32,
    width: f32,
    font_size: f32,
    color: Color,
    text_decoration: TextDecoration,
    background_color: Color,
}

/// Flatten the inline children tree into a sequence of items.
fn flatten_inline_tree(children: &[LayoutBox], doc: &Document, items: &mut Vec<InlineItemKind>) {
    for child in children {
        match &child.box_type {
            BoxType::TextRun { text, .. } => {
                let words = split_into_words(text);
                for segment in words {
                    match segment {
                        WordSegment::Word(w) => {
                            items.push(InlineItemKind::Word {
                                text: w,
                                font_size: child.font_size,
                                color: child.color,
                                text_decoration: child.text_decoration,
                                background_color: child.background_color,
                            });
                        }
                        WordSegment::Space => {
                            items.push(InlineItemKind::Space {
                                font_size: child.font_size,
                            });
                        }
                    }
                }
            }
            BoxType::Inline(node_id) => {
                if let NodeData::Element { tag_name, .. } = doc.node_data(*node_id) {
                    if tag_name == "br" {
                        items.push(InlineItemKind::ForcedBreak);
                        continue;
                    }
                }

                items.push(InlineItemKind::InlineStart {
                    margin_left: child.margin.left,
                    padding_left: child.padding.left,
                    border_left: child.border.left,
                });

                flatten_inline_tree(&child.children, doc, items);

                items.push(InlineItemKind::InlineEnd {
                    margin_right: child.margin.right,
                    padding_right: child.padding.right,
                    border_right: child.border.right,
                });
            }
            _ => {}
        }
    }
}

enum WordSegment {
    Word(String),
    Space,
}

/// Split text into alternating words and spaces.
fn split_into_words(text: &str) -> Vec<WordSegment> {
    let mut segments = Vec::new();
    let mut current_word = String::new();

    for ch in text.chars() {
        if ch == ' ' {
            if !current_word.is_empty() {
                segments.push(WordSegment::Word(std::mem::take(&mut current_word)));
            }
            segments.push(WordSegment::Space);
        } else {
            current_word.push(ch);
        }
    }

    if !current_word.is_empty() {
        segments.push(WordSegment::Word(current_word));
    }

    segments
}

/// Lay out inline children using a proper inline formatting context.
fn layout_inline_children(parent: &mut LayoutBox, font: &Font, doc: &Document) {
    let available_width = parent.rect.width;
    let text_align = parent.text_align;
    let line_height = parent.line_height;

    let mut items = Vec::new();
    flatten_inline_tree(&parent.children, doc, &mut items);

    if items.is_empty() {
        parent.rect.height = 0.0;
        return;
    }

    // Process items into line boxes.
    let mut all_lines: Vec<Vec<PendingFragment>> = Vec::new();
    let mut current_line: Vec<PendingFragment> = Vec::new();
    let mut cursor_x: f32 = 0.0;

    for item in &items {
        match item {
            InlineItemKind::Word {
                text,
                font_size,
                color,
                text_decoration,
                background_color,
            } => {
                let word_width = measure_text_width(font, text, *font_size);

                // If this word doesn't fit and the line isn't empty, break.
                if cursor_x > 0.0 && cursor_x + word_width > available_width {
                    all_lines.push(std::mem::take(&mut current_line));
                    cursor_x = 0.0;
                }

                current_line.push(PendingFragment {
                    text: text.clone(),
                    x: cursor_x,
                    width: word_width,
                    font_size: *font_size,
                    color: *color,
                    text_decoration: *text_decoration,
                    background_color: *background_color,
                });
                cursor_x += word_width;
            }
            InlineItemKind::Space { font_size } => {
                // Only add space if we have content on the line.
                if !current_line.is_empty() {
                    let space_width = measure_text_width(font, " ", *font_size);
                    if cursor_x + space_width <= available_width {
                        cursor_x += space_width;
                    }
                }
            }
            InlineItemKind::ForcedBreak => {
                all_lines.push(std::mem::take(&mut current_line));
                cursor_x = 0.0;
            }
            InlineItemKind::InlineStart {
                margin_left,
                padding_left,
                border_left,
            } => {
                cursor_x += margin_left + padding_left + border_left;
            }
            InlineItemKind::InlineEnd {
                margin_right,
                padding_right,
                border_right,
            } => {
                cursor_x += margin_right + padding_right + border_right;
            }
        }
    }

    // Flush the last line.
    if !current_line.is_empty() {
        all_lines.push(current_line);
    }

    if all_lines.is_empty() {
        parent.rect.height = 0.0;
        return;
    }

    // Position lines vertically and apply text-align.
    let mut text_lines = Vec::new();
    let mut y = parent.rect.y;
    let num_lines = all_lines.len();

    for (line_idx, line_fragments) in all_lines.iter().enumerate() {
        if line_fragments.is_empty() {
            y += line_height;
            continue;
        }

        // Compute line width from last fragment.
        let line_width = match line_fragments.last() {
            Some(last) => last.x + last.width,
            None => 0.0,
        };

        // Compute text-align offset.
        let is_last_line = line_idx == num_lines - 1;
        let align_offset =
            compute_align_offset(text_align, available_width, line_width, is_last_line);

        for frag in line_fragments {
            text_lines.push(TextLine {
                text: frag.text.clone(),
                x: parent.rect.x + frag.x + align_offset,
                y,
                width: frag.width,
                font_size: frag.font_size,
                color: frag.color,
                text_decoration: frag.text_decoration,
                background_color: frag.background_color,
            });
        }

        y += line_height;
    }

    parent.rect.height = num_lines as f32 * line_height;
    parent.lines = text_lines;
}

/// Compute the horizontal offset for text alignment.
fn compute_align_offset(
    align: TextAlign,
    available_width: f32,
    line_width: f32,
    is_last_line: bool,
) -> f32 {
    let extra_space = (available_width - line_width).max(0.0);
    match align {
        TextAlign::Left => 0.0,
        TextAlign::Center => extra_space / 2.0,
        TextAlign::Right => extra_space,
        TextAlign::Justify => {
            // Don't justify the last line (CSS spec behavior).
            if is_last_line {
                0.0
            } else {
                // For justify, we shift the whole line by 0 — the actual distribution
                // of space between words would need per-word spacing. For now, treat
                // as left-aligned; full justify support is a future enhancement.
                0.0
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Text measurement
// ---------------------------------------------------------------------------

/// Measure the total advance width of a text string at the given font size.
fn measure_text_width(font: &Font, text: &str, font_size: f32) -> f32 {
    let shaped = font.shape_text(text, font_size);
    match shaped.last() {
        Some(last) => last.x_offset + last.x_advance,
        None => 0.0,
    }
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

/// Build and lay out from a styled tree (produced by `resolve_styles`).
///
/// Returns a `LayoutTree` with positioned boxes ready for rendering.
pub fn layout(
    styled_root: &StyledNode,
    doc: &Document,
    viewport_width: f32,
    _viewport_height: f32,
    font: &Font,
) -> LayoutTree {
    let mut root = match build_box(styled_root, doc) {
        Some(b) => b,
        None => {
            return LayoutTree {
                root: LayoutBox::new(BoxType::Anonymous, &ComputedStyle::default()),
                width: viewport_width,
                height: 0.0,
            };
        }
    };

    compute_layout(&mut root, 0.0, 0.0, viewport_width, font, doc);

    let height = root.margin_box_height();
    LayoutTree {
        root,
        width: viewport_width,
        height,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use we_dom::Document;
    use we_style::computed::{extract_stylesheets, resolve_styles};

    fn test_font() -> Font {
        let paths = [
            "/System/Library/Fonts/Geneva.ttf",
            "/System/Library/Fonts/Monaco.ttf",
        ];
        for path in &paths {
            let p = std::path::Path::new(path);
            if p.exists() {
                return Font::from_file(p).expect("failed to parse font");
            }
        }
        panic!("no test font found");
    }

    fn layout_doc(doc: &Document) -> LayoutTree {
        let font = test_font();
        let sheets = extract_stylesheets(doc);
        let styled = resolve_styles(doc, &sheets).unwrap();
        layout(&styled, doc, 800.0, 600.0, &font)
    }

    #[test]
    fn empty_document() {
        let doc = Document::new();
        let font = test_font();
        let sheets = extract_stylesheets(&doc);
        let styled = resolve_styles(&doc, &sheets);
        if let Some(styled) = styled {
            let tree = layout(&styled, &doc, 800.0, 600.0, &font);
            assert_eq!(tree.width, 800.0);
        }
    }

    #[test]
    fn single_paragraph() {
        let mut doc = Document::new();
        let root = doc.root();
        let html = doc.create_element("html");
        let body = doc.create_element("body");
        let p = doc.create_element("p");
        let text = doc.create_text("Hello world");
        doc.append_child(root, html);
        doc.append_child(html, body);
        doc.append_child(body, p);
        doc.append_child(p, text);

        let tree = layout_doc(&doc);

        assert!(matches!(tree.root.box_type, BoxType::Block(_)));

        let body_box = &tree.root.children[0];
        assert!(matches!(body_box.box_type, BoxType::Block(_)));

        let p_box = &body_box.children[0];
        assert!(matches!(p_box.box_type, BoxType::Block(_)));

        assert!(!p_box.lines.is_empty(), "p should have text fragments");

        // Collect all text on the first visual line.
        let first_y = p_box.lines[0].y;
        let line_text: String = p_box
            .lines
            .iter()
            .filter(|l| (l.y - first_y).abs() < 0.01)
            .map(|l| l.text.as_str())
            .collect::<Vec<_>>()
            .join(" ");
        assert!(
            line_text.contains("Hello") && line_text.contains("world"),
            "line should contain Hello and world, got: {line_text}"
        );

        assert_eq!(p_box.margin.top, 16.0);
        assert_eq!(p_box.margin.bottom, 16.0);
    }

    #[test]
    fn blocks_stack_vertically() {
        let mut doc = Document::new();
        let root = doc.root();
        let html = doc.create_element("html");
        let body = doc.create_element("body");
        let p1 = doc.create_element("p");
        let t1 = doc.create_text("First");
        let p2 = doc.create_element("p");
        let t2 = doc.create_text("Second");
        doc.append_child(root, html);
        doc.append_child(html, body);
        doc.append_child(body, p1);
        doc.append_child(p1, t1);
        doc.append_child(body, p2);
        doc.append_child(p2, t2);

        let tree = layout_doc(&doc);
        let body_box = &tree.root.children[0];
        let first = &body_box.children[0];
        let second = &body_box.children[1];

        assert!(
            second.rect.y > first.rect.y,
            "second p (y={}) should be below first p (y={})",
            second.rect.y,
            first.rect.y
        );
    }

    #[test]
    fn heading_larger_than_body() {
        let mut doc = Document::new();
        let root = doc.root();
        let html = doc.create_element("html");
        let body = doc.create_element("body");
        let h1 = doc.create_element("h1");
        let h1_text = doc.create_text("Title");
        let p = doc.create_element("p");
        let p_text = doc.create_text("Text");
        doc.append_child(root, html);
        doc.append_child(html, body);
        doc.append_child(body, h1);
        doc.append_child(h1, h1_text);
        doc.append_child(body, p);
        doc.append_child(p, p_text);

        let tree = layout_doc(&doc);
        let body_box = &tree.root.children[0];
        let h1_box = &body_box.children[0];
        let p_box = &body_box.children[1];

        assert!(
            h1_box.font_size > p_box.font_size,
            "h1 font_size ({}) should be > p font_size ({})",
            h1_box.font_size,
            p_box.font_size
        );
        assert_eq!(h1_box.font_size, 32.0);

        assert!(
            h1_box.rect.height > p_box.rect.height,
            "h1 height ({}) should be > p height ({})",
            h1_box.rect.height,
            p_box.rect.height
        );
    }

    #[test]
    fn body_has_default_margin() {
        let mut doc = Document::new();
        let root = doc.root();
        let html = doc.create_element("html");
        let body = doc.create_element("body");
        let p = doc.create_element("p");
        let text = doc.create_text("Test");
        doc.append_child(root, html);
        doc.append_child(html, body);
        doc.append_child(body, p);
        doc.append_child(p, text);

        let tree = layout_doc(&doc);
        let body_box = &tree.root.children[0];

        assert_eq!(body_box.margin.top, 8.0);
        assert_eq!(body_box.margin.right, 8.0);
        assert_eq!(body_box.margin.bottom, 8.0);
        assert_eq!(body_box.margin.left, 8.0);

        assert_eq!(body_box.rect.x, 8.0);
        assert_eq!(body_box.rect.y, 8.0);
    }

    #[test]
    fn text_wraps_at_container_width() {
        let mut doc = Document::new();
        let root = doc.root();
        let html = doc.create_element("html");
        let body = doc.create_element("body");
        let p = doc.create_element("p");
        let text =
            doc.create_text("The quick brown fox jumps over the lazy dog and more words to wrap");
        doc.append_child(root, html);
        doc.append_child(html, body);
        doc.append_child(body, p);
        doc.append_child(p, text);

        let font = test_font();
        let sheets = extract_stylesheets(&doc);
        let styled = resolve_styles(&doc, &sheets).unwrap();
        let tree = layout(&styled, &doc, 100.0, 600.0, &font);
        let body_box = &tree.root.children[0];
        let p_box = &body_box.children[0];

        // Count distinct y-positions to count visual lines.
        let mut ys: Vec<f32> = p_box.lines.iter().map(|l| l.y).collect();
        ys.sort_by(|a, b| a.partial_cmp(b).unwrap());
        ys.dedup_by(|a, b| (*a - *b).abs() < 0.01);

        assert!(
            ys.len() > 1,
            "text should wrap to multiple lines, got {} visual lines",
            ys.len()
        );
    }

    #[test]
    fn layout_produces_positive_dimensions() {
        let mut doc = Document::new();
        let root = doc.root();
        let html = doc.create_element("html");
        let body = doc.create_element("body");
        let div = doc.create_element("div");
        let text = doc.create_text("Content");
        doc.append_child(root, html);
        doc.append_child(html, body);
        doc.append_child(body, div);
        doc.append_child(div, text);

        let tree = layout_doc(&doc);

        for b in tree.iter() {
            assert!(b.rect.width >= 0.0, "width should be >= 0");
            assert!(b.rect.height >= 0.0, "height should be >= 0");
        }

        assert!(tree.height > 0.0, "layout height should be > 0");
    }

    #[test]
    fn head_is_hidden() {
        let mut doc = Document::new();
        let root = doc.root();
        let html = doc.create_element("html");
        let head = doc.create_element("head");
        let title = doc.create_element("title");
        let title_text = doc.create_text("Page Title");
        let body = doc.create_element("body");
        let p = doc.create_element("p");
        let p_text = doc.create_text("Visible");
        doc.append_child(root, html);
        doc.append_child(html, head);
        doc.append_child(head, title);
        doc.append_child(title, title_text);
        doc.append_child(html, body);
        doc.append_child(body, p);
        doc.append_child(p, p_text);

        let tree = layout_doc(&doc);

        assert_eq!(
            tree.root.children.len(),
            1,
            "html should have 1 child (body), head should be hidden"
        );
    }

    #[test]
    fn mixed_block_and_inline() {
        let mut doc = Document::new();
        let root = doc.root();
        let html = doc.create_element("html");
        let body = doc.create_element("body");
        let div = doc.create_element("div");
        let text1 = doc.create_text("Text");
        let p = doc.create_element("p");
        let p_text = doc.create_text("Block");
        let text2 = doc.create_text("More");
        doc.append_child(root, html);
        doc.append_child(html, body);
        doc.append_child(body, div);
        doc.append_child(div, text1);
        doc.append_child(div, p);
        doc.append_child(p, p_text);
        doc.append_child(div, text2);

        let tree = layout_doc(&doc);
        let body_box = &tree.root.children[0];
        let div_box = &body_box.children[0];

        assert_eq!(
            div_box.children.len(),
            3,
            "div should have 3 children (anon, block, anon), got {}",
            div_box.children.len()
        );

        assert!(matches!(div_box.children[0].box_type, BoxType::Anonymous));
        assert!(matches!(div_box.children[1].box_type, BoxType::Block(_)));
        assert!(matches!(div_box.children[2].box_type, BoxType::Anonymous));
    }

    #[test]
    fn inline_elements_contribute_text() {
        let mut doc = Document::new();
        let root = doc.root();
        let html = doc.create_element("html");
        let body = doc.create_element("body");
        let p = doc.create_element("p");
        let t1 = doc.create_text("Hello ");
        let em = doc.create_element("em");
        let t2 = doc.create_text("world");
        let t3 = doc.create_text("!");
        doc.append_child(root, html);
        doc.append_child(html, body);
        doc.append_child(body, p);
        doc.append_child(p, t1);
        doc.append_child(p, em);
        doc.append_child(em, t2);
        doc.append_child(p, t3);

        let tree = layout_doc(&doc);
        let body_box = &tree.root.children[0];
        let p_box = &body_box.children[0];

        assert!(!p_box.lines.is_empty());

        let first_y = p_box.lines[0].y;
        let line_texts: Vec<&str> = p_box
            .lines
            .iter()
            .filter(|l| (l.y - first_y).abs() < 0.01)
            .map(|l| l.text.as_str())
            .collect();
        let combined = line_texts.join("");
        assert!(
            combined.contains("Hello") && combined.contains("world") && combined.contains("!"),
            "line should contain all text, got: {combined}"
        );
    }

    #[test]
    fn collapse_whitespace_works() {
        assert_eq!(collapse_whitespace("hello  world"), "hello world");
        assert_eq!(collapse_whitespace("  spaces  "), " spaces ");
        assert_eq!(collapse_whitespace("\n\ttabs\n"), " tabs ");
        assert_eq!(collapse_whitespace("no-extra"), "no-extra");
        assert_eq!(collapse_whitespace("   "), " ");
    }

    #[test]
    fn content_width_respects_body_margin() {
        let mut doc = Document::new();
        let root = doc.root();
        let html = doc.create_element("html");
        let body = doc.create_element("body");
        let div = doc.create_element("div");
        let text = doc.create_text("Content");
        doc.append_child(root, html);
        doc.append_child(html, body);
        doc.append_child(body, div);
        doc.append_child(div, text);

        let font = test_font();
        let sheets = extract_stylesheets(&doc);
        let styled = resolve_styles(&doc, &sheets).unwrap();
        let tree = layout(&styled, &doc, 800.0, 600.0, &font);
        let body_box = &tree.root.children[0];

        assert_eq!(body_box.rect.width, 784.0);

        let div_box = &body_box.children[0];
        assert_eq!(div_box.rect.width, 784.0);
    }

    #[test]
    fn multiple_heading_levels() {
        let mut doc = Document::new();
        let root = doc.root();
        let html = doc.create_element("html");
        let body = doc.create_element("body");
        doc.append_child(root, html);
        doc.append_child(html, body);

        let tags = ["h1", "h2", "h3"];
        for tag in &tags {
            let h = doc.create_element(tag);
            let t = doc.create_text(tag);
            doc.append_child(body, h);
            doc.append_child(h, t);
        }

        let tree = layout_doc(&doc);
        let body_box = &tree.root.children[0];

        let h1 = &body_box.children[0];
        let h2 = &body_box.children[1];
        let h3 = &body_box.children[2];
        assert!(h1.font_size > h2.font_size);
        assert!(h2.font_size > h3.font_size);

        assert!(h2.rect.y > h1.rect.y);
        assert!(h3.rect.y > h2.rect.y);
    }

    #[test]
    fn layout_tree_iteration() {
        let mut doc = Document::new();
        let root = doc.root();
        let html = doc.create_element("html");
        let body = doc.create_element("body");
        let p = doc.create_element("p");
        let text = doc.create_text("Test");
        doc.append_child(root, html);
        doc.append_child(html, body);
        doc.append_child(body, p);
        doc.append_child(p, text);

        let tree = layout_doc(&doc);
        let count = tree.iter().count();
        assert!(count >= 3, "should have at least html, body, p boxes");
    }

    #[test]
    fn css_style_affects_layout() {
        let html_str = r#"<!DOCTYPE html>
<html>
<head>
<style>
p { margin-top: 50px; margin-bottom: 50px; }
</style>
</head>
<body>
<p>First</p>
<p>Second</p>
</body>
</html>"#;
        let doc = we_html::parse_html(html_str);
        let font = test_font();
        let sheets = extract_stylesheets(&doc);
        let styled = resolve_styles(&doc, &sheets).unwrap();
        let tree = layout(&styled, &doc, 800.0, 600.0, &font);

        let body_box = &tree.root.children[0];
        let first = &body_box.children[0];
        let second = &body_box.children[1];

        assert_eq!(first.margin.top, 50.0);
        assert_eq!(first.margin.bottom, 50.0);

        assert!(second.rect.y > first.rect.y + 100.0);
    }

    #[test]
    fn inline_style_affects_layout() {
        let html_str = r#"<!DOCTYPE html>
<html>
<body>
<div style="padding-top: 20px; padding-bottom: 20px;">
<p>Content</p>
</div>
</body>
</html>"#;
        let doc = we_html::parse_html(html_str);
        let font = test_font();
        let sheets = extract_stylesheets(&doc);
        let styled = resolve_styles(&doc, &sheets).unwrap();
        let tree = layout(&styled, &doc, 800.0, 600.0, &font);

        let body_box = &tree.root.children[0];
        let div_box = &body_box.children[0];

        assert_eq!(div_box.padding.top, 20.0);
        assert_eq!(div_box.padding.bottom, 20.0);
    }

    #[test]
    fn css_color_propagates_to_layout() {
        let html_str = r#"<!DOCTYPE html>
<html>
<head><style>p { color: red; background-color: blue; }</style></head>
<body><p>Colored</p></body>
</html>"#;
        let doc = we_html::parse_html(html_str);
        let font = test_font();
        let sheets = extract_stylesheets(&doc);
        let styled = resolve_styles(&doc, &sheets).unwrap();
        let tree = layout(&styled, &doc, 800.0, 600.0, &font);

        let body_box = &tree.root.children[0];
        let p_box = &body_box.children[0];

        assert_eq!(p_box.color, Color::rgb(255, 0, 0));
        assert_eq!(p_box.background_color, Color::rgb(0, 0, 255));
    }

    // --- New inline layout tests ---

    #[test]
    fn inline_elements_have_per_fragment_styling() {
        let html_str = r#"<!DOCTYPE html>
<html>
<head><style>em { color: red; }</style></head>
<body><p>Hello <em>world</em></p></body>
</html>"#;
        let doc = we_html::parse_html(html_str);
        let font = test_font();
        let sheets = extract_stylesheets(&doc);
        let styled = resolve_styles(&doc, &sheets).unwrap();
        let tree = layout(&styled, &doc, 800.0, 600.0, &font);

        let body_box = &tree.root.children[0];
        let p_box = &body_box.children[0];

        let colors: Vec<Color> = p_box.lines.iter().map(|l| l.color).collect();
        assert!(
            colors.iter().any(|c| *c == Color::rgb(0, 0, 0)),
            "should have black text"
        );
        assert!(
            colors.iter().any(|c| *c == Color::rgb(255, 0, 0)),
            "should have red text from <em>"
        );
    }

    #[test]
    fn br_element_forces_line_break() {
        let html_str = r#"<!DOCTYPE html>
<html>
<body><p>Line one<br>Line two</p></body>
</html>"#;
        let doc = we_html::parse_html(html_str);
        let font = test_font();
        let sheets = extract_stylesheets(&doc);
        let styled = resolve_styles(&doc, &sheets).unwrap();
        let tree = layout(&styled, &doc, 800.0, 600.0, &font);

        let body_box = &tree.root.children[0];
        let p_box = &body_box.children[0];

        let mut ys: Vec<f32> = p_box.lines.iter().map(|l| l.y).collect();
        ys.sort_by(|a, b| a.partial_cmp(b).unwrap());
        ys.dedup_by(|a, b| (*a - *b).abs() < 0.01);

        assert!(
            ys.len() >= 2,
            "<br> should produce 2 visual lines, got {}",
            ys.len()
        );
    }

    #[test]
    fn text_align_center() {
        let html_str = r#"<!DOCTYPE html>
<html>
<head><style>p { text-align: center; }</style></head>
<body><p>Hi</p></body>
</html>"#;
        let doc = we_html::parse_html(html_str);
        let font = test_font();
        let sheets = extract_stylesheets(&doc);
        let styled = resolve_styles(&doc, &sheets).unwrap();
        let tree = layout(&styled, &doc, 800.0, 600.0, &font);

        let body_box = &tree.root.children[0];
        let p_box = &body_box.children[0];

        assert!(!p_box.lines.is_empty());
        let first = &p_box.lines[0];
        // Center-aligned: text should be noticeably offset from content x.
        assert!(
            first.x > p_box.rect.x + 10.0,
            "center-aligned text x ({}) should be offset from content x ({})",
            first.x,
            p_box.rect.x
        );
    }

    #[test]
    fn text_align_right() {
        let html_str = r#"<!DOCTYPE html>
<html>
<head><style>p { text-align: right; }</style></head>
<body><p>Hi</p></body>
</html>"#;
        let doc = we_html::parse_html(html_str);
        let font = test_font();
        let sheets = extract_stylesheets(&doc);
        let styled = resolve_styles(&doc, &sheets).unwrap();
        let tree = layout(&styled, &doc, 800.0, 600.0, &font);

        let body_box = &tree.root.children[0];
        let p_box = &body_box.children[0];

        assert!(!p_box.lines.is_empty());
        let first = &p_box.lines[0];
        let right_edge = p_box.rect.x + p_box.rect.width;
        assert!(
            (first.x + first.width - right_edge).abs() < 1.0,
            "right-aligned text end ({}) should be near right edge ({})",
            first.x + first.width,
            right_edge
        );
    }

    #[test]
    fn inline_padding_offsets_text() {
        let html_str = r#"<!DOCTYPE html>
<html>
<head><style>span { padding-left: 20px; padding-right: 20px; }</style></head>
<body><p>A<span>B</span>C</p></body>
</html>"#;
        let doc = we_html::parse_html(html_str);
        let font = test_font();
        let sheets = extract_stylesheets(&doc);
        let styled = resolve_styles(&doc, &sheets).unwrap();
        let tree = layout(&styled, &doc, 800.0, 600.0, &font);

        let body_box = &tree.root.children[0];
        let p_box = &body_box.children[0];

        // Should have at least 3 fragments: A, B, C
        assert!(
            p_box.lines.len() >= 3,
            "should have fragments for A, B, C, got {}",
            p_box.lines.len()
        );

        // B should be offset by the span's padding.
        let a_frag = &p_box.lines[0];
        let b_frag = &p_box.lines[1];
        let gap = b_frag.x - (a_frag.x + a_frag.width);
        // Gap should include the 20px padding-left from the span.
        assert!(
            gap >= 19.0,
            "gap between A and B ({gap}) should include span padding-left (20px)"
        );
    }

    #[test]
    fn text_fragments_have_correct_font_size() {
        let html_str = r#"<!DOCTYPE html>
<html>
<body><h1>Big</h1><p>Small</p></body>
</html>"#;
        let doc = we_html::parse_html(html_str);
        let font = test_font();
        let sheets = extract_stylesheets(&doc);
        let styled = resolve_styles(&doc, &sheets).unwrap();
        let tree = layout(&styled, &doc, 800.0, 600.0, &font);

        let body_box = &tree.root.children[0];
        let h1_box = &body_box.children[0];
        let p_box = &body_box.children[1];

        assert!(!h1_box.lines.is_empty());
        assert!(!p_box.lines.is_empty());
        assert_eq!(h1_box.lines[0].font_size, 32.0);
        assert_eq!(p_box.lines[0].font_size, 16.0);
    }

    #[test]
    fn line_height_from_computed_style() {
        let html_str = r#"<!DOCTYPE html>
<html>
<head><style>p { line-height: 30px; }</style></head>
<body><p>Line one Line two Line three</p></body>
</html>"#;
        let doc = we_html::parse_html(html_str);
        let font = test_font();
        let sheets = extract_stylesheets(&doc);
        let styled = resolve_styles(&doc, &sheets).unwrap();
        // Narrow viewport to force wrapping.
        let tree = layout(&styled, &doc, 100.0, 600.0, &font);

        let body_box = &tree.root.children[0];
        let p_box = &body_box.children[0];

        let mut ys: Vec<f32> = p_box.lines.iter().map(|l| l.y).collect();
        ys.sort_by(|a, b| a.partial_cmp(b).unwrap());
        ys.dedup_by(|a, b| (*a - *b).abs() < 0.01);

        if ys.len() >= 2 {
            let gap = ys[1] - ys[0];
            assert!(
                (gap - 30.0).abs() < 1.0,
                "line spacing ({gap}) should be ~30px from line-height"
            );
        }
    }
}