termusic 0.13.2

use std::{cell::RefCell, num::NonZeroUsize, rc::Rc};

use lru::LruCache;
use tuirealm::ratatui::{
    buffer::Buffer,
    layout::{Constraint, Layout, Offset, Rect},
    widgets::{Block, Paragraph, Widget},
};

type Rects = Rc<[Rect]>;
type Cache = LruCache<(Rect, DynamicHeightGrid), Rects>;

const CACHE_SIZE: NonZeroUsize = NonZeroUsize::new(10).unwrap();

// Cache the computations of the layout, so that we dont have to much work on re-prints if nothing changed.
thread_local! {
    static LAYOUT_CACHE: RefCell<Cache> = RefCell::new(Cache::new(
        CACHE_SIZE
    ));
}

// NOTE: the struct below could likely be changed to allow dynamic elems widths / heights, but currently only allows one size.

/// A dynamic grid where all elements are of `elem_width` and dynamic height.
/// Split automatically into rows and colums, if there is not enough space on the current row.
///
/// # Behavior
///
/// ## Low width
///
/// If there is lower width available than the set element width, ex `3` width is available, but element width is `10`,
/// then there will only be one element per row, which has `3` width.
///
/// ## Low height
///
/// If available height is lower than the set element height, ex `3` height is available, but element height is `10`,
/// then there will be only one element per area, which has `3` height.
///
/// If there is not enough height for each element, then if `draw_row_low_space` is set, the last row will have the remaining height.
/// Example: `5` height is available, but element height is `3`, then the first row will have `3` height, but the last row has `2` height.
///
/// If `draw_row_low_space` is unset, the last element will have height `0` and be effectively not drawn.
///
/// ## Uneven width
///
/// If there is more width available than set element width, then if `distribute_row_space` is set, elements will dynamically
/// get more width to consume the remainder width, distributed evenly among all elements in the row.
/// This will mean the width of each element will get up to `(element_width*2)-1` at most.
///
/// If `distribute_row_space` is unset, all elements in a row will have at most `element_width` width, leaving the remaining area empty.
/// (elements are *not* spaced)
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct DynamicHeightGrid {
    elems_height: Box<[u16]>,
    elem_width: u16,
    row_spacing: u16,
    draw_row_low_space: bool,
    distribute_row_space: bool,
    focus_node: Option<usize>,
}

impl DynamicHeightGrid {
    pub fn new<E: Into<Box<[u16]>>>(elems: E, elem_width: u16) -> Self {
        Self {
            elems_height: elems.into(),
            elem_width,
            row_spacing: 0,
            draw_row_low_space: false,
            distribute_row_space: false,
            focus_node: None,
        }
    }

    /// Set spacing between the dynamic rows.
    ///
    /// Default: `0`
    pub fn with_row_spacing(mut self, spacing: u16) -> Self {
        self.row_spacing = spacing;
        self
    }

    /// Draw the last row even if there is not enough space for a full element.
    ///
    /// Default: `false`
    pub fn draw_row_low_space(mut self) -> Self {
        self.draw_row_low_space = true;
        self
    }

    /// Distribute remaining row space among the elements in the row.
    ///
    /// Default: `false`
    pub fn distribute_row_space(mut self) -> Self {
        self.distribute_row_space = true;
        self
    }

    /// Set a node to skip to.
    ///
    /// Automatically figures out which row this node is on and does not allocate area to previous rows.
    pub fn focus_node(mut self, focus: Option<usize>) -> Self {
        self.focus_node = focus;
        self
    }

    /// Split `area` into `elems`, will always be `elems` amount.
    /// Elements areas that dont fit are `Rect(0,0,0,0)`.
    pub fn split(&self, area: Rect) -> Rects {
        LAYOUT_CACHE.with_borrow_mut(|c| {
            let key = (area, self.clone());
            c.get_or_insert(key, || self.get_areas_inner(area)).clone()
        })
    }

