oo-ide 0.0.4

//! Generic tree widget.
//!
//! [`TreeNode<T>`] is a recursive node carrying an arbitrary payload `T`.
//! [`TreeWidget<T>`] is the stateful widget: it owns the root nodes, tracks
//! the cursor, and handles expand / collapse.  Rendering is driven by a
//! caller-supplied closure so this module has no opinion about styling.
//!
//! ```text
//! TreeNode<T>          — recursive data node
//! TreeWidget<T>        — owns Vec<TreeNode<T>> + cursor, navigation, render
//! ```
//!
//! # Cursor representation
//!
//! The cursor is a `Vec<usize>` path from the root, e.g. `[1, 0, 2]` means
//! `roots[1].children[0].children[2]`.  Navigation is O(depth) rather than
//! O(n), which is effectively O(1) for any realistic tree.  No unsafe code is
//! required.
//!
//! A cached `cursor_flat_idx: usize` is maintained alongside the path so the
//! render scroll calculation never has to re-walk the tree.
//!
//! For the file-system specialisation see [`super::file_tree`].

use std::cell::Cell;

use ratatui::{
    buffer::Buffer,
    layout::Rect,
    text::Line,
    widgets::{List, ListItem, ListState, StatefulWidget},
};

/// Whether a directory node's children have been loaded from disk.
#[derive(Debug, Clone)]
pub enum ChildrenState<T> {
    /// Directory that has never been expanded — children not yet read.
    NotLoaded,
    /// File (always empty) or directory whose children have been read.
    Loaded(Vec<TreeNode<T>>),
}

/// A single node in the tree.  `T` is the caller's payload (e.g. `PathBuf`).
#[derive(Debug, Clone)]
pub(crate) struct TreeNode<T> {
    pub data: T,
    pub children: ChildrenState<T>,
    pub expanded: bool,
    /// Whether this node represents a directory (vs. a file leaf).
    pub is_dir: bool,
}

impl<T> TreeNode<T> {
    /// Leaf node (file — no children, never expandable).
    pub fn leaf(data: T) -> Self {
        Self {
            data,
            children: ChildrenState::Loaded(vec![]),
            expanded: false,
            is_dir: false,
        }
    }

    /// Interior node with already-loaded children (starts collapsed).
    pub fn inner(data: T, children: Vec<TreeNode<T>>) -> Self {
        Self {
            data,
            children: ChildrenState::Loaded(children),
            expanded: false,
            is_dir: true,
        }
    }

    /// Directory node whose children have not been loaded yet.
    pub fn unloaded(data: T) -> Self {
        Self {
            data,
            children: ChildrenState::NotLoaded,
            expanded: false,
            is_dir: true,
        }
    }

    pub fn is_leaf(&self) -> bool {
        !self.is_dir
    }

    /// Append all currently-visible nodes to `out` (depth-first, respecting
    /// `expanded`).  `depth` is the visual nesting level passed to callers.
    pub(super) fn flatten<'a>(&'a self, depth: usize, out: &mut Vec<FlatNode<'a, T>>) {
        out.push(FlatNode { node: self, depth });
        if self.expanded
            && let ChildrenState::Loaded(children) = &self.children {
                for child in children {
                    child.flatten(depth + 1, out);
                }
            }
    }
}

/// A shared reference to a visible node together with its nesting depth.
#[derive(Clone, Copy)]
pub struct FlatNode<'a, T> {
    pub node: &'a TreeNode<T>,
    pub depth: usize,
}

/// Path-based cursor.  Each element is a child index at that level, so
/// `path == [1, 0, 2]` identifies `roots[1].children[0].children[2]`.
///
/// Invariant: `path` is never empty while the tree has at least one root.
#[derive(Debug, Clone, Default)]
pub struct Cursor {
    path: Vec<usize>,
}

impl Cursor {
    fn root(idx: usize) -> Self {
        Self { path: vec![idx] }
    }

    /// How many levels below the root this node sits (0 = root-level).
    pub fn depth(&self) -> usize {
        self.path.len().saturating_sub(1)
    }

    /// The raw index path, e.g. for serialisation / restore.
    #[allow(dead_code)]
    pub fn path(&self) -> &[usize] {
        &self.path
    }
}

