use super::Document;
use super::Element;
use super::Heading;
use super::Object;
use super::Section;
use crate::parser::DocumentElement;

pub enum Token<'r, 's> {
    Document(&'r Document<'s>),
    Heading(&'r Heading<'s>),
    Section(&'r Section<'s>),
    Object(&'r Object<'s>),
    Element(&'r Element<'s>),
}

impl<'r, 's> Token<'r, 's> {
    pub fn iter_tokens(&self) -> impl Iterator<Item = Token<'r, 's>> {
        match self {
            Token::Document(document) => document
                .zeroth_section
                .iter()
                .map(Token::Section)
                .chain(document.children.iter().map(Token::Heading)),
            Token::Heading(heading) => {
                heading
                    .title
                    .iter()
                    .map(Token::Object)
                    .chain(heading.children.iter().map(|de| match de {
                        DocumentElement::Heading(ref obj) => Token::Heading(obj),
                        DocumentElement::Section(ref obj) => Token::Section(obj),
                    }))
            }
            Token::Section(section) => section.children.iter().map(Token::Element),
            Token::Object(_) => panic!(),
            Token::Element(_) => panic!(),
        }
    }
}