    /// Internal function for [`split`](Self::split) to reduce nesting.
    fn get_areas_inner(&self, area: Rect) -> Rects {
        let mut remaining_area = area;
        let mut cells = Vec::new();

        let mut remaining_elems = self.elems_height.iter().peekable();

        let elems_per_row = {
            let result = remaining_area
                .width
                .checked_div(self.elem_width)
                .unwrap_or_default();

            // have one element per row if available with is not 0
            // this is to avoid empty space if for example available width is 50, but wanted width is 55
            if result == 0 && remaining_area.width != 0 {
                1
            } else {
                result
            }
        };

        // too low area to even a single element, return all 0-length areas
        if elems_per_row == 0 {
            // fill the vec to be "elems" amount with 0-width/height rects.
            for _ in remaining_elems {
                cells.push(Rect::default());
            }

            return cells.into();
        }

        // only run total height calculation for skip if a skip is defined
        let mut rows_to_skip = if let Some(focus_idx) = self.focus_node {
            // get the total height this grid would consume if fully allocating all elements
            let total_height = self
                .elems_height
                .clone()
                .chunks(usize::from(elems_per_row))
                .fold(0usize, |acc, v| {
                    acc + usize::from(v.iter().copied().max().unwrap_or_default())
                });

            // dont skip rows if all elements fit into the area
            if total_height > usize::from(remaining_area.height) {
                focus_idx
                    .checked_div(usize::from(elems_per_row))
                    .unwrap_or_default()
            } else {
                0
            }
        } else {
            0
        };

        while remaining_elems
            .peek()
            .is_some_and(|v| (remaining_area.height >= **v) || self.draw_row_low_space)
            && !remaining_area.is_empty()
            && remaining_area.height > 0
        {
            if rows_to_skip > 0 {
                for _ in 0..elems_per_row {
                    if remaining_elems.next().is_none() {
                        break;
                    }

                    cells.push(Rect::default());
                }
                rows_to_skip -= 1;
                continue;
            }

            // dont add a spacing for the first row
            let spacing = if remaining_area == area {
                0
            } else {
                self.row_spacing
            };
            let [_row_spacer, spaced_area] =
                Layout::vertical([Constraint::Length(spacing), Constraint::Fill(0)])
                    .areas(remaining_area);
            remaining_area = spaced_area;

            let chunks = if self.distribute_row_space {
                let constraints = (0..elems_per_row).map(|_| Constraint::Min(self.elem_width));

                Layout::horizontal(constraints).split(remaining_area)
            } else {
                let constraints = (0..elems_per_row).map(|_| Constraint::Length(self.elem_width));

                Layout::horizontal(constraints).split(remaining_area)
            };

            let mut highest_height = 0;

            for chunk in chunks.iter() {
                let mut chunk = *chunk;
                // only add as many cells as there are requested elements
                let Some(elem_height) = remaining_elems.peek().copied().copied() else {
                    break;
                };

                // dont add the chunk if the height is not enough
                if remaining_area.height < elem_height && !self.draw_row_low_space {
                    break;
                }
                // actually advance the iterator
                remaining_elems.next();

                chunk.height = elem_height.min(remaining_area.height);
                highest_height = highest_height.max(chunk.height);

                cells.push(chunk);
            }

            remaining_area.height = remaining_area.height.saturating_sub(highest_height);
            remaining_area = remaining_area.offset(Offset {
                x: 0,
                y: i32::from(highest_height),
            });
        }

        // fill the vec to be "elems" amount with 0-width/height rects.
        for _ in remaining_elems {
            cells.push(Rect::default());
        }

        cells.into()
    }
}

/// Debug-only visualization of the areas
impl Widget for DynamicHeightGrid {
    fn render(self, area: Rect, buf: &mut Buffer) {
        let cells = self.split(area);

        for (i, cell) in cells.iter().enumerate() {
            Paragraph::new(format!("Area {:02}", i + 1))
                .block(Block::bordered())
                .render(*cell, buf);
        }
    }
}

#[cfg(test)]
mod tests {
    use pretty_assertions::assert_eq;
    use tuirealm::ratatui::layout::Rect;

    use super::DynamicHeightGrid;

