scheme-edit 0.1.0

use crate::cst::{Item, Node};
use crate::document::Document;

fn item_index_of_data(items: &[Item], data_index: usize) -> Option<usize> {
    let mut seen = 0;
    for (i, item) in items.iter().enumerate() {
        if matches!(item, Item::Node(_)) {
            if seen == data_index {
                return Some(i);
            }
            seen += 1;
        }
    }
    None
}

/// Separator to place before a node at ITEM_IDX: reuse the previous data
/// sibling's leading newline+indent when there is one, else a single space.
fn separator_before(items: &[Item], item_idx: usize) -> String {
    if item_idx > 0 {
        if let Some(Item::Ws(w)) = items.get(item_idx - 1) {
            if let Some(nl) = w.rfind('\n') {
                return format!("\n{}", &w[nl + 1..]);
            }
        }
    }
    " ".to_string()
}

/// Re-render NODE at INDENT so multi-line pretty output nests correctly
/// as stored text.
fn reflow(node: Node, indent: usize) -> Node {
    let text = node.to_pretty(indent);
    match Document::parse(&text) {
        Ok(doc) => doc
            .items
            .into_iter()
            .find_map(|i| match i {
                Item::Node(n) => Some(n),
                _ => None,
            })
            .unwrap_or(node),
        Err(_) => node,
    }
}

fn indent_of(sep: &str) -> usize {
    sep.strip_prefix('\n').map_or(0, str::len)
}

impl Node {
    fn items_mut(&mut self) -> Option<&mut Vec<Item>> {
        match self {
            Node::List { items, .. } => Some(items),
            _ => None,
        }
    }

    /// Number of data children (trivia excluded).
    pub fn data_len(&self) -> usize {
        self.list_nodes().count()
    }

    pub fn data_child(&self, data_index: usize) -> Option<&Node> {
        self.list_nodes().nth(data_index)
    }

    pub fn data_child_mut(&mut self, data_index: usize) -> Option<&mut Node> {
        let items = self.items_mut()?;
        let idx = item_index_of_data(items, data_index)?;
        match &mut items[idx] {
            Item::Node(n) => Some(n),
            _ => None,
        }
    }

    /// Data index of the first child matching PRED.
    pub fn position_of<F: Fn(&Node) -> bool>(&self, pred: F) -> Option<usize> {
        self.list_nodes().position(pred)
    }

    /// Insert NODE before the DATA_INDEX-th data child; trivia stays
    /// attached to its original neighbor. Out-of-range appends.
    pub fn insert_child(&mut self, data_index: usize, node: Node) {
        let Some(items) = self.items_mut() else {
            return;
        };
        let Some(idx) = item_index_of_data(items, data_index) else {
            self.push_child(node);
            return;
        };
        let sep = separator_before(items, idx);
        let node = reflow(node, indent_of(&sep));
        items.insert(idx, Item::Node(node));
        items.insert(idx + 1, Item::Ws(sep));
    }

    /// Insert NODE after the last data child (before trailing trivia/close).
    pub fn push_child(&mut self, node: Node) {
        let last = self.data_len().checked_sub(1);
        let Some(items) = self.items_mut() else {
            return;
        };
        match last.and_then(|d| item_index_of_data(items, d)) {
            Some(idx) => {
                let sep = separator_before(items, idx);
                let node = reflow(node, indent_of(&sep));
                items.insert(idx + 1, Item::Ws(sep));
                items.insert(idx + 2, Item::Node(node));
            }
            None => items.push(Item::Node(reflow(node, 0))),
        }
    }

    pub fn replace_child(&mut self, data_index: usize, node: Node) -> bool {
        let Some(items) = self.items_mut() else {
            return false;
        };
        let Some(idx) = item_index_of_data(items, data_index) else {
            return false;
        };
        let sep = separator_before(items, idx);
        items[idx] = Item::Node(reflow(node, indent_of(&sep)));
        true
    }

    /// Remove the DATA_INDEX-th data child. With TAKE_LEADING_TRIVIA, the
    /// comments and blanks between the previous data child and this one go
    /// with it; otherwise only directly adjacent whitespace is dropped.
    pub fn remove_child(&mut self, data_index: usize, take_leading_trivia: bool) -> Option<Node> {
        let items = self.items_mut()?;
        let idx = item_index_of_data(items, data_index)?;
        let mut start = idx;
        while start > 0 {
            match &items[start - 1] {
                Item::Ws(_) => start -= 1,
                Item::LineComment(_) | Item::BlockComment(_) | Item::DatumComment(_)
                    if take_leading_trivia =>
                {
                    start -= 1
                }
                _ => break,
            }
        }
        let mut removed = None;
        for item in items.drain(start..=idx) {
            if let Item::Node(n) = item {
                removed = Some(n);
            }
        }
        removed
    }
}

#[cfg(test)]
mod tests {
    use crate::{list, sym, Document};

    #[test]
    fn push_child_indents_like_siblings() {
        let mut doc = Document::parse("(list\n  (channel a)\n  (channel b))").unwrap();
        let form = doc.forms_mut().next().unwrap();
        form.push_child(list(vec![sym("channel"), sym("c")]));
        assert_eq!(
            doc.to_string(),
            "(list\n  (channel a)\n  (channel b)\n  (channel c))"
        );
    }

    #[test]
    fn remove_child_takes_attached_comment() {
        let mut doc =
            Document::parse("(list\n  (channel a)\n  ;; b's comment\n  (channel b))").unwrap();
        let form = doc.forms_mut().next().unwrap();
        let idx = form
            .position_of(|n| n.to_source().contains("channel b"))
            .unwrap();
        form.remove_child(idx, true);
        assert_eq!(doc.to_string(), "(list\n  (channel a))");
    }

    #[test]
    fn remove_child_keeps_unrelated_regions_verbatim() {
        let src = "(list\n  (channel   a)  ; weird   spacing preserved\n  (channel b))";
        let mut doc = Document::parse(src).unwrap();
        let form = doc.forms_mut().next().unwrap();
        let idx = form
            .position_of(|n| n.to_source().contains("channel b"))
            .unwrap();
        form.remove_child(idx, false);
        assert!(doc
            .to_string()
            .contains("(channel   a)  ; weird   spacing preserved"));
    }

    #[test]
    fn insert_before_tail_symbol() {
        let mut doc = Document::parse("(cons* (channel a)\n       %default-channels)").unwrap();
        let form = doc.forms_mut().next().unwrap();
        let tail = form
            .position_of(|n| n.as_symbol() == Some("%default-channels"))
            .unwrap();
        form.insert_child(tail, list(vec![sym("channel"), sym("b")]));
        assert_eq!(
            doc.to_string(),
            "(cons* (channel a)\n       (channel b)\n       %default-channels)"
        );
    }
}