package com.framsticks.util.swing;

import java.util.Iterator;

import javax.swing.tree.DefaultTreeModel;
import javax.swing.tree.MutableTreeNode;
import javax.swing.tree.TreeNode;

public final class TreeNodeUtils {

	private TreeNodeUtils() {
	}

	public static class TreeNodeChildrenIterator implements Iterator<TreeNode> {

		final TreeNode treeNode;
		final DefaultTreeModel treeModel;
		int index;

		/**
		 * @param treeNode
		 * @param number
		 */
		public TreeNodeChildrenIterator(TreeNode treeNode, DefaultTreeModel treeModel) {
			this.treeNode = treeNode;
			this.treeModel = treeModel;
			this.index = 0;
		}

		@Override
		public boolean hasNext() {
			return index < treeNode.getChildCount();
		}

		@Override
		public TreeNode next() {
			int i = index;
			++index;
			return (TreeNode) treeModel.getChild(treeNode, i);
		}

		@Override
		public void remove() {
			if (treeNode instanceof MutableTreeNode) {
				((MutableTreeNode) treeNode).remove(index);
				return;
			}
			throw new UnsupportedOperationException();
		}
	}

	public static Iterable<TreeNode> children(final TreeNode treeNode, final DefaultTreeModel treeModel) {
		return new Iterable<TreeNode>() {
			@Override
			public Iterator<TreeNode> iterator() {
				return new TreeNodeChildrenIterator(treeNode, treeModel);
			}
		};
	}

}