    #[test]
    fn should_zero_on_zero_area() {
        // test to know if there is a infinite loop on too low area
        let area = Rect::new(0, 0, 0, 0);
        let elems = [3, 3, 3];

        let areas = DynamicHeightGrid::new(elems, 10).split(area);

        assert_eq!(areas.len(), 3);
        assert_eq!(areas[0], Rect::new(0, 0, 0, 0));
        assert_eq!(areas[1], Rect::new(0, 0, 0, 0));
        assert_eq!(areas[2], Rect::new(0, 0, 0, 0));
    }

    #[test]
    fn should_split_all_single_row_uniform() {
        let area = Rect::new(0, 0, 30, 3);
        let elems = [3, 3, 3];

        let areas = DynamicHeightGrid::new(elems, 10).split(area);

        assert_eq!(areas.len(), 3);
        assert_eq!(areas[0], Rect::new(0, 0, 10, 3));
        assert_eq!(areas[1], Rect::new(10, 0, 10, 3));
        assert_eq!(areas[2], Rect::new(20, 0, 10, 3));
    }

    #[test]
    fn should_split_all_2_rows_uniform() {
        let area = Rect::new(0, 0, 20, 6);
        let elems = [3, 3, 3];

        let areas = DynamicHeightGrid::new(elems, 10).split(area);

        assert_eq!(areas.len(), 3);
        assert_eq!(areas[0], Rect::new(0, 0, 10, 3));
        assert_eq!(areas[1], Rect::new(10, 0, 10, 3));
        assert_eq!(areas[2], Rect::new(0, 3, 10, 3));
    }

    #[test]
    fn should_not_split_new_row_low_space_uniform() {
        let area = Rect::new(0, 0, 20, 3);
        let elems = [3, 3, 3];

        let areas = DynamicHeightGrid::new(elems, 10).split(area);

        assert_eq!(areas.len(), 3);
        assert_eq!(areas[0], Rect::new(0, 0, 10, 3));
        assert_eq!(areas[1], Rect::new(10, 0, 10, 3));
        assert_eq!(areas[2], Rect::new(0, 0, 0, 0));
    }

    #[test]
    fn should_split_new_row_low_space_uniform() {
        let area = Rect::new(0, 0, 20, 3);
        let elems = [3, 3, 3];

        let areas = DynamicHeightGrid::new(elems, 10)
            .draw_row_low_space()
            .split(area);

        assert_eq!(areas.len(), 3);
        assert_eq!(areas[0], Rect::new(0, 0, 10, 3));
        assert_eq!(areas[1], Rect::new(10, 0, 10, 3));
        assert_eq!(areas[2], Rect::new(0, 0, 0, 0));
    }

    #[test]
    fn should_have_row_spacing_uniform() {
        let area = Rect::new(0, 0, 20, 7);
        let elems = [3, 3, 3];

        let areas = DynamicHeightGrid::new(elems, 10)
            .with_row_spacing(1)
            .split(area);

        assert_eq!(areas.len(), 3);
        assert_eq!(areas[0], Rect::new(0, 0, 10, 3));
        assert_eq!(areas[1], Rect::new(10, 0, 10, 3));
        assert_eq!(areas[2], Rect::new(0, 4, 10, 3));
    }

    #[test]
    fn should_split_all_single_row_no_leftover_space_uniform() {
        let area = Rect::new(0, 0, 33, 3);
        let elems = [3, 3, 3];

        let areas = DynamicHeightGrid::new(elems, 10)
            .distribute_row_space()
            .split(area);

        assert_eq!(areas.len(), 3);
        assert_eq!(areas[0], Rect::new(0, 0, 11, 3));
        assert_eq!(areas[1], Rect::new(11, 0, 11, 3));
        assert_eq!(areas[2], Rect::new(22, 0, 11, 3));
    }

    #[test]
    fn should_split_all_single_row() {
        let area = Rect::new(0, 0, 30, 4);
        let elems = [3, 4, 3];

        let areas = DynamicHeightGrid::new(elems, 10).split(area);

        assert_eq!(areas.len(), 3);
        assert_eq!(areas[0], Rect::new(0, 0, 10, 3));
        assert_eq!(areas[1], Rect::new(10, 0, 10, 4));
        assert_eq!(areas[2], Rect::new(20, 0, 10, 3));
    }