/// Return value from [`TreeWidget::expand_or_enter`] and
/// [`TreeWidget::toggle_selected`] indicating whether the caller must load
/// the selected directory's children before the expansion can complete.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TreeAction {
    /// Operation completed; no further action needed.
    Done,
    /// The selected node is a directory whose children have not been loaded.
    /// The caller must read the directory and call `inject_children`.
    NeedsLoad,
}

/// Stateful tree: owns the root nodes and a path-based cursor.
///
/// `T` must be `Clone` only because `TreeNode` derives `Clone`; the widget
/// itself does not clone nodes during normal operation.
#[derive(Debug)]
pub struct TreeWidget<T> {
    roots: Vec<TreeNode<T>>,
    cursor: Cursor,
    /// Linear index of the selected node among currently-visible nodes.
    /// Maintained incrementally so `render` never has to recompute it.
    cursor_flat_idx: usize,
    /// Scroll offset (first visible row). Updated lazily in `render` so that
    /// the viewport only moves when the cursor would go out of bounds.
    /// Uses `Cell` so that `render(&self)` can update it without `&mut self`.
    scroll: Cell<usize>,
}

impl<T: Clone> TreeWidget<T> {
    pub fn new(roots: Vec<TreeNode<T>>) -> Self {
        let cursor = if roots.is_empty() {
            Cursor::default()
        } else {
            Cursor::root(0)
        };
        Self {
            roots,
            cursor,
            cursor_flat_idx: 0,
            scroll: Cell::new(0),
        }
    }

    #[allow(dead_code)]
    pub fn roots(&self) -> &[TreeNode<T>] {
        &self.roots
    }

    #[allow(dead_code)]
    pub fn cursor(&self) -> &Cursor {
        &self.cursor
    }

    #[allow(dead_code)]
    pub fn cursor_flat_idx(&self) -> usize {
        self.cursor_flat_idx
    }

    /// Resolve a path to a shared node reference.  Returns `None` if the path
    /// is empty or out-of-bounds (should not happen under normal operation).
    pub(crate) fn resolve_path(&self, path: &[usize]) -> Option<&TreeNode<T>> {
        let (&first, rest) = path.split_first()?;
        let mut node = self.roots.get(first)?;
        for &idx in rest {
            match &node.children {
                ChildrenState::Loaded(ch) => node = ch.get(idx)?,
                ChildrenState::NotLoaded => return None,
            }
        }
        Some(node)
    }

    /// Resolve a path to a mutable node reference.
    pub(crate) fn resolve_path_mut(&mut self, path: &[usize]) -> Option<&mut TreeNode<T>> {
        let (&first, rest) = path.split_first()?;
        let mut node = self.roots.get_mut(first)?;
        for &idx in rest {
            let ch = match &mut node.children {
                ChildrenState::Loaded(ch) => ch,
                ChildrenState::NotLoaded => return None,
            };
            node = ch.get_mut(idx)?;
        }
        Some(node)
    }

    /// Resolve the current cursor path.
    fn current(&self) -> Option<&TreeNode<T>> {
        self.resolve_path(&self.cursor.path)
    }

    /// Resolve the current cursor path mutably.
    fn current_mut(&mut self) -> Option<&mut TreeNode<T>> {
        let path = self.cursor.path.clone();
        self.resolve_path_mut(&path)
    }

    /// Return the sibling count for the node at `path`.
    /// For root-level nodes that is `roots.len()`; otherwise the parent's
    /// `children.len()`.
    fn sibling_count(&self, path: &[usize]) -> usize {
        if path.len() <= 1 {
            self.roots.len()
        } else {
            let parent_path = &path[..path.len() - 1];
            self.resolve_path(parent_path)
                .and_then(|p| match &p.children {
                    ChildrenState::Loaded(ch) => Some(ch.len()),
                    ChildrenState::NotLoaded => None,
                })
                .unwrap_or(0)
        }
    }

