#[cfg(test)]
mod test_element_iter;

use super::{Document, Element, Field, Value};

/// Walk the tree of `Element`s from the given root `Element` depth first.
pub(crate) struct DepthFirst<'a, 'buf> {
    stack: Vec<State<'a, 'buf>>,
}

impl<'a, 'buf> DepthFirst<'a, 'buf> {
    /// Create a new depth first walker with the given `Document` as the root.
    pub(crate) fn from_doc(doc: &'a Document<'buf>) -> Self {
        Self {
            stack: vec![State::Root(doc.root())],
        }
    }
}

/// The state of a stack level.
#[derive(Debug)]
enum State<'a, 'buf> {
    Root(&'a Element<'buf>),
    Array(std::slice::Iter<'a, Element<'buf>>),
    Object(std::slice::Iter<'a, Field<'buf>>),
}

impl<'a, 'buf> Iterator for DepthFirst<'a, 'buf> {
    type Item = &'a Element<'buf>;

    /// Return the next `Element` in the depth first walk.
    fn next(&mut self) -> Option<Self::Item> {
        let elem = loop {
            let mut state = self.stack.pop()?;

            let elem = match &mut state {
                State::Root(elem) => Some(*elem),
                State::Array(iter) => iter.next(),
                State::Object(iter) => {
                    let field = iter.next();
                    field.map(Field::element)
                }
            };

            let Some(elem) = elem else {
                continue;
            };

            match &state {
                State::Array(_) | State::Object(_) => self.stack.push(state),
                State::Root(_) => (),
            }

            break elem;
        };

        match elem.value() {
            Value::Array(elements) => self.stack.push(State::Array(elements.iter())),
            Value::Object(fields) => self.stack.push(State::Object(fields.iter())),
            Value::Null | Value::True | Value::False | Value::String(_) | Value::Number(_) => (),
        }

        Some(elem)
    }
}
