use egui::{pos2, vec2, Align, Layout, Rect, Ui, UiBuilder};

/// Extension trait to add rows functionality to egui::Ui
pub trait UiRowsExt {
    /// Divide the available height into equal rows.
    ///
    /// Similar to `ui.columns()` but for horizontal rows.
    /// Each row gets equal height and spans the full width.
    fn rows<R>(&mut self, num_rows: usize, add_contents: impl FnOnce(&mut [Ui]) -> R) -> R;
}

impl UiRowsExt for Ui {
    #[inline]
    fn rows<R>(&mut self, num_rows: usize, add_contents: impl FnOnce(&mut [Ui]) -> R) -> R {
        self.rows_dyn(num_rows, Box::new(add_contents))
    }
}

// Private helper trait to mirror egui's internal pattern
trait UiRowsExtPrivate {
    fn rows_dyn<'c, R>(
        &mut self,
        num_rows: usize,
        add_contents: Box<dyn FnOnce(&mut [Ui]) -> R + 'c>,
    ) -> R;
}

impl UiRowsExtPrivate for Ui {
    fn rows_dyn<'c, R>(
        &mut self,
        num_rows: usize,
        add_contents: Box<dyn FnOnce(&mut [Self]) -> R + 'c>,
    ) -> R {
        let spacing = self.spacing().item_spacing.y;
        let total_spacing = spacing * (num_rows as f32 - 1.0);
        let row_height = (self.available_height() - total_spacing) / (num_rows as f32);
        let top_left = self.cursor().min;

        let mut rows: Vec<Self> = (0..num_rows)
            .map(|row_idx| {
                let pos = top_left + vec2(0.0, (row_idx as f32) * (row_height + spacing));
                let child_rect = Rect::from_min_max(
                    pos,
                    pos2(self.max_rect().right_bottom().x, pos.y + row_height),
                );
                let mut row_ui = self.new_child(
                    UiBuilder::new()
                        .max_rect(child_rect)
                        .layout(Layout::left_to_right(Align::Center)),
                );
                row_ui.set_height(row_height);
                row_ui
            })
            .collect();

        let result = add_contents(&mut rows[..]);

        let mut max_row_height = row_height;
        let mut max_width = 0.0;
        for row in &rows {
            max_row_height = max_row_height.max(row.min_rect().height());
            max_width = row.min_size().x.max(max_width);
        }

        // Make sure we fit everything next frame:
        let total_required_height = total_spacing + max_row_height * (num_rows as f32);

        let size = vec2(
            max_width,
            self.available_height().max(total_required_height),
        );
        self.advance_cursor_after_rect(Rect::from_min_size(top_left, size));
        result
    }
}