    #[test]
    fn should_split_all_2_rows() {
        let area = Rect::new(0, 0, 20, 7);
        let elems = [3, 4, 3];

        let areas = DynamicHeightGrid::new(elems, 10).split(area);

        assert_eq!(areas.len(), 3);
        assert_eq!(areas[0], Rect::new(0, 0, 10, 3));
        assert_eq!(areas[1], Rect::new(10, 0, 10, 4));
        assert_eq!(areas[2], Rect::new(0, 4, 10, 3));
    }

    #[test]
    fn should_not_overflow_given_area() {
        let area = Rect::new(0, 0, 20, 6);
        let elems = [3, 4, 3];

        let areas = DynamicHeightGrid::new(elems, 10).split(area);

        assert_eq!(areas.len(), 3);
        assert_eq!(areas[0], Rect::new(0, 0, 10, 3));
        assert_eq!(areas[1], Rect::new(10, 0, 10, 4));
        assert_eq!(areas[2], Rect::new(0, 0, 0, 0));

        assert!(area.contains(areas[1].positions().last().unwrap()));
    }

    #[test]
    fn should_not_overflow_given_area_draw_last_row() {
        let area = Rect::new(0, 0, 20, 6);
        let elems = [3, 4, 3];

        let areas = DynamicHeightGrid::new(elems, 10)
            .draw_row_low_space()
            .split(area);

        assert_eq!(areas.len(), 3);
        assert_eq!(areas[0], Rect::new(0, 0, 10, 3));
        assert_eq!(areas[1], Rect::new(10, 0, 10, 4));
        assert_eq!(areas[2], Rect::new(0, 4, 10, 2));

        assert!(area.contains(areas[2].positions().last().unwrap()));
    }

    #[test]
    fn should_not_skip_if_enough_area() {
        let area = Rect::new(0, 0, 20, 7);
        let elems = [3, 4, 3];

        let areas = DynamicHeightGrid::new(elems, 10)
            .focus_node(Some(2))
            .split(area);

        assert_eq!(areas.len(), 3);
        assert_eq!(areas[0], Rect::new(0, 0, 10, 3));
        assert_eq!(areas[1], Rect::new(10, 0, 10, 4));
        assert_eq!(areas[2], Rect::new(0, 4, 10, 3));
    }

    #[test]
    fn should_skip_if_not_enough_area() {
        let area = Rect::new(0, 0, 30, 7);
        let elems = [3, 4, 3, 4, 3, 3, 4];

        let areas = DynamicHeightGrid::new(elems, 10)
            .focus_node(Some(3))
            .split(area);

        assert_eq!(areas.len(), 7);
        assert_eq!(areas[0], Rect::new(0, 0, 0, 0));
        assert_eq!(areas[1], Rect::new(0, 0, 0, 0));
        assert_eq!(areas[2], Rect::new(0, 0, 0, 0));
        assert_eq!(areas[3], Rect::new(0, 0, 10, 4));
        assert_eq!(areas[4], Rect::new(10, 0, 10, 3));
        assert_eq!(areas[5], Rect::new(20, 0, 10, 3));
        assert_eq!(areas[6], Rect::new(0, 0, 0, 0));
    }

    #[test]
    fn should_draw_in_lower_width() {
        let area = Rect::new(0, 0, 3, 10);
        let elems = [3, 4, 3];

        let areas = DynamicHeightGrid::new(elems, 10).split(area);

        assert_eq!(areas.len(), 3);
        assert_eq!(areas[0], Rect::new(0, 0, 3, 3));
        assert_eq!(areas[1], Rect::new(0, 3, 3, 4));
        assert_eq!(areas[2], Rect::new(0, 7, 3, 3));
    }

    #[test]
    fn should_not_panic_not_enough_width_for_one_elem_with_focus_node() {
        // test relies on having less width available than a element wants
        // and having a focus node set
        let area = Rect::new(0, 0, 3, 10);
        let elems = [3, 4, 3];

        let areas = DynamicHeightGrid::new(elems, 3)
            .focus_node(Some(0))
            .split(area);

        assert_eq!(areas.len(), 3);
    }
}