use std::iter::{empty, once};

use graphviz_rust::dot_structures::{Edge, Graph, Node, Stmt, Subgraph};

pub(super) enum Elem<'a> {
    Node(&'a Node),
    Edge(&'a Edge),
}

pub(super) trait GraphExt {
    fn iter_elems<'a>(&'a self) -> Box<dyn Iterator<Item = Elem<'a>> + 'a>;
}

impl GraphExt for Stmt {
    fn iter_elems<'a>(&'a self) -> Box<dyn Iterator<Item = Elem<'a>> + 'a> {
        match self {
            Stmt::Edge(edge) => Box::new(once(Elem::Edge(edge))),
            Stmt::Node(node) => Box::new(once(Elem::Node(node))),
            Stmt::Subgraph(sg) => sg.iter_elems(),
            _ => Box::new(empty()),
        }
    }
}

impl GraphExt for Subgraph {
    fn iter_elems<'a>(&'a self) -> Box<dyn Iterator<Item = Elem<'a>> + 'a> {
        Box::new(self.stmts.iter().flat_map(Stmt::iter_elems))
    }
}

impl GraphExt for Graph {
    fn iter_elems<'a>(&'a self) -> Box<dyn Iterator<Item = Elem<'a>> + 'a> {
        let stmts = match self {
            Graph::Graph { stmts, .. } => stmts,
            Graph::DiGraph { stmts, .. } => stmts,
        };
        Box::new(stmts.iter().flat_map(Stmt::iter_elems))
    }
}