    /// Move the cursor to the previous visible node.  O(depth).
    pub fn move_up(&mut self) {
        if self.cursor.path.is_empty() {
            return;
        }

        let last = *self.cursor.path.last().unwrap();

        if last == 0 {
            // First sibling: jump to parent, which is directly above us in
            // the visible list.
            if self.cursor.path.len() > 1 {
                self.cursor.path.pop();
                self.cursor_flat_idx -= 1;
            }
            return;
        }

        // Step to previous sibling.
        *self.cursor.path.last_mut().unwrap() -= 1;

        // Descend into the deepest expanded child of the new sibling.
        while let Some(node) = self.current() {
            let last_child = match (&node.children, node.expanded) {
                (ChildrenState::Loaded(ch), true) if !ch.is_empty() => ch.len() - 1,
                _ => break,
            };
            self.cursor.path.push(last_child);
        }

        self.sync_flat_idx();
    }

    /// Move the cursor to the next visible node.  O(depth).
    pub fn move_down(&mut self) {
        if self.cursor.path.is_empty() {
            return;
        }

        // Case 1: expanded node with loaded children → enter first child.
        {
            let node = match self.current() {
                Some(n) => n,
                None => return,
            };
            let has_visible_children = matches!(
                &node.children,
                ChildrenState::Loaded(ch) if node.expanded && !ch.is_empty()
            );
            if has_visible_children {
                self.cursor.path.push(0);
                self.cursor_flat_idx += 1;
                return;
            }
        }

        // Case 2: walk up until we find a level with a next sibling.
        let mut path = self.cursor.path.clone();
        loop {
            let sibling_count = self.sibling_count(&path);
            let idx = path.last_mut().unwrap();
            if *idx + 1 < sibling_count {
                *idx += 1;
                self.cursor.path = path;
                self.cursor_flat_idx += 1;
                return;
            }
            path.pop();
            if path.is_empty() {
                // Already at the last visible node.
                return;
            }
        }
    }

    pub fn move_up_n(&mut self, n: usize) {
        for _ in 0..n {
            self.move_up();
        }
    }

    pub fn move_down_n(&mut self, n: usize) {
        for _ in 0..n {
            self.move_down();
        }
    }

    /// Expand a collapsed interior node; if already expanded, move into its
    /// first child.  Returns [`TreeAction::NeedsLoad`] when the node is a
    /// directory whose children have not been read yet — the caller must load
    /// them and call [`Self::inject_children`].
    pub fn expand_or_enter(&mut self) -> TreeAction {
        let (is_leaf, expanded, not_loaded) = match self.current() {
            Some(n) => (
                n.is_leaf(),
                n.expanded,
                matches!(n.children, ChildrenState::NotLoaded),
            ),
            None => return TreeAction::Done,
        };
        if is_leaf {
            return TreeAction::Done;
        }
        if not_loaded {
            return TreeAction::NeedsLoad;
        }
        if !expanded {
            self.current_mut().unwrap().expanded = true;
        } else {
            // Already expanded — step into first child.
            self.cursor.path.push(0);
            self.cursor_flat_idx += 1;
        }
        TreeAction::Done
    }

    /// Collapse an expanded interior node; if already collapsed (or a leaf),
    /// jump to the parent.
    pub fn collapse_or_leave(&mut self) {
        let (is_leaf, expanded) = match self.current() {
            Some(n) => (n.is_leaf(), n.expanded),
            None => return,
        };
        if !is_leaf && expanded {
            self.current_mut().unwrap().expanded = false;
            // Cursor stays on the now-collapsed node; flat index is unchanged
            // because the node itself remains visible.
        } else if self.cursor.path.len() > 1 {
            // Jump to parent.
            self.cursor.path.pop();
            self.cursor_flat_idx -= 1;
        }
    }

    /// Toggle expand / collapse of the selected node.  Returns
    /// [`TreeAction::NeedsLoad`] when the node is an unloaded directory.
    ///
    /// If collapsing while the cursor is inside the subtree (possible via
    /// programmatic navigation), the path is truncated to the collapsed node
    /// and `cursor_flat_idx` is resynced.
    pub fn toggle_selected(&mut self) -> TreeAction {
        let (is_leaf, expanded, not_loaded) = match self.current() {
            Some(n) => (
                n.is_leaf(),
                n.expanded,
                matches!(n.children, ChildrenState::NotLoaded),
            ),
            None => return TreeAction::Done,
        };
        if is_leaf {
            return TreeAction::Done;
        }
        if not_loaded {
            return TreeAction::NeedsLoad;
        }
        let path_len = self.cursor.path.len();
        self.current_mut().unwrap().expanded = !expanded;

        if expanded {
            // Just collapsed — truncate any stale deeper path.
            self.cursor.path.truncate(path_len);
            // Recompute flat index: O(n), but only on explicit toggle.
            self.sync_flat_idx();
        }
        TreeAction::Done
    }

