use pg_query::{NodeEnum, protobuf::Node};

pub fn walk_all_nodes<F: FnMut(&NodeEnum)>(root: &NodeEnum, f: &mut F) {
    tracing::trace!(node = ?std::mem::discriminant(root), "visiting node");
    f(root);
    walk_children(root, f);
}

fn walk_children<F: FnMut(&NodeEnum)>(node: &NodeEnum, f: &mut F) {
    match node {
        NodeEnum::SelectStmt(s) => {
            walk_slice(&s.target_list, f);
            walk_slice(&s.from_clause, f);
            walk_slice(&s.group_clause, f);
            walk_slice(&s.sort_clause, f);
            if let Some(w) = s.where_clause.as_deref() {
                walk_node(w, f);
            }
            if let Some(h) = s.having_clause.as_deref() {
                walk_node(h, f);
            }
            if let Some(wc) = &s.with_clause {
                walk_slice(&wc.ctes, f);
            }
            if let Some(l) = s.larg.as_ref() {
                walk_all_nodes(&NodeEnum::SelectStmt(l.clone()), f);
            }
            if let Some(r) = s.rarg.as_ref() {
                walk_all_nodes(&NodeEnum::SelectStmt(r.clone()), f);
            }
        }
        NodeEnum::CommonTableExpr(c) => {
            if let Some(q) = c.ctequery.as_deref() {
                walk_node(q, f);
            }
        }
        NodeEnum::JoinExpr(j) => {
            if let Some(l) = j.larg.as_deref() {
                walk_node(l, f);
            }
            if let Some(r) = j.rarg.as_deref() {
                walk_node(r, f);
            }
            if let Some(q) = j.quals.as_deref() {
                walk_node(q, f);
            }
        }
        NodeEnum::RangeSubselect(rs) => {
            if let Some(q) = rs.subquery.as_deref() {
                walk_node(q, f);
            }
        }
        NodeEnum::ResTarget(rt) => {
            if let Some(v) = rt.val.as_deref() {
                walk_node(v, f);
            }
        }
        NodeEnum::FuncCall(fc) => {
            walk_slice(&fc.args, f);
        }
        NodeEnum::TypeCast(tc) => {
            if let Some(a) = tc.arg.as_deref() {
                walk_node(a, f);
            }
        }
        NodeEnum::AExpr(ae) => {
            if let Some(l) = ae.lexpr.as_deref() {
                walk_node(l, f);
            }
            if let Some(r) = ae.rexpr.as_deref() {
                walk_node(r, f);
            }
        }
        NodeEnum::BoolExpr(be) => {
            walk_slice(&be.args, f);
        }
        NodeEnum::SubLink(sl) => {
            if let Some(q) = sl.subselect.as_deref() {
                walk_node(q, f);
            }
        }
        NodeEnum::List(l) => {
            walk_slice(&l.items, f);
        }
        _ => {}
    }
}

pub(super) fn walk_slice<F: FnMut(&NodeEnum)>(nodes: &[Node], f: &mut F) {
    for n in nodes {
        walk_node(n, f);
    }
}

pub(super) fn walk_node<F: FnMut(&NodeEnum)>(node: &Node, f: &mut F) {
    if let Some(inner) = node.node.as_ref() {
        walk_all_nodes(inner, f);
    }
}
