<?php
namespace Tz\WordPress\Tools\Menu;
use Exception, Countable, Iterator;

function get_nav_menu($location) {
    $locations     = get_nav_menu_locations();
    if (!isset($locations[$location])) {
        throw new Exception("{$location} is not a defined menu location");
    }
    $menu_location = _get_nav_menu_object($locations[$location])->slug;
    $menu = _get_nav_menu_items($menu_location);
    _wp_menu_item_classes_by_context($menu);

    return $menu;
}


class HierarchicalMenu implements Countable, Iterator {
    private $contents = Array();
    private $current  = 0;

    private $id_lookup     = Array();
    private $target_lookup = Array();

    public function __construct(Array $menu, $parent = 0) {
        $i = 0;
        $last_id = -1;

        foreach ($menu as $key => $item) {
            $item->_children = false;

            if ($item->menu_item_parent == $parent) {
                $this->contents[$i] = $item;
                $this->id_lookup[$item->ID] = $i;
                $this->target_lookup[$item->object_id] = $i;

                $i++;
                $last_id = $item->ID;
            } elseif ($item->menu_item_parent == $last_id) {
                $this->contents[$this->id_lookup[$last_id]]->_children = true;
            }
        }

        foreach ($this->contents as $key => $item) {
            if (true === $item->_children) {
                $item->HierarchicalMenu = new HierarchicalMenu($menu, $item->ID);
            }
            unset($this->contents[$key]->_children);
        }
    }

    public function count() {
        return count($this->contents);
    }

    public function current() {
        return $this->contents[$this->current];
    }

    public function key() {
        return $this->current;
    }

    public function next() {
        $this->current++;
    }

    public function rewind() {
        $this->rewind = 0;
    }

    public function valid() {
        return (isset($this->contents[$this->current]) ? true : false);
    }

    public function getItemByID($id) {
        return $this->id_lookup[$id] ?: false;
    }

    public function getItemByTarget($id) {
        return $this->target_lookup[$id] ?: false;
    }
}
?>