    /// Inject loaded children into the node at `path` and mark it expanded.
    ///
    /// Called when an async directory-read result arrives for a previously
    /// `NotLoaded` node.  Cursor validity is preserved because injecting into
    /// a `NotLoaded` node does not shift any sibling indices.
    pub fn inject_children(&mut self, path: &[usize], new_children: Vec<TreeNode<T>>) {
        if let Some(node) = self.resolve_path_mut(path) {
            node.children = ChildrenState::Loaded(new_children);
            node.expanded = true;
        }
        self.sync_flat_idx();
    }

    /// Return a shared reference to the currently selected node and its depth.
    pub fn selected(&self) -> Option<FlatNode<'_, T>> {
        self.current().map(|node| FlatNode {
            node,
            depth: self.cursor.depth(),
        })
    }

    /// Recompute `cursor_flat_idx` by scanning the visible list for the
    /// node our path resolves to.  O(n) — call sparingly.
    pub(crate) fn sync_flat_idx(&mut self) {
        let target = match self.resolve_path(&self.cursor.path) {
            Some(n) => n as *const _,
            None => {
                self.cursor_flat_idx = 0;
                return;
            }
        };
        let mut flat = Vec::new();
        for root in &self.roots {
            root.flatten(0, &mut flat);
        }
        self.cursor_flat_idx = flat
            .iter()
            .position(|f| std::ptr::eq(f.node, target))
            .unwrap_or(0);
    }

    /// Set the internal cursor path (index path) directly and resync flat idx.
    pub(crate) fn set_cursor_path(&mut self, path: &[usize]) {
        self.cursor.path = path.to_vec();
        self.sync_flat_idx();
    }

    /// Collect all currently-visible nodes into a flat `Vec`.
    /// Only called by `render`; navigation never uses this.
    pub fn flatten(&self) -> Vec<FlatNode<'_, T>> {
        let mut out = Vec::new();
        for root in &self.roots {
            root.flatten(0, &mut out);
        }
        out
    }

    /// Render the visible tree into `area` using `buf`.
    ///
    /// `render_node(flat_node, is_selected) -> Line` controls all visual
    /// styling — icons, indentation, colours.
    pub fn render<F>(&self, area: Rect, buf: &mut Buffer, render_node: F)
    where
        F: Fn(&FlatNode<'_, T>, bool) -> Line<'static>,
    {
        let height = area.height as usize;
        let cursor = self.cursor_flat_idx;

        // Scroll only when the cursor would leave the visible area.
        let scroll = {
            let s = self.scroll.get();
            let s = if cursor < s { cursor } else { s };
            let s = if height > 0 && cursor >= s + height {
                cursor + 1 - height
            } else {
                s
            };
            self.scroll.set(s);
            s
        };

        let flat = self.flatten();

        let items: Vec<ListItem> = flat
            .iter()
            .enumerate()
            .skip(scroll)
            .take(height)
            .map(|(i, fn_)| ListItem::new(render_node(fn_, i == cursor)))
            .collect();

        let mut state = ListState::default();
        state.select(Some(cursor.saturating_sub(scroll)));
        StatefulWidget::render(List::new(items), area, buf, &mut state);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Build a small tree:
    /// ```text
    /// 0
    /// 1
    ///   1.0
    ///   1.1
    ///     1.1.0
    /// 2
    /// ```
    fn sample() -> TreeWidget<&'static str> {
        TreeWidget::new(vec![
            TreeNode::leaf("0"),
            TreeNode::inner(
                "1",
                vec![
                    TreeNode::leaf("1.0"),
                    TreeNode::inner("1.1", vec![TreeNode::leaf("1.1.0")]),
                ],
            ),
            TreeNode::leaf("2"),
        ])
    }

    fn data(w: &TreeWidget<&'static str>) -> &'static str {
        w.selected().unwrap().node.data
    }

    #[test]
    fn initial_cursor_at_first_root() {
        let w = sample();
        assert_eq!(data(&w), "0");
        assert_eq!(w.cursor_flat_idx(), 0);
    }

    #[test]
    fn move_down_across_roots_when_collapsed() {
        let mut w = sample();
        w.move_down();
        assert_eq!(data(&w), "1");
        w.move_down();
        assert_eq!(data(&w), "2");
        // At last node — should not move further.
        w.move_down();
        assert_eq!(data(&w), "2");
    }

    #[test]
    fn move_up_from_first_node_is_noop() {
        let mut w = sample();
        w.move_up();
        assert_eq!(data(&w), "0");
    }

    #[test]
    fn expand_descends_and_collapse_returns() {
        let mut w = sample();
        w.move_down(); // cursor → "1"
        w.expand_or_enter();
        assert_eq!(data(&w), "1"); // expanded, cursor stays
        w.expand_or_enter();
        assert_eq!(data(&w), "1.0"); // entered first child
        w.collapse_or_leave();
        assert_eq!(data(&w), "1"); // back to parent
    }

    #[test]
    fn move_down_into_expanded_subtree() {
        let mut w = sample();
        w.move_down(); // → "1"
        w.expand_or_enter(); // expand "1"
        w.move_down(); // → "1.0"
        assert_eq!(data(&w), "1.0");
        assert_eq!(w.cursor_flat_idx(), 2);
        w.move_down(); // → "1.1"
        assert_eq!(data(&w), "1.1");
        w.move_down(); // → "2" (1.1 is collapsed)
        assert_eq!(data(&w), "2");
    }

    #[test]
    fn move_up_skips_into_deepest_expanded_child() {
        let mut w = sample();
        // Expand 1 and 1.1, making visible order: 0, 1, 1.0, 1.1, 1.1.0, 2
        w.move_down(); // → "1"
        w.expand_or_enter(); // expand "1"
        w.move_down(); // → "1.0"
        w.move_down(); // → "1.1"
        w.expand_or_enter(); // expand "1.1"
        w.move_down(); // → "1.1.0"
        w.move_down(); // → "2"
        assert_eq!(data(&w), "2");
        assert_eq!(w.cursor_flat_idx(), 5);

        w.move_up();
        assert_eq!(data(&w), "1.1.0");
        assert_eq!(w.cursor_flat_idx(), 4);
    }

    #[test]
    fn toggle_collapses_and_resyncs_flat_idx() {
        let mut w = sample();
        w.move_down(); // → "1"
        w.expand_or_enter(); // expand "1"
        w.move_down(); // → "1.0"  flat_idx = 2
        w.move_up(); // → "1"     flat_idx = 1
        w.toggle_selected(); // collapse "1"
        assert_eq!(data(&w), "1");
        assert_eq!(w.cursor_flat_idx(), 1);
        // "2" is now flat_idx 2.
        w.move_down();
        assert_eq!(data(&w), "2");
        assert_eq!(w.cursor_flat_idx(), 2);
    }

    #[test]
    fn cursor_path_reflects_position() {
        let mut w = sample();
        w.move_down(); // → roots[1]
        w.expand_or_enter();
        w.move_down(); // → roots[1].children[0]
        assert_eq!(w.cursor().path(), &[1, 0]);
        w.move_down(); // → roots[1].children[1]
        assert_eq!(w.cursor().path(), &[1, 1]);
    }

    #[test]
    fn unloaded_node_returns_needs_load() {
        let mut w: TreeWidget<&'static str> = TreeWidget::new(vec![TreeNode::unloaded("dir")]);
        assert_eq!(w.expand_or_enter(), TreeAction::NeedsLoad);
        // Cursor should not have moved.
        assert_eq!(data(&w), "dir");
    }

    #[test]
    fn inject_children_expands_node() {
        let mut w: TreeWidget<&'static str> = TreeWidget::new(vec![TreeNode::unloaded("dir")]);
        assert_eq!(w.expand_or_enter(), TreeAction::NeedsLoad);
        w.inject_children(&[0], vec![TreeNode::leaf("child")]);
        // Node should now be expanded.
        assert_eq!(data(&w), "dir");
        w.move_down();
        assert_eq!(data(&w), "child");
    }
}